message_cache.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891
  1. package server
  2. import (
  3. "database/sql"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "net/netip"
  8. "strings"
  9. "time"
  10. _ "github.com/mattn/go-sqlite3" // SQLite driver
  11. "heckel.io/ntfy/log"
  12. "heckel.io/ntfy/util"
  13. )
  14. var (
  15. errUnexpectedMessageType = errors.New("unexpected message type")
  16. errMessageNotFound = errors.New("message not found")
  17. )
  18. // Messages cache
  19. const (
  20. createMessagesTableQuery = `
  21. BEGIN;
  22. CREATE TABLE IF NOT EXISTS messages (
  23. id INTEGER PRIMARY KEY AUTOINCREMENT,
  24. mid TEXT NOT NULL,
  25. time INT NOT NULL,
  26. expires INT NOT NULL,
  27. topic TEXT NOT NULL,
  28. message TEXT NOT NULL,
  29. title TEXT NOT NULL,
  30. priority INT NOT NULL,
  31. tags TEXT NOT NULL,
  32. click TEXT NOT NULL,
  33. icon TEXT NOT NULL,
  34. actions TEXT NOT NULL,
  35. attachment_name TEXT NOT NULL,
  36. attachment_type TEXT NOT NULL,
  37. attachment_size INT NOT NULL,
  38. attachment_expires INT NOT NULL,
  39. attachment_url TEXT NOT NULL,
  40. attachment_deleted INT NOT NULL,
  41. sender TEXT NOT NULL,
  42. user TEXT NOT NULL,
  43. encoding TEXT NOT NULL,
  44. published INT NOT NULL
  45. );
  46. CREATE INDEX IF NOT EXISTS idx_mid ON messages (mid);
  47. CREATE INDEX IF NOT EXISTS idx_time ON messages (time);
  48. CREATE INDEX IF NOT EXISTS idx_topic ON messages (topic);
  49. CREATE INDEX IF NOT EXISTS idx_expires ON messages (expires);
  50. CREATE INDEX IF NOT EXISTS idx_sender ON messages (sender);
  51. CREATE INDEX IF NOT EXISTS idx_user ON messages (user);
  52. CREATE INDEX IF NOT EXISTS idx_attachment_expires ON messages (attachment_expires);
  53. COMMIT;
  54. `
  55. insertMessageQuery = `
  56. INSERT INTO messages (mid, time, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, attachment_deleted, sender, user, encoding, published)
  57. VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
  58. `
  59. deleteMessageQuery = `DELETE FROM messages WHERE mid = ?`
  60. updateMessagesForTopicExpiryQuery = `UPDATE messages SET expires = ? WHERE topic = ?`
  61. selectRowIDFromMessageID = `SELECT id FROM messages WHERE mid = ?` // Do not include topic, see #336 and TestServer_PollSinceID_MultipleTopics
  62. selectMessagesByIDQuery = `
  63. SELECT mid, time, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user, encoding
  64. FROM messages
  65. WHERE mid = ?
  66. `
  67. selectMessagesSinceTimeQuery = `
  68. SELECT mid, time, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user, encoding
  69. FROM messages
  70. WHERE topic = ? AND time >= ? AND published = 1
  71. ORDER BY time, id
  72. `
  73. selectMessagesSinceTimeIncludeScheduledQuery = `
  74. SELECT mid, time, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user, encoding
  75. FROM messages
  76. WHERE topic = ? AND time >= ?
  77. ORDER BY time, id
  78. `
  79. selectMessagesSinceIDQuery = `
  80. SELECT mid, time, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user, encoding
  81. FROM messages
  82. WHERE topic = ? AND id > ? AND published = 1
  83. ORDER BY time, id
  84. `
  85. selectMessagesSinceIDIncludeScheduledQuery = `
  86. SELECT mid, time, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user, encoding
  87. FROM messages
  88. WHERE topic = ? AND (id > ? OR published = 0)
  89. ORDER BY time, id
  90. `
  91. selectMessagesDueQuery = `
  92. SELECT mid, time, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user, encoding
  93. FROM messages
  94. WHERE time <= ? AND published = 0
  95. ORDER BY time, id
  96. `
  97. selectMessagesExpiredQuery = `SELECT mid FROM messages WHERE expires <= ? AND published = 1`
  98. updateMessagePublishedQuery = `UPDATE messages SET published = 1 WHERE mid = ?`
  99. selectMessagesCountQuery = `SELECT COUNT(*) FROM messages`
  100. selectMessageCountPerTopicQuery = `SELECT topic, COUNT(*) FROM messages GROUP BY topic`
  101. selectTopicsQuery = `SELECT topic FROM messages GROUP BY topic`
  102. updateAttachmentDeleted = `UPDATE messages SET attachment_deleted = 1 WHERE mid = ?`
  103. selectAttachmentsExpiredQuery = `SELECT mid FROM messages WHERE attachment_expires > 0 AND attachment_expires <= ? AND attachment_deleted = 0`
  104. selectAttachmentsSizeBySenderQuery = `SELECT IFNULL(SUM(attachment_size), 0) FROM messages WHERE user = '' AND sender = ? AND attachment_expires >= ?`
  105. selectAttachmentsSizeByUserIDQuery = `SELECT IFNULL(SUM(attachment_size), 0) FROM messages WHERE user = ? AND attachment_expires >= ?`
  106. )
  107. // Schema management queries
  108. const (
  109. currentSchemaVersion = 10
  110. createSchemaVersionTableQuery = `
  111. CREATE TABLE IF NOT EXISTS schemaVersion (
  112. id INT PRIMARY KEY,
  113. version INT NOT NULL
  114. );
  115. `
  116. insertSchemaVersion = `INSERT INTO schemaVersion VALUES (1, ?)`
  117. updateSchemaVersion = `UPDATE schemaVersion SET version = ? WHERE id = 1`
  118. selectSchemaVersionQuery = `SELECT version FROM schemaVersion WHERE id = 1`
  119. // 0 -> 1
  120. migrate0To1AlterMessagesTableQuery = `
  121. BEGIN;
  122. ALTER TABLE messages ADD COLUMN title TEXT NOT NULL DEFAULT('');
  123. ALTER TABLE messages ADD COLUMN priority INT NOT NULL DEFAULT(0);
  124. ALTER TABLE messages ADD COLUMN tags TEXT NOT NULL DEFAULT('');
  125. COMMIT;
  126. `
  127. // 1 -> 2
  128. migrate1To2AlterMessagesTableQuery = `
  129. ALTER TABLE messages ADD COLUMN published INT NOT NULL DEFAULT(1);
  130. `
  131. // 2 -> 3
  132. migrate2To3AlterMessagesTableQuery = `
  133. BEGIN;
  134. ALTER TABLE messages ADD COLUMN click TEXT NOT NULL DEFAULT('');
  135. ALTER TABLE messages ADD COLUMN attachment_name TEXT NOT NULL DEFAULT('');
  136. ALTER TABLE messages ADD COLUMN attachment_type TEXT NOT NULL DEFAULT('');
  137. ALTER TABLE messages ADD COLUMN attachment_size INT NOT NULL DEFAULT('0');
  138. ALTER TABLE messages ADD COLUMN attachment_expires INT NOT NULL DEFAULT('0');
  139. ALTER TABLE messages ADD COLUMN attachment_owner TEXT NOT NULL DEFAULT('');
  140. ALTER TABLE messages ADD COLUMN attachment_url TEXT NOT NULL DEFAULT('');
  141. COMMIT;
  142. `
  143. // 3 -> 4
  144. migrate3To4AlterMessagesTableQuery = `
  145. ALTER TABLE messages ADD COLUMN encoding TEXT NOT NULL DEFAULT('');
  146. `
  147. // 4 -> 5
  148. migrate4To5AlterMessagesTableQuery = `
  149. BEGIN;
  150. CREATE TABLE IF NOT EXISTS messages_new (
  151. id INTEGER PRIMARY KEY AUTOINCREMENT,
  152. mid TEXT NOT NULL,
  153. time INT NOT NULL,
  154. topic TEXT NOT NULL,
  155. message TEXT NOT NULL,
  156. title TEXT NOT NULL,
  157. priority INT NOT NULL,
  158. tags TEXT NOT NULL,
  159. click TEXT NOT NULL,
  160. attachment_name TEXT NOT NULL,
  161. attachment_type TEXT NOT NULL,
  162. attachment_size INT NOT NULL,
  163. attachment_expires INT NOT NULL,
  164. attachment_url TEXT NOT NULL,
  165. attachment_owner TEXT NOT NULL,
  166. encoding TEXT NOT NULL,
  167. published INT NOT NULL
  168. );
  169. CREATE INDEX IF NOT EXISTS idx_mid ON messages_new (mid);
  170. CREATE INDEX IF NOT EXISTS idx_topic ON messages_new (topic);
  171. INSERT
  172. INTO messages_new (
  173. mid, time, topic, message, title, priority, tags, click, attachment_name, attachment_type,
  174. attachment_size, attachment_expires, attachment_url, attachment_owner, encoding, published)
  175. SELECT
  176. id, time, topic, message, title, priority, tags, click, attachment_name, attachment_type,
  177. attachment_size, attachment_expires, attachment_url, attachment_owner, encoding, published
  178. FROM messages;
  179. DROP TABLE messages;
  180. ALTER TABLE messages_new RENAME TO messages;
  181. COMMIT;
  182. `
  183. // 5 -> 6
  184. migrate5To6AlterMessagesTableQuery = `
  185. ALTER TABLE messages ADD COLUMN actions TEXT NOT NULL DEFAULT('');
  186. `
  187. // 6 -> 7
  188. migrate6To7AlterMessagesTableQuery = `
  189. ALTER TABLE messages RENAME COLUMN attachment_owner TO sender;
  190. `
  191. // 7 -> 8
  192. migrate7To8AlterMessagesTableQuery = `
  193. ALTER TABLE messages ADD COLUMN icon TEXT NOT NULL DEFAULT('');
  194. `
  195. // 8 -> 9
  196. migrate8To9AlterMessagesTableQuery = `
  197. CREATE INDEX IF NOT EXISTS idx_time ON messages (time);
  198. `
  199. // 9 -> 10
  200. migrate9To10AlterMessagesTableQuery = `
  201. ALTER TABLE messages ADD COLUMN user TEXT NOT NULL DEFAULT('');
  202. ALTER TABLE messages ADD COLUMN attachment_deleted INT NOT NULL DEFAULT('0');
  203. ALTER TABLE messages ADD COLUMN expires INT NOT NULL DEFAULT('0');
  204. CREATE INDEX IF NOT EXISTS idx_expires ON messages (expires);
  205. CREATE INDEX IF NOT EXISTS idx_sender ON messages (sender);
  206. CREATE INDEX IF NOT EXISTS idx_user ON messages (user);
  207. CREATE INDEX IF NOT EXISTS idx_attachment_expires ON messages (attachment_expires);
  208. `
  209. migrate9To10UpdateMessageExpiryQuery = `UPDATE messages SET expires = time + ?`
  210. )
  211. var (
  212. migrations = map[int]func(db *sql.DB, cacheDuration time.Duration) error{
  213. 0: migrateFrom0,
  214. 1: migrateFrom1,
  215. 2: migrateFrom2,
  216. 3: migrateFrom3,
  217. 4: migrateFrom4,
  218. 5: migrateFrom5,
  219. 6: migrateFrom6,
  220. 7: migrateFrom7,
  221. 8: migrateFrom8,
  222. 9: migrateFrom9,
  223. }
  224. )
  225. type messageCache struct {
  226. db *sql.DB
  227. queue *util.BatchingQueue[*message]
  228. nop bool
  229. }
  230. // newSqliteCache creates a SQLite file-backed cache
  231. func newSqliteCache(filename, startupQueries string, cacheDuration time.Duration, batchSize int, batchTimeout time.Duration, nop bool) (*messageCache, error) {
  232. db, err := sql.Open("sqlite3", filename)
  233. if err != nil {
  234. return nil, err
  235. }
  236. if err := setupDB(db, startupQueries, cacheDuration); err != nil {
  237. return nil, err
  238. }
  239. var queue *util.BatchingQueue[*message]
  240. if batchSize > 0 || batchTimeout > 0 {
  241. queue = util.NewBatchingQueue[*message](batchSize, batchTimeout)
  242. }
  243. cache := &messageCache{
  244. db: db,
  245. queue: queue,
  246. nop: nop,
  247. }
  248. go cache.processMessageBatches()
  249. return cache, nil
  250. }
  251. // newMemCache creates an in-memory cache
  252. func newMemCache() (*messageCache, error) {
  253. return newSqliteCache(createMemoryFilename(), "", 0, 0, 0, false)
  254. }
  255. // newNopCache creates an in-memory cache that discards all messages;
  256. // it is always empty and can be used if caching is entirely disabled
  257. func newNopCache() (*messageCache, error) {
  258. return newSqliteCache(createMemoryFilename(), "", 0, 0, 0, true)
  259. }
  260. // createMemoryFilename creates a unique memory filename to use for the SQLite backend.
  261. // From mattn/go-sqlite3: "Each connection to ":memory:" opens a brand new in-memory
  262. // sql database, so if the stdlib's sql engine happens to open another connection and
  263. // you've only specified ":memory:", that connection will see a brand new database.
  264. // A workaround is to use "file::memory:?cache=shared" (or "file:foobar?mode=memory&cache=shared").
  265. // Every connection to this string will point to the same in-memory database."
  266. func createMemoryFilename() string {
  267. return fmt.Sprintf("file:%s?mode=memory&cache=shared", util.RandomString(10))
  268. }
  269. // AddMessage stores a message to the message cache synchronously, or queues it to be stored at a later date asyncronously.
  270. // The message is queued only if "batchSize" or "batchTimeout" are passed to the constructor.
  271. func (c *messageCache) AddMessage(m *message) error {
  272. if c.queue != nil {
  273. c.queue.Enqueue(m)
  274. return nil
  275. }
  276. return c.addMessages([]*message{m})
  277. }
  278. // addMessages synchronously stores a match of messages. If the database is locked, the transaction waits until
  279. // SQLite's busy_timeout is exceeded before erroring out.
  280. func (c *messageCache) addMessages(ms []*message) error {
  281. if c.nop {
  282. return nil
  283. }
  284. if len(ms) == 0 {
  285. return nil
  286. }
  287. start := time.Now()
  288. tx, err := c.db.Begin()
  289. if err != nil {
  290. return err
  291. }
  292. defer tx.Rollback()
  293. stmt, err := tx.Prepare(insertMessageQuery)
  294. if err != nil {
  295. return err
  296. }
  297. defer stmt.Close()
  298. for _, m := range ms {
  299. if m.Event != messageEvent {
  300. return errUnexpectedMessageType
  301. }
  302. published := m.Time <= time.Now().Unix()
  303. tags := strings.Join(m.Tags, ",")
  304. var attachmentName, attachmentType, attachmentURL string
  305. var attachmentSize, attachmentExpires, attachmentDeleted int64
  306. if m.Attachment != nil {
  307. attachmentName = m.Attachment.Name
  308. attachmentType = m.Attachment.Type
  309. attachmentSize = m.Attachment.Size
  310. attachmentExpires = m.Attachment.Expires
  311. attachmentURL = m.Attachment.URL
  312. }
  313. var actionsStr string
  314. if len(m.Actions) > 0 {
  315. actionsBytes, err := json.Marshal(m.Actions)
  316. if err != nil {
  317. return err
  318. }
  319. actionsStr = string(actionsBytes)
  320. }
  321. var sender string
  322. if m.Sender.IsValid() {
  323. sender = m.Sender.String()
  324. }
  325. _, err := stmt.Exec(
  326. m.ID,
  327. m.Time,
  328. m.Expires,
  329. m.Topic,
  330. m.Message,
  331. m.Title,
  332. m.Priority,
  333. tags,
  334. m.Click,
  335. m.Icon,
  336. actionsStr,
  337. attachmentName,
  338. attachmentType,
  339. attachmentSize,
  340. attachmentExpires,
  341. attachmentURL,
  342. attachmentDeleted, // Always zero
  343. sender,
  344. m.User,
  345. m.Encoding,
  346. published,
  347. )
  348. if err != nil {
  349. return err
  350. }
  351. }
  352. if err := tx.Commit(); err != nil {
  353. log.Tag(tagMessageCache).Err(err).Error("Writing %d message(s) failed (took %v)", len(ms), time.Since(start))
  354. return err
  355. }
  356. log.Tag(tagMessageCache).Debug("Wrote %d message(s) in %v", len(ms), time.Since(start))
  357. return nil
  358. }
  359. func (c *messageCache) Messages(topic string, since sinceMarker, scheduled bool) ([]*message, error) {
  360. if since.IsNone() {
  361. return make([]*message, 0), nil
  362. } else if since.IsID() {
  363. return c.messagesSinceID(topic, since, scheduled)
  364. }
  365. return c.messagesSinceTime(topic, since, scheduled)
  366. }
  367. func (c *messageCache) messagesSinceTime(topic string, since sinceMarker, scheduled bool) ([]*message, error) {
  368. var rows *sql.Rows
  369. var err error
  370. if scheduled {
  371. rows, err = c.db.Query(selectMessagesSinceTimeIncludeScheduledQuery, topic, since.Time().Unix())
  372. } else {
  373. rows, err = c.db.Query(selectMessagesSinceTimeQuery, topic, since.Time().Unix())
  374. }
  375. if err != nil {
  376. return nil, err
  377. }
  378. return readMessages(rows)
  379. }
  380. func (c *messageCache) messagesSinceID(topic string, since sinceMarker, scheduled bool) ([]*message, error) {
  381. idrows, err := c.db.Query(selectRowIDFromMessageID, since.ID())
  382. if err != nil {
  383. return nil, err
  384. }
  385. defer idrows.Close()
  386. if !idrows.Next() {
  387. return c.messagesSinceTime(topic, sinceAllMessages, scheduled)
  388. }
  389. var rowID int64
  390. if err := idrows.Scan(&rowID); err != nil {
  391. return nil, err
  392. }
  393. idrows.Close()
  394. var rows *sql.Rows
  395. if scheduled {
  396. rows, err = c.db.Query(selectMessagesSinceIDIncludeScheduledQuery, topic, rowID)
  397. } else {
  398. rows, err = c.db.Query(selectMessagesSinceIDQuery, topic, rowID)
  399. }
  400. if err != nil {
  401. return nil, err
  402. }
  403. return readMessages(rows)
  404. }
  405. func (c *messageCache) MessagesDue() ([]*message, error) {
  406. rows, err := c.db.Query(selectMessagesDueQuery, time.Now().Unix())
  407. if err != nil {
  408. return nil, err
  409. }
  410. return readMessages(rows)
  411. }
  412. // MessagesExpired returns a list of IDs for messages that have expires (should be deleted)
  413. func (c *messageCache) MessagesExpired() ([]string, error) {
  414. rows, err := c.db.Query(selectMessagesExpiredQuery, time.Now().Unix())
  415. if err != nil {
  416. return nil, err
  417. }
  418. defer rows.Close()
  419. ids := make([]string, 0)
  420. for rows.Next() {
  421. var id string
  422. if err := rows.Scan(&id); err != nil {
  423. return nil, err
  424. }
  425. ids = append(ids, id)
  426. }
  427. if err := rows.Err(); err != nil {
  428. return nil, err
  429. }
  430. return ids, nil
  431. }
  432. func (c *messageCache) Message(id string) (*message, error) {
  433. rows, err := c.db.Query(selectMessagesByIDQuery, id)
  434. if err != nil {
  435. return nil, err
  436. }
  437. if !rows.Next() {
  438. return nil, errMessageNotFound
  439. }
  440. defer rows.Close()
  441. return readMessage(rows)
  442. }
  443. func (c *messageCache) MarkPublished(m *message) error {
  444. _, err := c.db.Exec(updateMessagePublishedQuery, m.ID)
  445. return err
  446. }
  447. func (c *messageCache) MessageCounts() (map[string]int, error) {
  448. rows, err := c.db.Query(selectMessageCountPerTopicQuery)
  449. if err != nil {
  450. return nil, err
  451. }
  452. defer rows.Close()
  453. var topic string
  454. var count int
  455. counts := make(map[string]int)
  456. for rows.Next() {
  457. if err := rows.Scan(&topic, &count); err != nil {
  458. return nil, err
  459. } else if err := rows.Err(); err != nil {
  460. return nil, err
  461. }
  462. counts[topic] = count
  463. }
  464. return counts, nil
  465. }
  466. func (c *messageCache) Topics() (map[string]*topic, error) {
  467. rows, err := c.db.Query(selectTopicsQuery)
  468. if err != nil {
  469. return nil, err
  470. }
  471. defer rows.Close()
  472. topics := make(map[string]*topic)
  473. for rows.Next() {
  474. var id string
  475. if err := rows.Scan(&id); err != nil {
  476. return nil, err
  477. }
  478. topics[id] = newTopic(id)
  479. }
  480. if err := rows.Err(); err != nil {
  481. return nil, err
  482. }
  483. return topics, nil
  484. }
  485. func (c *messageCache) DeleteMessages(ids ...string) error {
  486. tx, err := c.db.Begin()
  487. if err != nil {
  488. return err
  489. }
  490. defer tx.Rollback()
  491. for _, id := range ids {
  492. if _, err := tx.Exec(deleteMessageQuery, id); err != nil {
  493. return err
  494. }
  495. }
  496. return tx.Commit()
  497. }
  498. func (c *messageCache) ExpireMessages(topics ...string) error {
  499. tx, err := c.db.Begin()
  500. if err != nil {
  501. return err
  502. }
  503. defer tx.Rollback()
  504. for _, t := range topics {
  505. if _, err := tx.Exec(updateMessagesForTopicExpiryQuery, time.Now().Unix(), t); err != nil {
  506. return err
  507. }
  508. }
  509. return tx.Commit()
  510. }
  511. func (c *messageCache) AttachmentsExpired() ([]string, error) {
  512. rows, err := c.db.Query(selectAttachmentsExpiredQuery, time.Now().Unix())
  513. if err != nil {
  514. return nil, err
  515. }
  516. defer rows.Close()
  517. ids := make([]string, 0)
  518. for rows.Next() {
  519. var id string
  520. if err := rows.Scan(&id); err != nil {
  521. return nil, err
  522. }
  523. ids = append(ids, id)
  524. }
  525. if err := rows.Err(); err != nil {
  526. return nil, err
  527. }
  528. return ids, nil
  529. }
  530. func (c *messageCache) MarkAttachmentsDeleted(ids ...string) error {
  531. tx, err := c.db.Begin()
  532. if err != nil {
  533. return err
  534. }
  535. defer tx.Rollback()
  536. for _, id := range ids {
  537. if _, err := tx.Exec(updateAttachmentDeleted, id); err != nil {
  538. return err
  539. }
  540. }
  541. return tx.Commit()
  542. }
  543. func (c *messageCache) AttachmentBytesUsedBySender(sender string) (int64, error) {
  544. rows, err := c.db.Query(selectAttachmentsSizeBySenderQuery, sender, time.Now().Unix())
  545. if err != nil {
  546. return 0, err
  547. }
  548. return c.readAttachmentBytesUsed(rows)
  549. }
  550. func (c *messageCache) AttachmentBytesUsedByUser(userID string) (int64, error) {
  551. rows, err := c.db.Query(selectAttachmentsSizeByUserIDQuery, userID, time.Now().Unix())
  552. if err != nil {
  553. return 0, err
  554. }
  555. return c.readAttachmentBytesUsed(rows)
  556. }
  557. func (c *messageCache) readAttachmentBytesUsed(rows *sql.Rows) (int64, error) {
  558. defer rows.Close()
  559. var size int64
  560. if !rows.Next() {
  561. return 0, errors.New("no rows found")
  562. }
  563. if err := rows.Scan(&size); err != nil {
  564. return 0, err
  565. } else if err := rows.Err(); err != nil {
  566. return 0, err
  567. }
  568. return size, nil
  569. }
  570. func (c *messageCache) processMessageBatches() {
  571. if c.queue == nil {
  572. return
  573. }
  574. for messages := range c.queue.Dequeue() {
  575. if err := c.addMessages(messages); err != nil {
  576. log.Tag(tagMessageCache).Err(err).Error("Cannot write message batch")
  577. }
  578. }
  579. }
  580. func readMessages(rows *sql.Rows) ([]*message, error) {
  581. defer rows.Close()
  582. messages := make([]*message, 0)
  583. for rows.Next() {
  584. m, err := readMessage(rows)
  585. if err != nil {
  586. return nil, err
  587. }
  588. messages = append(messages, m)
  589. }
  590. if err := rows.Err(); err != nil {
  591. return nil, err
  592. }
  593. return messages, nil
  594. }
  595. func readMessage(rows *sql.Rows) (*message, error) {
  596. var timestamp, expires, attachmentSize, attachmentExpires int64
  597. var priority int
  598. var id, topic, msg, title, tagsStr, click, icon, actionsStr, attachmentName, attachmentType, attachmentURL, sender, user, encoding string
  599. err := rows.Scan(
  600. &id,
  601. &timestamp,
  602. &expires,
  603. &topic,
  604. &msg,
  605. &title,
  606. &priority,
  607. &tagsStr,
  608. &click,
  609. &icon,
  610. &actionsStr,
  611. &attachmentName,
  612. &attachmentType,
  613. &attachmentSize,
  614. &attachmentExpires,
  615. &attachmentURL,
  616. &sender,
  617. &user,
  618. &encoding,
  619. )
  620. if err != nil {
  621. return nil, err
  622. }
  623. var tags []string
  624. if tagsStr != "" {
  625. tags = strings.Split(tagsStr, ",")
  626. }
  627. var actions []*action
  628. if actionsStr != "" {
  629. if err := json.Unmarshal([]byte(actionsStr), &actions); err != nil {
  630. return nil, err
  631. }
  632. }
  633. senderIP, err := netip.ParseAddr(sender)
  634. if err != nil {
  635. senderIP = netip.Addr{} // if no IP stored in database, return invalid address
  636. }
  637. var att *attachment
  638. if attachmentName != "" && attachmentURL != "" {
  639. att = &attachment{
  640. Name: attachmentName,
  641. Type: attachmentType,
  642. Size: attachmentSize,
  643. Expires: attachmentExpires,
  644. URL: attachmentURL,
  645. }
  646. }
  647. return &message{
  648. ID: id,
  649. Time: timestamp,
  650. Expires: expires,
  651. Event: messageEvent,
  652. Topic: topic,
  653. Message: msg,
  654. Title: title,
  655. Priority: priority,
  656. Tags: tags,
  657. Click: click,
  658. Icon: icon,
  659. Actions: actions,
  660. Attachment: att,
  661. Sender: senderIP, // Must parse assuming database must be correct
  662. User: user,
  663. Encoding: encoding,
  664. }, nil
  665. }
  666. func (c *messageCache) Close() error {
  667. return c.db.Close()
  668. }
  669. func setupDB(db *sql.DB, startupQueries string, cacheDuration time.Duration) error {
  670. // Run startup queries
  671. if startupQueries != "" {
  672. if _, err := db.Exec(startupQueries); err != nil {
  673. return err
  674. }
  675. }
  676. // If 'messages' table does not exist, this must be a new database
  677. rowsMC, err := db.Query(selectMessagesCountQuery)
  678. if err != nil {
  679. return setupNewCacheDB(db)
  680. }
  681. rowsMC.Close()
  682. // If 'messages' table exists, check 'schemaVersion' table
  683. schemaVersion := 0
  684. rowsSV, err := db.Query(selectSchemaVersionQuery)
  685. if err == nil {
  686. defer rowsSV.Close()
  687. if !rowsSV.Next() {
  688. return errors.New("cannot determine schema version: cache file may be corrupt")
  689. }
  690. if err := rowsSV.Scan(&schemaVersion); err != nil {
  691. return err
  692. }
  693. rowsSV.Close()
  694. }
  695. // Do migrations
  696. if schemaVersion == currentSchemaVersion {
  697. return nil
  698. } else if schemaVersion > currentSchemaVersion {
  699. return fmt.Errorf("unexpected schema version: version %d is higher than current version %d", schemaVersion, currentSchemaVersion)
  700. }
  701. for i := schemaVersion; i < currentSchemaVersion; i++ {
  702. fn, ok := migrations[i]
  703. if !ok {
  704. return fmt.Errorf("cannot find migration step from schema version %d to %d", i, i+1)
  705. } else if err := fn(db, cacheDuration); err != nil {
  706. return err
  707. }
  708. }
  709. return nil
  710. }
  711. func setupNewCacheDB(db *sql.DB) error {
  712. if _, err := db.Exec(createMessagesTableQuery); err != nil {
  713. return err
  714. }
  715. if _, err := db.Exec(createSchemaVersionTableQuery); err != nil {
  716. return err
  717. }
  718. if _, err := db.Exec(insertSchemaVersion, currentSchemaVersion); err != nil {
  719. return err
  720. }
  721. return nil
  722. }
  723. func migrateFrom0(db *sql.DB, _ time.Duration) error {
  724. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 0 to 1")
  725. if _, err := db.Exec(migrate0To1AlterMessagesTableQuery); err != nil {
  726. return err
  727. }
  728. if _, err := db.Exec(createSchemaVersionTableQuery); err != nil {
  729. return err
  730. }
  731. if _, err := db.Exec(insertSchemaVersion, 1); err != nil {
  732. return err
  733. }
  734. return nil
  735. }
  736. func migrateFrom1(db *sql.DB, _ time.Duration) error {
  737. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 1 to 2")
  738. if _, err := db.Exec(migrate1To2AlterMessagesTableQuery); err != nil {
  739. return err
  740. }
  741. if _, err := db.Exec(updateSchemaVersion, 2); err != nil {
  742. return err
  743. }
  744. return nil
  745. }
  746. func migrateFrom2(db *sql.DB, _ time.Duration) error {
  747. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 2 to 3")
  748. if _, err := db.Exec(migrate2To3AlterMessagesTableQuery); err != nil {
  749. return err
  750. }
  751. if _, err := db.Exec(updateSchemaVersion, 3); err != nil {
  752. return err
  753. }
  754. return nil
  755. }
  756. func migrateFrom3(db *sql.DB, _ time.Duration) error {
  757. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 3 to 4")
  758. if _, err := db.Exec(migrate3To4AlterMessagesTableQuery); err != nil {
  759. return err
  760. }
  761. if _, err := db.Exec(updateSchemaVersion, 4); err != nil {
  762. return err
  763. }
  764. return nil
  765. }
  766. func migrateFrom4(db *sql.DB, _ time.Duration) error {
  767. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 4 to 5")
  768. if _, err := db.Exec(migrate4To5AlterMessagesTableQuery); err != nil {
  769. return err
  770. }
  771. if _, err := db.Exec(updateSchemaVersion, 5); err != nil {
  772. return err
  773. }
  774. return nil
  775. }
  776. func migrateFrom5(db *sql.DB, _ time.Duration) error {
  777. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 5 to 6")
  778. if _, err := db.Exec(migrate5To6AlterMessagesTableQuery); err != nil {
  779. return err
  780. }
  781. if _, err := db.Exec(updateSchemaVersion, 6); err != nil {
  782. return err
  783. }
  784. return nil
  785. }
  786. func migrateFrom6(db *sql.DB, _ time.Duration) error {
  787. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 6 to 7")
  788. if _, err := db.Exec(migrate6To7AlterMessagesTableQuery); err != nil {
  789. return err
  790. }
  791. if _, err := db.Exec(updateSchemaVersion, 7); err != nil {
  792. return err
  793. }
  794. return nil
  795. }
  796. func migrateFrom7(db *sql.DB, _ time.Duration) error {
  797. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 7 to 8")
  798. if _, err := db.Exec(migrate7To8AlterMessagesTableQuery); err != nil {
  799. return err
  800. }
  801. if _, err := db.Exec(updateSchemaVersion, 8); err != nil {
  802. return err
  803. }
  804. return nil
  805. }
  806. func migrateFrom8(db *sql.DB, _ time.Duration) error {
  807. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 8 to 9")
  808. if _, err := db.Exec(migrate8To9AlterMessagesTableQuery); err != nil {
  809. return err
  810. }
  811. if _, err := db.Exec(updateSchemaVersion, 9); err != nil {
  812. return err
  813. }
  814. return nil
  815. }
  816. func migrateFrom9(db *sql.DB, cacheDuration time.Duration) error {
  817. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 9 to 10")
  818. tx, err := db.Begin()
  819. if err != nil {
  820. return err
  821. }
  822. defer tx.Rollback()
  823. if _, err := tx.Exec(migrate9To10AlterMessagesTableQuery); err != nil {
  824. return err
  825. }
  826. if _, err := tx.Exec(migrate9To10UpdateMessageExpiryQuery, int64(cacheDuration.Seconds())); err != nil {
  827. return err
  828. }
  829. if _, err := tx.Exec(updateSchemaVersion, 10); err != nil {
  830. return err
  831. }
  832. return tx.Commit()
  833. }