source: code/trunk/db_postgres.go@ 774

package soju

import (
        "context"
        "database/sql"
        "errors"
        "fmt"
        "math"
        "strings"
        "time"

        _ "github.com/lib/pq"
        "github.com/prometheus/client_golang/prometheus"
        promcollectors "github.com/prometheus/client_golang/prometheus/collectors"
)

const postgresQueryTimeout = 5 * time.Second

const postgresConfigSchema = `
CREATE TABLE IF NOT EXISTS "Config" (
        id SMALLINT PRIMARY KEY,
        version INTEGER NOT NULL,
        CHECK(id = 1)
);
`

const postgresSchema = `
CREATE TABLE "User" (
        id SERIAL PRIMARY KEY,
        username VARCHAR(255) NOT NULL UNIQUE,
        password VARCHAR(255),
        admin BOOLEAN NOT NULL DEFAULT FALSE,
        realname VARCHAR(255)
);

CREATE TYPE sasl_mechanism AS ENUM ('PLAIN', 'EXTERNAL');

CREATE TABLE "Network" (
        id SERIAL PRIMARY KEY,
        name VARCHAR(255),
        "user" INTEGER NOT NULL REFERENCES "User"(id) ON DELETE CASCADE,
        addr VARCHAR(255) NOT NULL,
        nick VARCHAR(255),
        username VARCHAR(255),
        realname VARCHAR(255),
        pass VARCHAR(255),
        connect_commands VARCHAR(1023),
        sasl_mechanism sasl_mechanism,
        sasl_plain_username VARCHAR(255),
        sasl_plain_password VARCHAR(255),
        sasl_external_cert BYTEA,
        sasl_external_key BYTEA,
        enabled BOOLEAN NOT NULL DEFAULT TRUE,
        UNIQUE("user", addr, nick),
        UNIQUE("user", name)
);

CREATE TABLE "Channel" (
        id SERIAL PRIMARY KEY,
        network INTEGER NOT NULL REFERENCES "Network"(id) ON DELETE CASCADE,
        name VARCHAR(255) NOT NULL,
        key VARCHAR(255),
        detached BOOLEAN NOT NULL DEFAULT FALSE,
        detached_internal_msgid VARCHAR(255),
        relay_detached INTEGER NOT NULL DEFAULT 0,
        reattach_on INTEGER NOT NULL DEFAULT 0,
        detach_after INTEGER NOT NULL DEFAULT 0,
        detach_on INTEGER NOT NULL DEFAULT 0,
        UNIQUE(network, name)
);

CREATE TABLE "DeliveryReceipt" (
        id SERIAL PRIMARY KEY,
        network INTEGER NOT NULL REFERENCES "Network"(id) ON DELETE CASCADE,
        target VARCHAR(255) NOT NULL,
        client VARCHAR(255) NOT NULL DEFAULT '',
        internal_msgid VARCHAR(255) NOT NULL,
        UNIQUE(network, target, client)
);
`

var postgresMigrations = []string{
        "", // migration #0 is reserved for schema initialization
        `ALTER TABLE "Network" ALTER COLUMN nick DROP NOT NULL`,
        `
                CREATE TYPE sasl_mechanism AS ENUM ('PLAIN', 'EXTERNAL');
                ALTER TABLE "Network"
                        ALTER COLUMN sasl_mechanism
                        TYPE sasl_mechanism
                        USING sasl_mechanism::sasl_mechanism;
        `,
}

type PostgresDB struct {
        db *sql.DB
}

func OpenPostgresDB(source string) (Database, error) {
        sqlPostgresDB, err := sql.Open("postgres", source)
        if err != nil {
                return nil, err
        }

        db := &PostgresDB{db: sqlPostgresDB}
        if err := db.upgrade(); err != nil {
                sqlPostgresDB.Close()
                return nil, err
        }

        return db, nil
}

