Move Matrix state store to main database. Fixes #64
This commit is contained in:
Tulir Asokan
@ -29,7 +29,8 @@ import (
|
||||
|
||||
type Database struct {
|
||||
*sql.DB
|
||||
log log.Logger
|
||||
log log.Logger
|
||||
dialect string
|
||||
|
||||
User *UserQuery
|
||||
Portal *PortalQuery
|
||||
@ -48,8 +49,9 @@ func New(dbType string, uri string) (*Database, error) {
|
||||
}
|
||||
|
||||
db := &Database{
|
||||
DB: conn,
|
||||
log: log.Sub("Database"),
|
||||
DB: conn,
|
||||
log: log.Sub("Database"),
|
||||
dialect: dbType,
|
||||
}
|
||||
db.User = &UserQuery{
|
||||
db: db,
|
||||
@ -70,8 +72,8 @@ func New(dbType string, uri string) (*Database, error) {
|
||||
return db, nil
|
||||
}
|
||||
|
||||
func (db *Database) Init(dialectName string) error {
|
||||
return upgrades.Run(db.log.Sub("Upgrade"), dialectName, db.DB)
|
||||
func (db *Database) Init() error {
|
||||
return upgrades.Run(db.log.Sub("Upgrade"), db.dialect, db.DB)
|
||||
}
|
||||
|
||||
type Scannable interface {
|
||||
|
||||
240
database/statestore.go
Normal file
240
database/statestore.go
Normal file
@ -0,0 +1,240 @@
|
||||
// mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
|
||||
// Copyright (C) 2019 Tulir Asokan
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
package database
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
log "maunium.net/go/maulogger/v2"
|
||||
|
||||
"maunium.net/go/mautrix"
|
||||
"maunium.net/go/mautrix-appservice"
|
||||
)
|
||||
|
||||
type SQLStateStore struct {
|
||||
*appservice.TypingStateStore
|
||||
|
||||
db *Database
|
||||
log log.Logger
|
||||
|
||||
Typing map[string]map[string]int64
|
||||
typingLock sync.RWMutex
|
||||
}
|
||||
|
||||
func NewSQLStateStore(db *Database) *SQLStateStore {
|
||||
return &SQLStateStore{
|
||||
TypingStateStore: appservice.NewTypingStateStore(),
|
||||
db: db,
|
||||
log: log.Sub("StateStore"),
|
||||
}
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) IsRegistered(userID string) bool {
|
||||
row := store.db.QueryRow("SELECT EXISTS(SELECT 1 FROM mx_registrations WHERE user_id=$1)", userID)
|
||||
var isRegistered bool
|
||||
err := row.Scan(&isRegistered)
|
||||
if err != nil {
|
||||
store.log.Warnfln("Failed to scan registration existence for %s: %v", userID, err)
|
||||
}
|
||||
return isRegistered
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) MarkRegistered(userID string) {
|
||||
var err error
|
||||
if store.db.dialect == "postgres" {
|
||||
_, err = store.db.Exec("INSERT INTO mx_registrations (user_id) VALUES ($1) ON CONFLICT (user_id) DO NOTHING", userID)
|
||||
} else if store.db.dialect == "sqlite3" {
|
||||
_, err = store.db.Exec("INSERT OR REPLACE INTO mx_registrations (user_id) VALUES ($1)", userID)
|
||||
} else {
|
||||
err = fmt.Errorf("unsupported dialect %s", store.db.dialect)
|
||||
}
|
||||
if err != nil {
|
||||
store.log.Warnfln("Failed to mark %s as registered: %v", userID, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) GetRoomMemberships(roomID string) map[string]mautrix.Membership {
|
||||
memberships := make(map[string]mautrix.Membership)
|
||||
rows, err := store.db.Query("SELECT user_id, membership FROM mx_user_profile WHERE room_id=$1", roomID)
|
||||
if err != nil {
|
||||
return memberships
|
||||
}
|
||||
var userID string
|
||||
var membership mautrix.Membership
|
||||
for rows.Next() {
|
||||
err := rows.Scan(&userID, &membership)
|
||||
if err != nil {
|
||||
store.log.Warnfln("Failed to scan membership in %s: %v", roomID, err)
|
||||
} else {
|
||||
memberships[userID] = membership
|
||||
}
|
||||
}
|
||||
return memberships
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) GetMembership(roomID, userID string) mautrix.Membership {
|
||||
row := store.db.QueryRow("SELECT membership FROM mx_user_profile WHERE room_id=$1 AND user_id=$2", roomID, userID)
|
||||
membership := mautrix.MembershipLeave
|
||||
if row != nil && row.Scan(&membership) != nil {
|
||||
store.log.Warnln("Failed to scan membership of %s in %s: %v", userID, roomID, membership)
|
||||
}
|
||||
return membership
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) IsInRoom(roomID, userID string) bool {
|
||||
return store.IsMembership(roomID, userID, "join")
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) IsInvited(roomID, userID string) bool {
|
||||
return store.IsMembership(roomID, userID, "join", "invite")
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) IsMembership(roomID, userID string, allowedMemberships ...mautrix.Membership) bool {
|
||||
membership := store.GetMembership(roomID, userID)
|
||||
for _, allowedMembership := range allowedMemberships {
|
||||
if allowedMembership == membership {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
func (store *SQLStateStore) SetMembership(roomID, userID string, membership mautrix.Membership) {
|
||||
var err error
|
||||
if store.db.dialect == "postgres" {
|
||||
_, err = store.db.Exec(`INSERT INTO mx_user_profile (room_id, user_id, membership) VALUES ($1, $2, $3)
|
||||
ON CONFLICT (room_id, user_id) DO UPDATE SET membership=$3`, roomID, userID, membership)
|
||||
} else if store.db.dialect == "sqlite3" {
|
||||
_, err = store.db.Exec("INSERT OR REPLACE INTO mx_user_profile (room_id, user_id, membership) VALUES ($1, $2, $3)", roomID, userID, membership)
|
||||
} else {
|
||||
err = fmt.Errorf("unsupported dialect %s", store.db.dialect)
|
||||
}
|
||||
if err != nil {
|
||||
store.log.Warnfln("Failed to set membership of %s in %s to %s: %v", userID, roomID, membership, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) SetPowerLevels(roomID string, levels *mautrix.PowerLevels) {
|
||||
levelsBytes, err := json.Marshal(levels)
|
||||
if err != nil {
|
||||
store.log.Errorfln("Failed to marshal power levels of %s: %v", roomID, err)
|
||||
return
|
||||
}
|
||||
if store.db.dialect == "postgres" {
|
||||
_, err = store.db.Exec(`INSERT INTO mx_room_state (room_id, power_levels) VALUES ($1, $2)
|
||||
ON CONFLICT (room_id) DO UPDATE SET power_levels=$2`, roomID, levelsBytes)
|
||||
} else if store.db.dialect == "sqlite3" {
|
||||
_, err = store.db.Exec("INSERT OR REPLACE INTO mx_room_state (room_id, power_levels) VALUES ($1, $2)", roomID, levelsBytes)
|
||||
} else {
|
||||
err = fmt.Errorf("unsupported dialect %s", store.db.dialect)
|
||||
}
|
||||
if err != nil {
|
||||
store.log.Warnfln("Failed to store power levels of %s: %v", roomID, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) GetPowerLevels(roomID string) (levels *mautrix.PowerLevels) {
|
||||
row := store.db.QueryRow("SELECT power_levels FROM mx_room_state WHERE room_id=$1", roomID)
|
||||
if row == nil {
|
||||
return
|
||||
}
|
||||
var data []byte
|
||||
err := row.Scan(&data)
|
||||
if err != nil {
|
||||
store.log.Errorln("Failed to scan power levels of %s: %v", roomID, err)
|
||||
return
|
||||
}
|
||||
levels = &mautrix.PowerLevels{}
|
||||
err = json.Unmarshal(data, levels)
|
||||
if err != nil {
|
||||
store.log.Errorln("Failed to parse power levels of %s: %v", roomID, err)
|
||||
return nil
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) GetPowerLevel(roomID, userID string) int {
|
||||
if store.db.dialect == "postgres" {
|
||||
row := store.db.QueryRow(`SELECT
|
||||
COALESCE((power_levels->'users'->$2)::int, (power_levels->'users_default')::int, 0)
|
||||
FROM mx_room_state WHERE room_id=$1`, roomID, userID)
|
||||
if row == nil {
|
||||
// Power levels not in db
|
||||
return 0
|
||||
}
|
||||
var powerLevel int
|
||||
err := row.Scan(&powerLevel)
|
||||
if err != nil {
|
||||
store.log.Errorln("Failed to scan power level of %s in %s: %v", userID, roomID, err)
|
||||
}
|
||||
return powerLevel
|
||||
}
|
||||
return store.GetPowerLevels(roomID).GetUserLevel(userID)
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) GetPowerLevelRequirement(roomID string, eventType mautrix.EventType) int {
|
||||
if store.db.dialect == "postgres" {
|
||||
defaultType := "events_default"
|
||||
defaultValue := 0
|
||||
if eventType.IsState() {
|
||||
defaultType = "state_default"
|
||||
defaultValue = 50
|
||||
}
|
||||
row := store.db.QueryRow(`SELECT
|
||||
COALESCE((power_levels->'events'->$2)::int, (power_levels->'$3')::int, $4)
|
||||
FROM mx_room_state WHERE room_id=$1`, roomID, eventType.Type, defaultType, defaultValue)
|
||||
if row == nil {
|
||||
// Power levels not in db
|
||||
return defaultValue
|
||||
}
|
||||
var powerLevel int
|
||||
err := row.Scan(&powerLevel)
|
||||
if err != nil {
|
||||
store.log.Errorln("Failed to scan power level for %s in %s: %v", eventType, roomID, err)
|
||||
}
|
||||
return powerLevel
|
||||
}
|
||||
return store.GetPowerLevels(roomID).GetEventLevel(eventType)
|
||||
}
|
||||
|
||||
func (store *SQLStateStore) HasPowerLevel(roomID, userID string, eventType mautrix.EventType) bool {
|
||||
if store.db.dialect == "postgres" {
|
||||
defaultType := "events_default"
|
||||
defaultValue := 0
|
||||
if eventType.IsState() {
|
||||
defaultType = "state_default"
|
||||
defaultValue = 50
|
||||
}
|
||||
row := store.db.QueryRow(`SELECT
|
||||
COALESCE((power_levels->'users'->$2)::int, (power_levels->'users_default')::int, 0)
|
||||
>=
|
||||
COALESCE((power_levels->'events'->$3)::int, (power_levels->'$4')::int, $5)
|
||||
FROM mx_room_state WHERE room_id=$1`, roomID, userID, eventType.Type, defaultType, defaultValue)
|
||||
if row == nil {
|
||||
// Power levels not in db
|
||||
return defaultValue == 0
|
||||
}
|
||||
var hasPower bool
|
||||
err := row.Scan(&hasPower)
|
||||
if err != nil {
|
||||
store.log.Errorln("Failed to scan power level for %s in %s: %v", eventType, roomID, err)
|
||||
}
|
||||
return hasPower
|
||||
}
|
||||
return store.GetPowerLevel(roomID, userID) >= store.GetPowerLevelRequirement(roomID, eventType)
|
||||
}
|
||||
107
database/upgrades/2019-08-25-move-state-store-to-db.go
Normal file
107
database/upgrades/2019-08-25-move-state-store-to-db.go
Normal file
@ -0,0 +1,107 @@
|
||||
package upgrades
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"strings"
|
||||
|
||||
"maunium.net/go/mautrix"
|
||||
"maunium.net/go/mautrix-appservice"
|
||||
)
|
||||
|
||||
func init() {
|
||||
migrateRegistrations := func(tx *sql.Tx, registrations map[string]bool) error {
|
||||
values := make([]interface{}, 0, len(registrations))
|
||||
valueStrings := make([]string, 0, len(registrations))
|
||||
i := 1
|
||||
for userID, registered := range registrations {
|
||||
if registered {
|
||||
values = append(values, userID)
|
||||
valueStrings = append(valueStrings, fmt.Sprintf("($%d)", i))
|
||||
i++
|
||||
}
|
||||
}
|
||||
valueString := strings.Join(valueStrings, ",")
|
||||
_, err := tx.Exec("INSERT INTO mx_registrations (user_id) VALUES "+valueString, values...)
|
||||
return err
|
||||
}
|
||||
|
||||
migrateMemberships := func(tx *sql.Tx, rooms map[string]map[string]mautrix.Membership) error {
|
||||
var values []interface{}
|
||||
var valueStrings []string
|
||||
i := 1
|
||||
for roomID, members := range rooms {
|
||||
for userID, membership := range members {
|
||||
values = append(values, roomID, userID, membership)
|
||||
valueStrings = append(valueStrings, fmt.Sprintf("($%d, $%d, $%d)", i, i+1, i+2))
|
||||
i += 3
|
||||
}
|
||||
}
|
||||
valueString := strings.Join(valueStrings, ",")
|
||||
_, err := tx.Exec("INSERT INTO mx_user_profile (room_id, user_id, membership) VALUES "+valueString, values...)
|
||||
return err
|
||||
}
|
||||
|
||||
migratePowerLevels := func(tx *sql.Tx, rooms map[string]*mautrix.PowerLevels) error {
|
||||
var values []interface{}
|
||||
var valueStrings []string
|
||||
i := 1
|
||||
for roomID, powerLevels := range rooms {
|
||||
powerLevelBytes, err := json.Marshal(powerLevels)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
values = append(values, roomID, powerLevelBytes)
|
||||
valueStrings = append(valueStrings, fmt.Sprintf("($%d, $%d)", i, i+1))
|
||||
i += 2
|
||||
}
|
||||
valueString := strings.Join(valueStrings, ",")
|
||||
_, err := tx.Exec("INSERT INTO mx_room_state (room_id, power_levels) VALUES "+valueString, values...)
|
||||
return err
|
||||
}
|
||||
|
||||
userProfileTable := `CREATE TABLE mx_user_profile (
|
||||
room_id VARCHAR(255),
|
||||
user_id VARCHAR(255),
|
||||
membership VARCHAR(15) NOT NULL,
|
||||
PRIMARY KEY (room_id, user_id)
|
||||
)`
|
||||
|
||||
roomStateTable := `CREATE TABLE mx_room_state (
|
||||
room_id VARCHAR(255) PRIMARY KEY,
|
||||
power_levels TEXT
|
||||
)`
|
||||
|
||||
registrationsTable := `CREATE TABLE mx_registrations (
|
||||
user_id VARCHAR(255) PRIMARY KEY
|
||||
)`
|
||||
|
||||
upgrades[9] = upgrade{"Move state store to main DB", func(dialect Dialect, tx *sql.Tx, db *sql.DB) error {
|
||||
store := appservice.NewBasicStateStore().(*appservice.BasicStateStore)
|
||||
|
||||
if dialect == Postgres {
|
||||
roomStateTable = strings.Replace(roomStateTable, "TEXT", "JSONB", 1)
|
||||
}
|
||||
|
||||
if data, err := ioutil.ReadFile("mx-state.json"); err != nil {
|
||||
return err
|
||||
} else if err = json.Unmarshal(data, &store); err != nil {
|
||||
return err
|
||||
} else if _, err := tx.Exec(userProfileTable); err != nil {
|
||||
return err
|
||||
} else if _, err = tx.Exec(roomStateTable); err != nil {
|
||||
return err
|
||||
} else if _, err = tx.Exec(registrationsTable); err != nil {
|
||||
return err
|
||||
} else if err = migrateRegistrations(tx, store.Registrations); err != nil {
|
||||
return err
|
||||
} else if err = migrateMemberships(tx, store.Memberships); err != nil {
|
||||
return err
|
||||
} else if err = migratePowerLevels(tx, store.PowerLevels); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}}
|
||||
}
|
||||
@ -22,7 +22,7 @@ type upgrade struct {
|
||||
fn upgradeFunc
|
||||
}
|
||||
|
||||
const NumberOfUpgrades = 9
|
||||
const NumberOfUpgrades = 10
|
||||
|
||||
var upgrades [NumberOfUpgrades]upgrade
|
||||
|
||||
|
||||
Reference in New Issue
Block a user