portal.go 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
  2. // Copyright (C) 2019 Tulir Asokan
  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. "strings"
  20. log "maunium.net/go/maulogger/v2"
  21. "maunium.net/go/mautrix-whatsapp/types"
  22. )
  23. type PortalKey struct {
  24. JID types.WhatsAppID
  25. Receiver types.WhatsAppID
  26. }
  27. func GroupPortalKey(jid types.WhatsAppID) PortalKey {
  28. return PortalKey{
  29. JID: jid,
  30. Receiver: jid,
  31. }
  32. }
  33. func NewPortalKey(jid, receiver types.WhatsAppID) PortalKey {
  34. if strings.HasSuffix(jid, "@g.us") {
  35. receiver = jid
  36. }
  37. return PortalKey{
  38. JID: jid,
  39. Receiver: receiver,
  40. }
  41. }
  42. func (key PortalKey) String() string {
  43. if key.Receiver == key.JID {
  44. return key.JID
  45. }
  46. return key.JID + "-" + key.Receiver
  47. }
  48. type PortalQuery struct {
  49. db *Database
  50. log log.Logger
  51. }
  52. func (pq *PortalQuery) CreateTable(dbType string) error {
  53. _, err := pq.db.Exec(`CREATE TABLE IF NOT EXISTS portal (
  54. jid VARCHAR(255),
  55. receiver VARCHAR(255),
  56. mxid VARCHAR(255) UNIQUE,
  57. name VARCHAR(255) NOT NULL,
  58. topic VARCHAR(255) NOT NULL,
  59. avatar VARCHAR(255) NOT NULL,
  60. PRIMARY KEY (jid, receiver)
  61. )`)
  62. return err
  63. }
  64. func (pq *PortalQuery) New() *Portal {
  65. return &Portal{
  66. db: pq.db,
  67. log: pq.log,
  68. }
  69. }
  70. func (pq *PortalQuery) GetAll() (portals []*Portal) {
  71. rows, err := pq.db.Query("SELECT * FROM portal")
  72. if err != nil || rows == nil {
  73. return nil
  74. }
  75. defer rows.Close()
  76. for rows.Next() {
  77. portals = append(portals, pq.New().Scan(rows))
  78. }
  79. return
  80. }
  81. func (pq *PortalQuery) GetByJID(key PortalKey) *Portal {
  82. return pq.get("SELECT * FROM portal WHERE jid=$1 AND receiver=$2", key.JID, key.Receiver)
  83. }
  84. func (pq *PortalQuery) GetByMXID(mxid types.MatrixRoomID) *Portal {
  85. return pq.get("SELECT * FROM portal WHERE mxid=$1", mxid)
  86. }
  87. func (pq *PortalQuery) get(query string, args ...interface{}) *Portal {
  88. row := pq.db.QueryRow(query, args...)
  89. if row == nil {
  90. return nil
  91. }
  92. return pq.New().Scan(row)
  93. }
  94. type Portal struct {
  95. db *Database
  96. log log.Logger
  97. Key PortalKey
  98. MXID types.MatrixRoomID
  99. Name string
  100. Topic string
  101. Avatar string
  102. }
  103. func (portal *Portal) Scan(row Scannable) *Portal {
  104. var mxid sql.NullString
  105. err := row.Scan(&portal.Key.JID, &portal.Key.Receiver, &mxid, &portal.Name, &portal.Topic, &portal.Avatar)
  106. if err != nil {
  107. if err != sql.ErrNoRows {
  108. portal.log.Errorln("Database scan failed:", err)
  109. }
  110. return nil
  111. }
  112. portal.MXID = mxid.String
  113. return portal
  114. }
  115. func (portal *Portal) mxidPtr() *string {
  116. if len(portal.MXID) > 0 {
  117. return &portal.MXID
  118. }
  119. return nil
  120. }
  121. func (portal *Portal) Insert() {
  122. _, err := portal.db.Exec("INSERT INTO portal VALUES ($1, $2, $3, $4, $5, $6)",
  123. portal.Key.JID, portal.Key.Receiver, portal.mxidPtr(), portal.Name, portal.Topic, portal.Avatar)
  124. if err != nil {
  125. portal.log.Warnfln("Failed to insert %s: %v", portal.Key, err)
  126. }
  127. }
  128. func (portal *Portal) Update() {
  129. var mxid *string
  130. if len(portal.MXID) > 0 {
  131. mxid = &portal.MXID
  132. }
  133. _, err := portal.db.Exec("UPDATE portal SET mxid=$1, name=$2, topic=$3, avatar=$4 WHERE jid=$5 AND receiver=$6",
  134. mxid, portal.Name, portal.Topic, portal.Avatar, portal.Key.JID, portal.Key.Receiver)
  135. if err != nil {
  136. portal.log.Warnfln("Failed to update %s: %v", portal.Key, err)
  137. }
  138. }