1
0

message_cache.go 29 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018
  1. package server
  2. import (
  3. "database/sql"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "net/netip"
  8. "path/filepath"
  9. "strings"
  10. "time"
  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. nop bool
  265. }
  266. // newSqliteCache creates a SQLite file-backed cache
  267. func newSqliteCache(filename, startupQueries string, cacheDuration time.Duration, batchSize int, batchTimeout time.Duration, nop bool) (*messageCache, error) {
  268. // Check the parent directory of the database file (makes for friendly error messages)
  269. parentDir := filepath.Dir(filename)
  270. if !util.FileExists(parentDir) {
  271. return nil, fmt.Errorf("cache database directory %s does not exist or is not accessible", parentDir)
  272. }
  273. // Open database
  274. db, err := sql.Open("sqlite3", filename)
  275. if err != nil {
  276. return nil, err
  277. }
  278. if err := setupMessagesDB(db, startupQueries, cacheDuration); err != nil {
  279. return nil, err
  280. }
  281. var queue *util.BatchingQueue[*message]
  282. if batchSize > 0 || batchTimeout > 0 {
  283. queue = util.NewBatchingQueue[*message](batchSize, batchTimeout)
  284. }
  285. cache := &messageCache{
  286. db: db,
  287. queue: queue,
  288. nop: nop,
  289. }
  290. go cache.processMessageBatches()
  291. return cache, nil
  292. }
  293. // newMemCache creates an in-memory cache
  294. func newMemCache() (*messageCache, error) {
  295. return newSqliteCache(createMemoryFilename(), "", 0, 0, 0, false)
  296. }
  297. // newNopCache creates an in-memory cache that discards all messages;
  298. // it is always empty and can be used if caching is entirely disabled
  299. func newNopCache() (*messageCache, error) {
  300. return newSqliteCache(createMemoryFilename(), "", 0, 0, 0, true)
  301. }
  302. // createMemoryFilename creates a unique memory filename to use for the SQLite backend.
  303. // From mattn/go-sqlite3: "Each connection to ":memory:" opens a brand new in-memory
  304. // sql database, so if the stdlib's sql engine happens to open another connection and
  305. // you've only specified ":memory:", that connection will see a brand new database.
  306. // A workaround is to use "file::memory:?cache=shared" (or "file:foobar?mode=memory&cache=shared").
  307. // Every connection to this string will point to the same in-memory database."
  308. func createMemoryFilename() string {
  309. return fmt.Sprintf("file:%s?mode=memory&cache=shared", util.RandomString(10))
  310. }
  311. // AddMessage stores a message to the message cache synchronously, or queues it to be stored at a later date asyncronously.
  312. // The message is queued only if "batchSize" or "batchTimeout" are passed to the constructor.
  313. func (c *messageCache) AddMessage(m *message) error {
  314. if c.queue != nil {
  315. c.queue.Enqueue(m)
  316. return nil
  317. }
  318. return c.addMessages([]*message{m})
  319. }
  320. // addMessages synchronously stores a match of messages. If the database is locked, the transaction waits until
  321. // SQLite's busy_timeout is exceeded before erroring out.
  322. func (c *messageCache) addMessages(ms []*message) error {
  323. if c.nop {
  324. return nil
  325. }
  326. if len(ms) == 0 {
  327. return nil
  328. }
  329. start := time.Now()
  330. tx, err := c.db.Begin()
  331. if err != nil {
  332. return err
  333. }
  334. defer tx.Rollback()
  335. stmt, err := tx.Prepare(insertMessageQuery)
  336. if err != nil {
  337. return err
  338. }
  339. defer stmt.Close()
  340. for _, m := range ms {
  341. if m.Event != messageEvent {
  342. return errUnexpectedMessageType
  343. }
  344. published := m.Time <= time.Now().Unix()
  345. tags := strings.Join(m.Tags, ",")
  346. var attachmentName, attachmentType, attachmentURL string
  347. var attachmentSize, attachmentExpires, attachmentDeleted int64
  348. if m.Attachment != nil {
  349. attachmentName = m.Attachment.Name
  350. attachmentType = m.Attachment.Type
  351. attachmentSize = m.Attachment.Size
  352. attachmentExpires = m.Attachment.Expires
  353. attachmentURL = m.Attachment.URL
  354. }
  355. var actionsStr string
  356. if len(m.Actions) > 0 {
  357. actionsBytes, err := json.Marshal(m.Actions)
  358. if err != nil {
  359. return err
  360. }
  361. actionsStr = string(actionsBytes)
  362. }
  363. var sender string
  364. if m.Sender.IsValid() {
  365. sender = m.Sender.String()
  366. }
  367. _, err := stmt.Exec(
  368. m.ID,
  369. m.Time,
  370. m.Expires,
  371. m.Topic,
  372. m.Message,
  373. m.Title,
  374. m.Priority,
  375. tags,
  376. m.Click,
  377. m.Icon,
  378. actionsStr,
  379. attachmentName,
  380. attachmentType,
  381. attachmentSize,
  382. attachmentExpires,
  383. attachmentURL,
  384. attachmentDeleted, // Always zero
  385. sender,
  386. m.User,
  387. m.ContentType,
  388. m.Encoding,
  389. published,
  390. )
  391. if err != nil {
  392. return err
  393. }
  394. }
  395. if err := tx.Commit(); err != nil {
  396. log.Tag(tagMessageCache).Err(err).Error("Writing %d message(s) failed (took %v)", len(ms), time.Since(start))
  397. return err
  398. }
  399. log.Tag(tagMessageCache).Debug("Wrote %d message(s) in %v", len(ms), time.Since(start))
  400. return nil
  401. }
  402. func (c *messageCache) Messages(topic string, since sinceMarker, scheduled bool) ([]*message, error) {
  403. if since.IsNone() {
  404. return make([]*message, 0), nil
  405. } else if since.IsLatest() {
  406. return c.messagesLatest(topic)
  407. } else if since.IsID() {
  408. return c.messagesSinceID(topic, since, scheduled)
  409. }
  410. return c.messagesSinceTime(topic, since, scheduled)
  411. }
  412. func (c *messageCache) messagesSinceTime(topic string, since sinceMarker, scheduled bool) ([]*message, error) {
  413. var rows *sql.Rows
  414. var err error
  415. if scheduled {
  416. rows, err = c.db.Query(selectMessagesSinceTimeIncludeScheduledQuery, topic, since.Time().Unix())
  417. } else {
  418. rows, err = c.db.Query(selectMessagesSinceTimeQuery, topic, since.Time().Unix())
  419. }
  420. if err != nil {
  421. return nil, err
  422. }
  423. return readMessages(rows)
  424. }
  425. func (c *messageCache) messagesSinceID(topic string, since sinceMarker, scheduled bool) ([]*message, error) {
  426. idrows, err := c.db.Query(selectRowIDFromMessageID, since.ID())
  427. if err != nil {
  428. return nil, err
  429. }
  430. defer idrows.Close()
  431. if !idrows.Next() {
  432. return c.messagesSinceTime(topic, sinceAllMessages, scheduled)
  433. }
  434. var rowID int64
  435. if err := idrows.Scan(&rowID); err != nil {
  436. return nil, err
  437. }
  438. idrows.Close()
  439. var rows *sql.Rows
  440. if scheduled {
  441. rows, err = c.db.Query(selectMessagesSinceIDIncludeScheduledQuery, topic, rowID)
  442. } else {
  443. rows, err = c.db.Query(selectMessagesSinceIDQuery, topic, rowID)
  444. }
  445. if err != nil {
  446. return nil, err
  447. }
  448. return readMessages(rows)
  449. }
  450. func (c *messageCache) messagesLatest(topic string) ([]*message, error) {
  451. rows, err := c.db.Query(selectMessagesLatestQuery, topic)
  452. if err != nil {
  453. return nil, err
  454. }
  455. return readMessages(rows)
  456. }
  457. func (c *messageCache) MessagesDue() ([]*message, error) {
  458. rows, err := c.db.Query(selectMessagesDueQuery, time.Now().Unix())
  459. if err != nil {
  460. return nil, err
  461. }
  462. return readMessages(rows)
  463. }
  464. // MessagesExpired returns a list of IDs for messages that have expires (should be deleted)
  465. func (c *messageCache) MessagesExpired() ([]string, error) {
  466. rows, err := c.db.Query(selectMessagesExpiredQuery, time.Now().Unix())
  467. if err != nil {
  468. return nil, err
  469. }
  470. defer rows.Close()
  471. ids := make([]string, 0)
  472. for rows.Next() {
  473. var id string
  474. if err := rows.Scan(&id); err != nil {
  475. return nil, err
  476. }
  477. ids = append(ids, id)
  478. }
  479. if err := rows.Err(); err != nil {
  480. return nil, err
  481. }
  482. return ids, nil
  483. }
  484. func (c *messageCache) Message(id string) (*message, error) {
  485. rows, err := c.db.Query(selectMessagesByIDQuery, id)
  486. if err != nil {
  487. return nil, err
  488. }
  489. if !rows.Next() {
  490. return nil, errMessageNotFound
  491. }
  492. defer rows.Close()
  493. return readMessage(rows)
  494. }
  495. func (c *messageCache) MarkPublished(m *message) error {
  496. _, err := c.db.Exec(updateMessagePublishedQuery, m.ID)
  497. return err
  498. }
  499. func (c *messageCache) MessageCounts() (map[string]int, error) {
  500. rows, err := c.db.Query(selectMessageCountPerTopicQuery)
  501. if err != nil {
  502. return nil, err
  503. }
  504. defer rows.Close()
  505. var topic string
  506. var count int
  507. counts := make(map[string]int)
  508. for rows.Next() {
  509. if err := rows.Scan(&topic, &count); err != nil {
  510. return nil, err
  511. } else if err := rows.Err(); err != nil {
  512. return nil, err
  513. }
  514. counts[topic] = count
  515. }
  516. return counts, nil
  517. }
  518. func (c *messageCache) Topics() (map[string]*topic, error) {
  519. rows, err := c.db.Query(selectTopicsQuery)
  520. if err != nil {
  521. return nil, err
  522. }
  523. defer rows.Close()
  524. topics := make(map[string]*topic)
  525. for rows.Next() {
  526. var id string
  527. if err := rows.Scan(&id); err != nil {
  528. return nil, err
  529. }
  530. topics[id] = newTopic(id)
  531. }
  532. if err := rows.Err(); err != nil {
  533. return nil, err
  534. }
  535. return topics, nil
  536. }
  537. func (c *messageCache) DeleteMessages(ids ...string) error {
  538. tx, err := c.db.Begin()
  539. if err != nil {
  540. return err
  541. }
  542. defer tx.Rollback()
  543. for _, id := range ids {
  544. if _, err := tx.Exec(deleteMessageQuery, id); err != nil {
  545. return err
  546. }
  547. }
  548. return tx.Commit()
  549. }
  550. func (c *messageCache) ExpireMessages(topics ...string) error {
  551. tx, err := c.db.Begin()
  552. if err != nil {
  553. return err
  554. }
  555. defer tx.Rollback()
  556. for _, t := range topics {
  557. if _, err := tx.Exec(updateMessagesForTopicExpiryQuery, time.Now().Unix()-1, t); err != nil {
  558. return err
  559. }
  560. }
  561. return tx.Commit()
  562. }
  563. func (c *messageCache) AttachmentsExpired() ([]string, error) {
  564. rows, err := c.db.Query(selectAttachmentsExpiredQuery, time.Now().Unix())
  565. if err != nil {
  566. return nil, err
  567. }
  568. defer rows.Close()
  569. ids := make([]string, 0)
  570. for rows.Next() {
  571. var id string
  572. if err := rows.Scan(&id); err != nil {
  573. return nil, err
  574. }
  575. ids = append(ids, id)
  576. }
  577. if err := rows.Err(); err != nil {
  578. return nil, err
  579. }
  580. return ids, nil
  581. }
  582. func (c *messageCache) MarkAttachmentsDeleted(ids ...string) error {
  583. tx, err := c.db.Begin()
  584. if err != nil {
  585. return err
  586. }
  587. defer tx.Rollback()
  588. for _, id := range ids {
  589. if _, err := tx.Exec(updateAttachmentDeleted, id); err != nil {
  590. return err
  591. }
  592. }
  593. return tx.Commit()
  594. }
  595. func (c *messageCache) AttachmentBytesUsedBySender(sender string) (int64, error) {
  596. rows, err := c.db.Query(selectAttachmentsSizeBySenderQuery, sender, time.Now().Unix())
  597. if err != nil {
  598. return 0, err
  599. }
  600. return c.readAttachmentBytesUsed(rows)
  601. }
  602. func (c *messageCache) AttachmentBytesUsedByUser(userID string) (int64, error) {
  603. rows, err := c.db.Query(selectAttachmentsSizeByUserIDQuery, userID, time.Now().Unix())
  604. if err != nil {
  605. return 0, err
  606. }
  607. return c.readAttachmentBytesUsed(rows)
  608. }
  609. func (c *messageCache) readAttachmentBytesUsed(rows *sql.Rows) (int64, error) {
  610. defer rows.Close()
  611. var size int64
  612. if !rows.Next() {
  613. return 0, errors.New("no rows found")
  614. }
  615. if err := rows.Scan(&size); err != nil {
  616. return 0, err
  617. } else if err := rows.Err(); err != nil {
  618. return 0, err
  619. }
  620. return size, nil
  621. }
  622. func (c *messageCache) processMessageBatches() {
  623. if c.queue == nil {
  624. return
  625. }
  626. for messages := range c.queue.Dequeue() {
  627. if err := c.addMessages(messages); err != nil {
  628. log.Tag(tagMessageCache).Err(err).Error("Cannot write message batch")
  629. }
  630. }
  631. }
  632. func readMessages(rows *sql.Rows) ([]*message, error) {
  633. defer rows.Close()
  634. messages := make([]*message, 0)
  635. for rows.Next() {
  636. m, err := readMessage(rows)
  637. if err != nil {
  638. return nil, err
  639. }
  640. messages = append(messages, m)
  641. }
  642. if err := rows.Err(); err != nil {
  643. return nil, err
  644. }
  645. return messages, nil
  646. }
  647. func readMessage(rows *sql.Rows) (*message, error) {
  648. var timestamp, expires, attachmentSize, attachmentExpires int64
  649. var priority int
  650. var id, topic, msg, title, tagsStr, click, icon, actionsStr, attachmentName, attachmentType, attachmentURL, sender, user, contentType, encoding string
  651. err := rows.Scan(
  652. &id,
  653. &timestamp,
  654. &expires,
  655. &topic,
  656. &msg,
  657. &title,
  658. &priority,
  659. &tagsStr,
  660. &click,
  661. &icon,
  662. &actionsStr,
  663. &attachmentName,
  664. &attachmentType,
  665. &attachmentSize,
  666. &attachmentExpires,
  667. &attachmentURL,
  668. &sender,
  669. &user,
  670. &contentType,
  671. &encoding,
  672. )
  673. if err != nil {
  674. return nil, err
  675. }
  676. var tags []string
  677. if tagsStr != "" {
  678. tags = strings.Split(tagsStr, ",")
  679. }
  680. var actions []*action
  681. if actionsStr != "" {
  682. if err := json.Unmarshal([]byte(actionsStr), &actions); err != nil {
  683. return nil, err
  684. }
  685. }
  686. senderIP, err := netip.ParseAddr(sender)
  687. if err != nil {
  688. senderIP = netip.Addr{} // if no IP stored in database, return invalid address
  689. }
  690. var att *attachment
  691. if attachmentName != "" && attachmentURL != "" {
  692. att = &attachment{
  693. Name: attachmentName,
  694. Type: attachmentType,
  695. Size: attachmentSize,
  696. Expires: attachmentExpires,
  697. URL: attachmentURL,
  698. }
  699. }
  700. return &message{
  701. ID: id,
  702. Time: timestamp,
  703. Expires: expires,
  704. Event: messageEvent,
  705. Topic: topic,
  706. Message: msg,
  707. Title: title,
  708. Priority: priority,
  709. Tags: tags,
  710. Click: click,
  711. Icon: icon,
  712. Actions: actions,
  713. Attachment: att,
  714. Sender: senderIP, // Must parse assuming database must be correct
  715. User: user,
  716. ContentType: contentType,
  717. Encoding: encoding,
  718. }, nil
  719. }
  720. func (c *messageCache) UpdateStats(messages int64) error {
  721. _, err := c.db.Exec(updateStatsQuery, messages)
  722. return err
  723. }
  724. func (c *messageCache) Stats() (messages int64, err error) {
  725. rows, err := c.db.Query(selectStatsQuery)
  726. if err != nil {
  727. return 0, err
  728. }
  729. defer rows.Close()
  730. if !rows.Next() {
  731. return 0, errNoRows
  732. }
  733. if err := rows.Scan(&messages); err != nil {
  734. return 0, err
  735. }
  736. return messages, nil
  737. }
  738. func (c *messageCache) Close() error {
  739. return c.db.Close()
  740. }
  741. func setupMessagesDB(db *sql.DB, startupQueries string, cacheDuration time.Duration) error {
  742. // Run startup queries
  743. if startupQueries != "" {
  744. if _, err := db.Exec(startupQueries); err != nil {
  745. return err
  746. }
  747. }
  748. // If 'messages' table does not exist, this must be a new database
  749. rowsMC, err := db.Query(selectMessagesCountQuery)
  750. if err != nil {
  751. return setupNewCacheDB(db)
  752. }
  753. rowsMC.Close()
  754. // If 'messages' table exists, check 'schemaVersion' table
  755. schemaVersion := 0
  756. rowsSV, err := db.Query(selectSchemaVersionQuery)
  757. if err == nil {
  758. defer rowsSV.Close()
  759. if !rowsSV.Next() {
  760. return errors.New("cannot determine schema version: cache file may be corrupt")
  761. }
  762. if err := rowsSV.Scan(&schemaVersion); err != nil {
  763. return err
  764. }
  765. rowsSV.Close()
  766. }
  767. // Do migrations
  768. if schemaVersion == currentSchemaVersion {
  769. return nil
  770. } else if schemaVersion > currentSchemaVersion {
  771. return fmt.Errorf("unexpected schema version: version %d is higher than current version %d", schemaVersion, currentSchemaVersion)
  772. }
  773. for i := schemaVersion; i < currentSchemaVersion; i++ {
  774. fn, ok := migrations[i]
  775. if !ok {
  776. return fmt.Errorf("cannot find migration step from schema version %d to %d", i, i+1)
  777. } else if err := fn(db, cacheDuration); err != nil {
  778. return err
  779. }
  780. }
  781. return nil
  782. }
  783. func setupNewCacheDB(db *sql.DB) error {
  784. if _, err := db.Exec(createMessagesTableQuery); err != nil {
  785. return err
  786. }
  787. if _, err := db.Exec(createSchemaVersionTableQuery); err != nil {
  788. return err
  789. }
  790. if _, err := db.Exec(insertSchemaVersion, currentSchemaVersion); err != nil {
  791. return err
  792. }
  793. return nil
  794. }
  795. func migrateFrom0(db *sql.DB, _ time.Duration) error {
  796. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 0 to 1")
  797. if _, err := db.Exec(migrate0To1AlterMessagesTableQuery); err != nil {
  798. return err
  799. }
  800. if _, err := db.Exec(createSchemaVersionTableQuery); err != nil {
  801. return err
  802. }
  803. if _, err := db.Exec(insertSchemaVersion, 1); err != nil {
  804. return err
  805. }
  806. return nil
  807. }
  808. func migrateFrom1(db *sql.DB, _ time.Duration) error {
  809. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 1 to 2")
  810. if _, err := db.Exec(migrate1To2AlterMessagesTableQuery); err != nil {
  811. return err
  812. }
  813. if _, err := db.Exec(updateSchemaVersion, 2); err != nil {
  814. return err
  815. }
  816. return nil
  817. }
  818. func migrateFrom2(db *sql.DB, _ time.Duration) error {
  819. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 2 to 3")
  820. if _, err := db.Exec(migrate2To3AlterMessagesTableQuery); err != nil {
  821. return err
  822. }
  823. if _, err := db.Exec(updateSchemaVersion, 3); err != nil {
  824. return err
  825. }
  826. return nil
  827. }
  828. func migrateFrom3(db *sql.DB, _ time.Duration) error {
  829. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 3 to 4")
  830. if _, err := db.Exec(migrate3To4AlterMessagesTableQuery); err != nil {
  831. return err
  832. }
  833. if _, err := db.Exec(updateSchemaVersion, 4); err != nil {
  834. return err
  835. }
  836. return nil
  837. }
  838. func migrateFrom4(db *sql.DB, _ time.Duration) error {
  839. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 4 to 5")
  840. if _, err := db.Exec(migrate4To5AlterMessagesTableQuery); err != nil {
  841. return err
  842. }
  843. if _, err := db.Exec(updateSchemaVersion, 5); err != nil {
  844. return err
  845. }
  846. return nil
  847. }
  848. func migrateFrom5(db *sql.DB, _ time.Duration) error {
  849. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 5 to 6")
  850. if _, err := db.Exec(migrate5To6AlterMessagesTableQuery); err != nil {
  851. return err
  852. }
  853. if _, err := db.Exec(updateSchemaVersion, 6); err != nil {
  854. return err
  855. }
  856. return nil
  857. }
  858. func migrateFrom6(db *sql.DB, _ time.Duration) error {
  859. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 6 to 7")
  860. if _, err := db.Exec(migrate6To7AlterMessagesTableQuery); err != nil {
  861. return err
  862. }
  863. if _, err := db.Exec(updateSchemaVersion, 7); err != nil {
  864. return err
  865. }
  866. return nil
  867. }
  868. func migrateFrom7(db *sql.DB, _ time.Duration) error {
  869. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 7 to 8")
  870. if _, err := db.Exec(migrate7To8AlterMessagesTableQuery); err != nil {
  871. return err
  872. }
  873. if _, err := db.Exec(updateSchemaVersion, 8); err != nil {
  874. return err
  875. }
  876. return nil
  877. }
  878. func migrateFrom8(db *sql.DB, _ time.Duration) error {
  879. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 8 to 9")
  880. if _, err := db.Exec(migrate8To9AlterMessagesTableQuery); err != nil {
  881. return err
  882. }
  883. if _, err := db.Exec(updateSchemaVersion, 9); err != nil {
  884. return err
  885. }
  886. return nil
  887. }
  888. func migrateFrom9(db *sql.DB, cacheDuration time.Duration) error {
  889. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 9 to 10")
  890. tx, err := db.Begin()
  891. if err != nil {
  892. return err
  893. }
  894. defer tx.Rollback()
  895. if _, err := tx.Exec(migrate9To10AlterMessagesTableQuery); err != nil {
  896. return err
  897. }
  898. if _, err := tx.Exec(migrate9To10UpdateMessageExpiryQuery, int64(cacheDuration.Seconds())); err != nil {
  899. return err
  900. }
  901. if _, err := tx.Exec(updateSchemaVersion, 10); err != nil {
  902. return err
  903. }
  904. return tx.Commit()
  905. }
  906. func migrateFrom10(db *sql.DB, _ time.Duration) error {
  907. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 10 to 11")
  908. tx, err := db.Begin()
  909. if err != nil {
  910. return err
  911. }
  912. defer tx.Rollback()
  913. if _, err := tx.Exec(migrate10To11AlterMessagesTableQuery); err != nil {
  914. return err
  915. }
  916. if _, err := tx.Exec(updateSchemaVersion, 11); err != nil {
  917. return err
  918. }
  919. return tx.Commit()
  920. }
  921. func migrateFrom11(db *sql.DB, _ time.Duration) error {
  922. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 11 to 12")
  923. tx, err := db.Begin()
  924. if err != nil {
  925. return err
  926. }
  927. defer tx.Rollback()
  928. if _, err := tx.Exec(migrate11To12AlterMessagesTableQuery); err != nil {
  929. return err
  930. }
  931. if _, err := tx.Exec(updateSchemaVersion, 12); err != nil {
  932. return err
  933. }
  934. return tx.Commit()
  935. }
  936. func migrateFrom12(db *sql.DB, _ time.Duration) error {
  937. log.Tag(tagMessageCache).Info("Migrating cache database schema: from 12 to 13")
  938. tx, err := db.Begin()
  939. if err != nil {
  940. return err
  941. }
  942. defer tx.Rollback()
  943. if _, err := tx.Exec(migrate12To13AlterMessagesTableQuery); err != nil {
  944. return err
  945. }
  946. if _, err := tx.Exec(updateSchemaVersion, 13); err != nil {
  947. return err
  948. }
  949. return tx.Commit()
  950. }