message_cache.go 29 KB

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