message_cache.go 29 KB

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