puppet.go 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  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. log "maunium.net/go/maulogger/v2"
  20. "maunium.net/go/mautrix-whatsapp/types"
  21. )
  22. type PuppetQuery struct {
  23. db *Database
  24. log log.Logger
  25. }
  26. func (pq *PuppetQuery) CreateTable(dbType string) error {
  27. _, err := pq.db.Exec(`CREATE TABLE IF NOT EXISTS puppet (
  28. jid VARCHAR(255) PRIMARY KEY,
  29. avatar VARCHAR(255),
  30. displayname VARCHAR(255),
  31. name_quality SMALLINT
  32. )`)
  33. return err
  34. }
  35. func (pq *PuppetQuery) New() *Puppet {
  36. return &Puppet{
  37. db: pq.db,
  38. log: pq.log,
  39. }
  40. }
  41. func (pq *PuppetQuery) GetAll() (puppets []*Puppet) {
  42. rows, err := pq.db.Query("SELECT * FROM puppet")
  43. if err != nil || rows == nil {
  44. return nil
  45. }
  46. defer rows.Close()
  47. for rows.Next() {
  48. puppets = append(puppets, pq.New().Scan(rows))
  49. }
  50. return
  51. }
  52. func (pq *PuppetQuery) Get(jid types.WhatsAppID) *Puppet {
  53. row := pq.db.QueryRow("SELECT * FROM puppet WHERE jid=$1", jid)
  54. if row == nil {
  55. return nil
  56. }
  57. return pq.New().Scan(row)
  58. }
  59. type Puppet struct {
  60. db *Database
  61. log log.Logger
  62. JID types.WhatsAppID
  63. Avatar string
  64. Displayname string
  65. NameQuality int8
  66. }
  67. func (puppet *Puppet) Scan(row Scannable) *Puppet {
  68. var displayname, avatar sql.NullString
  69. var quality sql.NullInt64
  70. err := row.Scan(&puppet.JID, &avatar, &displayname, &quality)
  71. if err != nil {
  72. if err != sql.ErrNoRows {
  73. puppet.log.Errorln("Database scan failed:", err)
  74. }
  75. return nil
  76. }
  77. puppet.Displayname = displayname.String
  78. puppet.Avatar = avatar.String
  79. puppet.NameQuality = int8(quality.Int64)
  80. return puppet
  81. }
  82. func (puppet *Puppet) Insert() {
  83. _, err := puppet.db.Exec("INSERT INTO puppet VALUES ($1, $2, $3, $4)",
  84. puppet.JID, puppet.Avatar, puppet.Displayname, puppet.NameQuality)
  85. if err != nil {
  86. puppet.log.Warnfln("Failed to insert %s: %v", puppet.JID, err)
  87. }
  88. }
  89. func (puppet *Puppet) Update() {
  90. _, err := puppet.db.Exec("UPDATE puppet SET displayname=$1, name_quality=$2, avatar=$3 WHERE jid=$4",
  91. puppet.Displayname, puppet.NameQuality, puppet.Avatar, puppet.JID)
  92. if err != nil {
  93. puppet.log.Warnfln("Failed to update %s->%s: %v", puppet.JID, err)
  94. }
  95. }