message_cache.go 32 KB

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