backfillqueue.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
  2. // Copyright (C) 2021 Tulir Asokan, Sumner Evans
  3. //
  4. // This program is free software: you can redistribute it and/or modify
  5. // it under the terms of the GNU Affero General Public License as published by
  6. // the Free Software Foundation, either version 3 of the License, or
  7. // (at your option) any later version.
  8. //
  9. // This program is distributed in the hope that it will be useful,
  10. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  11. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  12. // GNU Affero General Public License for more details.
  13. //
  14. // You should have received a copy of the GNU Affero General Public License
  15. // along with this program. If not, see <https://www.gnu.org/licenses/>.
  16. package database
  17. import (
  18. "database/sql"
  19. "errors"
  20. "time"
  21. log "maunium.net/go/maulogger/v2"
  22. "maunium.net/go/mautrix/id"
  23. )
  24. type BackfillType int
  25. const (
  26. BackfillImmediate BackfillType = 0
  27. BackfillDeferred = 1
  28. BackfillMedia = 2
  29. )
  30. type BackfillQuery struct {
  31. db *Database
  32. log log.Logger
  33. }
  34. func (bq *BackfillQuery) New() *Backfill {
  35. return &Backfill{
  36. db: bq.db,
  37. log: bq.log,
  38. Portal: &PortalKey{},
  39. }
  40. }
  41. func (bq *BackfillQuery) NewWithValues(userID id.UserID, backfillType BackfillType, priority int, portal *PortalKey, timeStart *time.Time, timeEnd *time.Time, maxBatchEvents, maxTotalEvents, batchDelay int) *Backfill {
  42. return &Backfill{
  43. db: bq.db,
  44. log: bq.log,
  45. UserID: userID,
  46. BackfillType: backfillType,
  47. Priority: priority,
  48. Portal: portal,
  49. TimeStart: timeStart,
  50. TimeEnd: timeEnd,
  51. MaxBatchEvents: maxBatchEvents,
  52. MaxTotalEvents: maxTotalEvents,
  53. BatchDelay: batchDelay,
  54. }
  55. }
  56. const (
  57. getNextBackfillQuery = `
  58. SELECT queue_id, user_mxid, type, priority, portal_jid, portal_receiver, time_start, time_end, max_batch_events, max_total_events, batch_delay
  59. FROM backfill_queue
  60. WHERE user_mxid=$1
  61. AND type=$2
  62. AND completed_at IS NULL
  63. ORDER BY priority, queue_id
  64. LIMIT 1
  65. `
  66. )
  67. // GetNext returns the next backfill to perform
  68. func (bq *BackfillQuery) GetNext(userID id.UserID, backfillType BackfillType) (backfill *Backfill) {
  69. rows, err := bq.db.Query(getNextBackfillQuery, userID, backfillType)
  70. defer rows.Close()
  71. if err != nil || rows == nil {
  72. bq.log.Error(err)
  73. return
  74. }
  75. if rows.Next() {
  76. backfill = bq.New().Scan(rows)
  77. }
  78. return
  79. }
  80. func (bq *BackfillQuery) DeleteAll(userID id.UserID) error {
  81. _, err := bq.db.Exec("DELETE FROM backfill_queue WHERE user_mxid=$1", userID)
  82. return err
  83. }
  84. type Backfill struct {
  85. db *Database
  86. log log.Logger
  87. // Fields
  88. QueueID int
  89. UserID id.UserID
  90. BackfillType BackfillType
  91. Priority int
  92. Portal *PortalKey
  93. TimeStart *time.Time
  94. TimeEnd *time.Time
  95. MaxBatchEvents int
  96. MaxTotalEvents int
  97. BatchDelay int
  98. CompletedAt *time.Time
  99. }
  100. func (b *Backfill) Scan(row Scannable) *Backfill {
  101. err := row.Scan(&b.QueueID, &b.UserID, &b.BackfillType, &b.Priority, &b.Portal.JID, &b.Portal.Receiver, &b.TimeStart, &b.TimeEnd, &b.MaxBatchEvents, &b.MaxTotalEvents, &b.BatchDelay)
  102. if err != nil {
  103. if !errors.Is(err, sql.ErrNoRows) {
  104. b.log.Errorln("Database scan failed:", err)
  105. }
  106. return nil
  107. }
  108. return b
  109. }
  110. func (b *Backfill) Insert() {
  111. rows, err := b.db.Query(`
  112. INSERT INTO backfill_queue
  113. (user_mxid, type, priority, portal_jid, portal_receiver, time_start, time_end, max_batch_events, max_total_events, batch_delay, completed_at)
  114. VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
  115. RETURNING queue_id
  116. `, b.UserID, b.BackfillType, b.Priority, b.Portal.JID, b.Portal.Receiver, b.TimeStart, b.TimeEnd, b.MaxBatchEvents, b.MaxTotalEvents, b.BatchDelay, b.CompletedAt)
  117. defer rows.Close()
  118. if err != nil || !rows.Next() {
  119. b.log.Warnfln("Failed to insert %v/%s with priority %d: %v", b.BackfillType, b.Portal.JID, b.Priority, err)
  120. return
  121. }
  122. err = rows.Scan(&b.QueueID)
  123. if err != nil {
  124. b.log.Warnfln("Failed to insert %s/%s with priority %s: %v", b.BackfillType, b.Portal.JID, b.Priority, err)
  125. }
  126. }
  127. func (b *Backfill) MarkDone() {
  128. if b.QueueID == 0 {
  129. b.log.Errorf("Cannot delete backfill without queue_id. Maybe it wasn't actually inserted in the database?")
  130. return
  131. }
  132. _, err := b.db.Exec("UPDATE backfill_queue SET completed_at=$1 WHERE queue_id=$2", time.Now(), b.QueueID)
  133. if err != nil {
  134. b.log.Warnfln("Failed to mark %s/%s as complete: %v", b.BackfillType, b.Priority, err)
  135. }
  136. }