puppet.go 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
  2. // Copyright (C) 2018 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"
  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() error {
  27. _, err := pq.db.Exec(`CREATE TABLE IF NOT EXISTS puppet (
  28. jid VARCHAR(25) PRIMARY KEY,
  29. avatar VARCHAR(255),
  30. displayname VARCHAR(255),
  31. name_quality TINYINT
  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=?", 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. err := row.Scan(&puppet.JID, &displayname, &avatar)
  70. if err != nil {
  71. if err != sql.ErrNoRows {
  72. puppet.log.Errorln("Database scan failed:", err)
  73. }
  74. return nil
  75. }
  76. puppet.Displayname = displayname.String
  77. puppet.Avatar = avatar.String
  78. return puppet
  79. }
  80. func (puppet *Puppet) Insert() error {
  81. _, err := puppet.db.Exec("INSERT INTO puppet VALUES (?, ?, ?, ?)",
  82. puppet.JID, puppet.Avatar, puppet.Displayname, puppet.NameQuality)
  83. if err != nil {
  84. puppet.log.Warnfln("Failed to insert %s: %v", puppet.JID, err)
  85. }
  86. return err
  87. }
  88. func (puppet *Puppet) Update() error {
  89. _, err := puppet.db.Exec("UPDATE puppet SET displayname=?, name_quality=?, avatar=? WHERE jid=?",
  90. puppet.Displayname, puppet.NameQuality, puppet.Avatar, puppet.JID)
  91. if err != nil {
  92. puppet.log.Warnfln("Failed to update %s->%s: %v", puppet.JID, err)
  93. }
  94. return err
  95. }