func (db *PostgresDB) upgrade() error {
        tx, err := db.db.Begin()
        if err != nil {
                return err
        }
        defer tx.Rollback()

        if _, err := tx.Exec(postgresConfigSchema); err != nil {
                return fmt.Errorf("failed to create Config table: %s", err)
        }

        var version int
        err = tx.QueryRow(`SELECT version FROM "Config"`).Scan(&version)
        if err != nil && !errors.Is(err, sql.ErrNoRows) {
                return fmt.Errorf("failed to query schema version: %s", err)
        }

        if version == len(postgresMigrations) {
                return nil
        }
        if version > len(postgresMigrations) {
                return fmt.Errorf("soju (version %d) older than schema (version %d)", len(postgresMigrations), version)
        }

        if version == 0 {
                if _, err := tx.Exec(postgresSchema); err != nil {
                        return fmt.Errorf("failed to initialize schema: %s", err)
                }
        } else {
                for i := version; i < len(postgresMigrations); i++ {
                        if _, err := tx.Exec(postgresMigrations[i]); err != nil {
                                return fmt.Errorf("failed to execute migration #%v: %v", i, err)
                        }
                }
        }

        _, err = tx.Exec(`INSERT INTO "Config" (id, version) VALUES (1, $1)
                ON CONFLICT (id) DO UPDATE SET version = $1`, len(postgresMigrations))
        if err != nil {
                return fmt.Errorf("failed to bump schema version: %v", err)
        }

        return tx.Commit()
}

func (db *PostgresDB) Close() error {
        return db.db.Close()
}

func (db *PostgresDB) MetricsCollector() prometheus.Collector {
        return promcollectors.NewDBStatsCollector(db.db, "main")
}

func (db *PostgresDB) Stats(ctx context.Context) (*DatabaseStats, error) {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        var stats DatabaseStats
        row := db.db.QueryRowContext(ctx, `SELECT
                (SELECT COUNT(*) FROM "User") AS users,
                (SELECT COUNT(*) FROM "Network") AS networks,
                (SELECT COUNT(*) FROM "Channel") AS channels`)
        if err := row.Scan(&stats.Users, &stats.Networks, &stats.Channels); err != nil {
                return nil, err
        }

        return &stats, nil
}

func (db *PostgresDB) ListUsers(ctx context.Context) ([]User, error) {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        rows, err := db.db.QueryContext(ctx,
                `SELECT id, username, password, admin, realname FROM "User"`)
        if err != nil {
                return nil, err
        }
        defer rows.Close()

        var users []User
        for rows.Next() {
                var user User
                var password, realname sql.NullString
                if err := rows.Scan(&user.ID, &user.Username, &password, &user.Admin, &realname); err != nil {
                        return nil, err
                }
                user.Password = password.String
                user.Realname = realname.String
                users = append(users, user)
        }
        if err := rows.Err(); err != nil {
                return nil, err
        }

        return users, nil
}

func (db *PostgresDB) GetUser(ctx context.Context, username string) (*User, error) {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        user := &User{Username: username}

        var password, realname sql.NullString
        row := db.db.QueryRowContext(ctx,
                `SELECT id, password, admin, realname FROM "User" WHERE username = $1`,
                username)
        if err := row.Scan(&user.ID, &password, &user.Admin, &realname); err != nil {
                return nil, err
        }
        user.Password = password.String
        user.Realname = realname.String
        return user, nil
}

func (db *PostgresDB) StoreUser(ctx context.Context, user *User) error {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        password := toNullString(user.Password)
        realname := toNullString(user.Realname)

        var err error
        if user.ID == 0 {
                err = db.db.QueryRowContext(ctx, `
                        INSERT INTO "User" (username, password, admin, realname)
                        VALUES ($1, $2, $3, $4)
                        RETURNING id`,
                        user.Username, password, user.Admin, realname).Scan(&user.ID)
        } else {
                _, err = db.db.ExecContext(ctx, `
                        UPDATE "User"
                        SET password = $1, admin = $2, realname = $3
                        WHERE id = $4`,
                        password, user.Admin, realname, user.ID)
        }
        return err
}

func (db *PostgresDB) DeleteUser(ctx context.Context, id int64) error {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        _, err := db.db.ExecContext(ctx, `DELETE FROM "User" WHERE id = $1`, id)
        return err
}

