2019-08-25-move-state-store-to-db.go 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. package upgrades
  2. import (
  3. "database/sql"
  4. "encoding/json"
  5. "fmt"
  6. "io/ioutil"
  7. "os"
  8. "strings"
  9. "maunium.net/go/mautrix"
  10. "maunium.net/go/mautrix-appservice"
  11. )
  12. func init() {
  13. migrateRegistrations := func(tx *sql.Tx, registrations map[string]bool) error {
  14. values := make([]interface{}, 0, len(registrations))
  15. valueStrings := make([]string, 0, len(registrations))
  16. i := 1
  17. for userID, registered := range registrations {
  18. if registered {
  19. values = append(values, userID)
  20. valueStrings = append(valueStrings, fmt.Sprintf("($%d)", i))
  21. i++
  22. }
  23. }
  24. valueString := strings.Join(valueStrings, ",")
  25. _, err := tx.Exec("INSERT INTO mx_registrations (user_id) VALUES "+valueString, values...)
  26. return err
  27. }
  28. migrateMemberships := func(tx *sql.Tx, rooms map[string]map[string]mautrix.Membership) error {
  29. var values []interface{}
  30. var valueStrings []string
  31. i := 1
  32. for roomID, members := range rooms {
  33. for userID, membership := range members {
  34. values = append(values, roomID, userID, membership)
  35. valueStrings = append(valueStrings, fmt.Sprintf("($%d, $%d, $%d)", i, i+1, i+2))
  36. i += 3
  37. }
  38. }
  39. valueString := strings.Join(valueStrings, ",")
  40. _, err := tx.Exec("INSERT INTO mx_user_profile (room_id, user_id, membership) VALUES "+valueString, values...)
  41. return err
  42. }
  43. migratePowerLevels := func(tx *sql.Tx, rooms map[string]*mautrix.PowerLevels) error {
  44. var values []interface{}
  45. var valueStrings []string
  46. i := 1
  47. for roomID, powerLevels := range rooms {
  48. powerLevelBytes, err := json.Marshal(powerLevels)
  49. if err != nil {
  50. return err
  51. }
  52. values = append(values, roomID, powerLevelBytes)
  53. valueStrings = append(valueStrings, fmt.Sprintf("($%d, $%d)", i, i+1))
  54. i += 2
  55. }
  56. valueString := strings.Join(valueStrings, ",")
  57. _, err := tx.Exec("INSERT INTO mx_room_state (room_id, power_levels) VALUES "+valueString, values...)
  58. return err
  59. }
  60. userProfileTable := `CREATE TABLE mx_user_profile (
  61. room_id VARCHAR(255),
  62. user_id VARCHAR(255),
  63. membership VARCHAR(15) NOT NULL,
  64. PRIMARY KEY (room_id, user_id)
  65. )`
  66. roomStateTable := `CREATE TABLE mx_room_state (
  67. room_id VARCHAR(255) PRIMARY KEY,
  68. power_levels TEXT
  69. )`
  70. registrationsTable := `CREATE TABLE mx_registrations (
  71. user_id VARCHAR(255) PRIMARY KEY
  72. )`
  73. upgrades[9] = upgrade{"Move state store to main DB", func(tx *sql.Tx, ctx context) error {
  74. store := appservice.NewBasicStateStore().(*appservice.BasicStateStore)
  75. if ctx.dialect == Postgres {
  76. roomStateTable = strings.Replace(roomStateTable, "TEXT", "JSONB", 1)
  77. }
  78. if _, err := tx.Exec(userProfileTable); err != nil {
  79. return err
  80. } else if _, err = tx.Exec(roomStateTable); err != nil {
  81. return err
  82. } else if _, err = tx.Exec(registrationsTable); err != nil {
  83. return err
  84. } else if data, err := ioutil.ReadFile("mx-state.json"); err != nil {
  85. ctx.log.Debugln("mx-state.json not found, not migrating state store")
  86. } else if err = json.Unmarshal(data, &store); err != nil {
  87. return err
  88. } else if err = migrateRegistrations(tx, store.Registrations); err != nil {
  89. return err
  90. } else if err = migrateMemberships(tx, store.Memberships); err != nil {
  91. return err
  92. } else if err = migratePowerLevels(tx, store.PowerLevels); err != nil {
  93. return err
  94. } else if err = os.Rename("mx-state.json", "mx-state.json.bak"); err != nil {
  95. return err
  96. }
  97. return nil
  98. }}
  99. }