1
0

message_cache.go 25 KB

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