func (db *PostgresDB) ListNetworks(ctx context.Context, userID int64) ([]Network, error) {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        rows, err := db.db.QueryContext(ctx, `
                SELECT id, name, addr, nick, username, realname, pass, connect_commands, sasl_mechanism,
                        sasl_plain_username, sasl_plain_password, sasl_external_cert, sasl_external_key, enabled
                FROM "Network"
                WHERE "user" = $1`, userID)
        if err != nil {
                return nil, err
        }
        defer rows.Close()

        var networks []Network
        for rows.Next() {
                var net Network
                var name, nick, username, realname, pass, connectCommands sql.NullString
                var saslMechanism, saslPlainUsername, saslPlainPassword sql.NullString
                err := rows.Scan(&net.ID, &name, &net.Addr, &nick, &username, &realname,
                        &pass, &connectCommands, &saslMechanism, &saslPlainUsername, &saslPlainPassword,
                        &net.SASL.External.CertBlob, &net.SASL.External.PrivKeyBlob, &net.Enabled)
                if err != nil {
                        return nil, err
                }
                net.Name = name.String
                net.Nick = nick.String
                net.Username = username.String
                net.Realname = realname.String
                net.Pass = pass.String
                if connectCommands.Valid {
                        net.ConnectCommands = strings.Split(connectCommands.String, "\r\n")
                }
                net.SASL.Mechanism = saslMechanism.String
                net.SASL.Plain.Username = saslPlainUsername.String
                net.SASL.Plain.Password = saslPlainPassword.String
                networks = append(networks, net)
        }
        if err := rows.Err(); err != nil {
                return nil, err
        }

        return networks, nil
}

func (db *PostgresDB) StoreNetwork(ctx context.Context, userID int64, network *Network) error {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        netName := toNullString(network.Name)
        nick := toNullString(network.Nick)
        netUsername := toNullString(network.Username)
        realname := toNullString(network.Realname)
        pass := toNullString(network.Pass)
        connectCommands := toNullString(strings.Join(network.ConnectCommands, "\r\n"))

        var saslMechanism, saslPlainUsername, saslPlainPassword sql.NullString
        if network.SASL.Mechanism != "" {
                saslMechanism = toNullString(network.SASL.Mechanism)
                switch network.SASL.Mechanism {
                case "PLAIN":
                        saslPlainUsername = toNullString(network.SASL.Plain.Username)
                        saslPlainPassword = toNullString(network.SASL.Plain.Password)
                        network.SASL.External.CertBlob = nil
                        network.SASL.External.PrivKeyBlob = nil
                case "EXTERNAL":
                        // keep saslPlain* nil
                default:
                        return fmt.Errorf("soju: cannot store network: unsupported SASL mechanism %q", network.SASL.Mechanism)
                }
        }

        var err error
        if network.ID == 0 {
                err = db.db.QueryRowContext(ctx, `
                        INSERT INTO "Network" ("user", name, addr, nick, username, realname, pass, connect_commands,
                                sasl_mechanism, sasl_plain_username, sasl_plain_password, sasl_external_cert,
                                sasl_external_key, enabled)
                        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
                        RETURNING id`,
                        userID, netName, network.Addr, nick, netUsername, realname, pass, connectCommands,
                        saslMechanism, saslPlainUsername, saslPlainPassword, network.SASL.External.CertBlob,
                        network.SASL.External.PrivKeyBlob, network.Enabled).Scan(&network.ID)
        } else {
                _, err = db.db.ExecContext(ctx, `
                        UPDATE "Network"
                        SET name = $2, addr = $3, nick = $4, username = $5, realname = $6, pass = $7,
                                connect_commands = $8, sasl_mechanism = $9, sasl_plain_username = $10,
                                sasl_plain_password = $11, sasl_external_cert = $12, sasl_external_key = $13,
                                enabled = $14
                        WHERE id = $1`,
                        network.ID, netName, network.Addr, nick, netUsername, realname, pass, connectCommands,
                        saslMechanism, saslPlainUsername, saslPlainPassword, network.SASL.External.CertBlob,
                        network.SASL.External.PrivKeyBlob, network.Enabled)
        }
        return err
}

func (db *PostgresDB) DeleteNetwork(ctx context.Context, id int64) error {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        _, err := db.db.ExecContext(ctx, `DELETE FROM "Network" WHERE id = $1`, id)
        return err
}

func (db *PostgresDB) ListChannels(ctx context.Context, networkID int64) ([]Channel, error) {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        rows, err := db.db.QueryContext(ctx, `
                SELECT id, name, key, detached, detached_internal_msgid, relay_detached, reattach_on, detach_after,
                        detach_on
                FROM "Channel"
                WHERE network = $1`, networkID)
        if err != nil {
                return nil, err
        }
        defer rows.Close()

        var channels []Channel
        for rows.Next() {
                var ch Channel
                var key, detachedInternalMsgID sql.NullString
                var detachAfter int64
                if err := rows.Scan(&ch.ID, &ch.Name, &key, &ch.Detached, &detachedInternalMsgID, &ch.RelayDetached, &ch.ReattachOn, &detachAfter, &ch.DetachOn); err != nil {
                        return nil, err
                }
                ch.Key = key.String
                ch.DetachedInternalMsgID = detachedInternalMsgID.String
                ch.DetachAfter = time.Duration(detachAfter) * time.Second
                channels = append(channels, ch)
        }
        if err := rows.Err(); err != nil {
                return nil, err
        }

        return channels, nil
}

func (db *PostgresDB) StoreChannel(ctx context.Context, networkID int64, ch *Channel) error {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        key := toNullString(ch.Key)
        detachAfter := int64(math.Ceil(ch.DetachAfter.Seconds()))

        var err error
        if ch.ID == 0 {
                err = db.db.QueryRowContext(ctx, `
                        INSERT INTO "Channel" (network, name, key, detached, detached_internal_msgid, relay_detached, reattach_on,
                                detach_after, detach_on)
                        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
                        RETURNING id`,
                        networkID, ch.Name, key, ch.Detached, toNullString(ch.DetachedInternalMsgID),
                        ch.RelayDetached, ch.ReattachOn, detachAfter, ch.DetachOn).Scan(&ch.ID)
        } else {
                _, err = db.db.ExecContext(ctx, `
                        UPDATE "Channel"
                        SET name = $2, key = $3, detached = $4, detached_internal_msgid = $5,
                                relay_detached = $6, reattach_on = $7, detach_after = $8, detach_on = $9
                        WHERE id = $1`,
                        ch.ID, ch.Name, key, ch.Detached, toNullString(ch.DetachedInternalMsgID),
                        ch.RelayDetached, ch.ReattachOn, detachAfter, ch.DetachOn)
        }
        return err
}

func (db *PostgresDB) DeleteChannel(ctx context.Context, id int64) error {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        _, err := db.db.ExecContext(ctx, `DELETE FROM "Channel" WHERE id = $1`, id)
        return err
}

func (db *PostgresDB) ListDeliveryReceipts(ctx context.Context, networkID int64) ([]DeliveryReceipt, error) {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        rows, err := db.db.QueryContext(ctx, `
                SELECT id, target, client, internal_msgid
                FROM "DeliveryReceipt"
                WHERE network = $1`, networkID)
        if err != nil {
                return nil, err
        }
        defer rows.Close()

        var receipts []DeliveryReceipt
        for rows.Next() {
                var rcpt DeliveryReceipt
                if err := rows.Scan(&rcpt.ID, &rcpt.Target, &rcpt.Client, &rcpt.InternalMsgID); err != nil {
                        return nil, err
                }
                receipts = append(receipts, rcpt)
        }
        if err := rows.Err(); err != nil {
                return nil, err
        }

        return receipts, nil
}

func (db *PostgresDB) StoreClientDeliveryReceipts(ctx context.Context, networkID int64, client string, receipts []DeliveryReceipt) error {
        ctx, cancel := context.WithTimeout(ctx, postgresQueryTimeout)
        defer cancel()

        tx, err := db.db.Begin()
        if err != nil {
                return err
        }
        defer tx.Rollback()

        _, err = tx.ExecContext(ctx,
                `DELETE FROM "DeliveryReceipt" WHERE network = $1 AND client = $2`,
                networkID, client)
        if err != nil {
                return err
        }

        stmt, err := tx.PrepareContext(ctx, `
                INSERT INTO "DeliveryReceipt" (network, target, client, internal_msgid)
                VALUES ($1, $2, $3, $4)
                RETURNING id`)
        if err != nil {
                return err
        }
        defer stmt.Close()

        for i := range receipts {
                rcpt := &receipts[i]
                err := stmt.
                        QueryRowContext(ctx, networkID, rcpt.Target, client, rcpt.InternalMsgID).
                        Scan(&rcpt.ID)
                if err != nil {
                        return err
                }
        }

        return tx.Commit()
}