mesh-catalog/modules/mesh-mail/cmd/mesh-mail/pgstore.go

package main

// The index on the mesh's postgres (novox/hq ADR 0359 §3): the database the mesh provisions for this module through
// the postgres-database provision, its address in the file DATABASE_URL_FILE names. Its schema is numbered migrations
// in migrations/, applied by the preparation step (`mesh-mail prepare`, ADR 0135) and recorded in the database itself;
// a fresh database is given the consolidated schema, schema.sql, at once, with every migration recorded as applied.
// Nothing else changes the schema: the keeper, the watch and the tools refuse a schema that is behind, and say so.

import (
	"context"
	"crypto/sha256"
	"embed"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

//go:embed migrations/*.sql
var migrationFiles embed.FS

//go:embed schema.sql
var consolidated string

// Migration is one numbered change of the schema.
type Migration struct {
	Number   int
	Name     string
	SQL      string
	Checksum string
}

var migrationFile = regexp.MustCompile(`^(\d{4})-([a-z0-9-]+)\.sql$`)

// Migrations are the migrations this binary carries, in order, numbered from 0001 without a gap.
func Migrations() ([]Migration, error) {
	entries, err := fs.ReadDir(migrationFiles, "migrations")
	if err != nil {
		return nil, err
	}
	var out []Migration
	for _, e := range entries {
		m := migrationFile.FindStringSubmatch(e.Name())
		if m == nil {
			return nil, fmt.Errorf("%s is not a migration's name (NNNN-what-it-does.sql)", e.Name())
		}
		n, _ := strconv.Atoi(m[1])
		body, err := fs.ReadFile(migrationFiles, path.Join("migrations", e.Name()))
		if err != nil {
			return nil, err
		}
		if strings.TrimSpace(string(body)) == "" {
			return nil, fmt.Errorf("%s is empty", e.Name())
		}
		sum := sha256.Sum256(body)
		out = append(out, Migration{Number: n, Name: m[2], SQL: string(body), Checksum: hex.EncodeToString(sum[:])})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Number < out[j].Number })
	for i, m := range out {
		if m.Number != i+1 {
			return nil, fmt.Errorf("the migrations are not numbered 0001 on without a gap: %04d is the %d. one", m.Number, i+1)
		}
	}
	return out, nil
}

const ledger = `create table if not exists migration (
    number   integer     primary key,
    name     text        not null,
    checksum text        not null,
    applied  timestamptz not null default now())`

// lockKey is the advisory lock every migration run takes, so two preparations never interleave.
const lockKey int64 = 4_250_359_001

// PgStore is the Store on postgres.
type PgStore struct{ pool *pgxpool.Pool }

// PgStoreFromEnv opens the index the mesh provisioned.
func PgStoreFromEnv(ctx context.Context) (*PgStore, error) {
	file := os.Getenv("DATABASE_URL_FILE")
	if file == "" {
		return nil, errors.New("DATABASE_URL_FILE is not set: the index's database is a requirement the mesh resolves")
	}
	raw, err := os.ReadFile(file)
	if err != nil {
		return nil, fmt.Errorf("the database's address cannot be read: %w", err)
	}
	return OpenPgStore(ctx, strings.TrimSpace(string(raw)))
}

// OpenPgStore opens the index at an address.
func OpenPgStore(ctx context.Context, url string) (*PgStore, error) {
	pool, err := pgxpool.New(ctx, url)
	if err != nil {
		return nil, err
	}
	return &PgStore{pool: pool}, nil
}

// Close closes the store.
func (s *PgStore) Close() { s.pool.Close() }

// Prepare brings the schema to what this version needs, and says what it did.
func (s *PgStore) Prepare(ctx context.Context) ([]string, error) {
	migrations, err := Migrations()
	if err != nil {
		return nil, err
	}
	conn, err := s.pool.Acquire(ctx)
	if err != nil {
		return nil, err
	}
	defer conn.Release()
	if _, err := conn.Exec(ctx, `select pg_advisory_lock($1)`, lockKey); err != nil {
		return nil, fmt.Errorf("the migration lock cannot be taken: %w", err)
	}
	defer func() { _, _ = conn.Exec(context.WithoutCancel(ctx), `select pg_advisory_unlock($1)`, lockKey) }()
	if _, err := conn.Exec(ctx, ledger); err != nil {
		return nil, err
	}
	applied := map[int]string{}
	rows, err := conn.Query(ctx, `select number, checksum from migration`)
	if err != nil {
		return nil, err
	}
	highest := 0
	for rows.Next() {
		var n int
		var sum string
		if err := rows.Scan(&n, &sum); err != nil {
			rows.Close()
			return nil, err
		}
		applied[n] = sum
		highest = max(highest, n)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return nil, err
	}
	record := func(tx pgx.Tx, m Migration) error {
		_, err := tx.Exec(ctx, `insert into migration (number, name, checksum) values ($1, $2, $3)`, m.Number, m.Name, m.Checksum)
		return err
	}
	if len(applied) == 0 {
		var exists bool
		if err := conn.QueryRow(ctx, `select to_regclass('message') is not null`).Scan(&exists); err != nil {
			return nil, err
		}
		if exists {
			return nil, errors.New("the database holds the index's tables and no record of a migration: its contents are not known, so nothing is changed")
		}
		tx, err := conn.Begin(ctx)
		if err != nil {
			return nil, err
		}
		defer func() { _ = tx.Rollback(context.WithoutCancel(ctx)) }()
		if _, err := tx.Exec(ctx, consolidated); err != nil {
			return nil, fmt.Errorf("the consolidated schema failed, and nothing of it was kept: %w", err)
		}
		for _, m := range migrations {
			if err := record(tx, m); err != nil {
				return nil, err
			}
		}
		if err := tx.Commit(ctx); err != nil {
			return nil, err
		}
		return []string{fmt.Sprintf("a fresh database: the consolidated schema, through migration %04d", len(migrations))}, nil
	}
	var problems []string
	var pending []Migration
	for _, m := range migrations {
		sum, ran := applied[m.Number]
		switch {
		case ran && sum != m.Checksum:
			problems = append(problems, fmt.Sprintf("migration %04d-%s ran, and its file has changed since: write a new one", m.Number, m.Name))
		case !ran && m.Number < highest:
			problems = append(problems, fmt.Sprintf("migration %04d-%s never ran, and %04d did: renumber it above %04d", m.Number, m.Name, highest, highest))
		case !ran:
			pending = append(pending, m)
		}
	}
	if len(problems) > 0 {
		return nil, errors.New("this database and these migrations disagree: " + strings.Join(problems, "; "))
	}
	var said []string
	for _, m := range pending {
		err := func() error {
			tx, err := conn.Begin(ctx)
			if err != nil {
				return err
			}
			defer func() { _ = tx.Rollback(context.WithoutCancel(ctx)) }()
			if _, err := tx.Exec(ctx, m.SQL); err != nil {
				return fmt.Errorf("migration %04d-%s failed, and nothing it did was kept: %w", m.Number, m.Name, err)
			}
			if err := record(tx, m); err != nil {
				return err
			}
			return tx.Commit(ctx)
		}()
		if err != nil {
			return said, err
		}
		said = append(said, fmt.Sprintf("migration %04d-%s applied", m.Number, m.Name))
	}
	if len(said) == 0 {
		said = []string{fmt.Sprintf("the schema is at migration %04d, as this version needs", len(migrations))}
	}
	return said, nil
}

// Ready says why the index cannot serve this version, or nil.
func (s *PgStore) Ready(ctx context.Context) error {
	migrations, err := Migrations()
	if err != nil {
		return err
	}
	var at int
	err = s.pool.QueryRow(ctx, `select coalesce(max(number), 0) from migration`).Scan(&at)
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) && pgErr.Code == "42P01" {
		at = 0
	} else if err != nil {
		return fmt.Errorf("the database cannot be reached: %w", err)
	}
	if at != len(migrations) {
		return fmt.Errorf("the database's schema is at migration %04d and this version needs %04d: the preparation step "+
			"(mesh-mail prepare) has not run", at, len(migrations))
	}
	return nil
}

func (s *PgStore) Folder(ctx context.Context, account, folder string) (Folder, bool, error) {
	f := Folder{Account: account, Name: folder}
	var validity int64
	err := s.pool.QueryRow(ctx, `select uidvalidity, messages, synced from folder where account = $1 and name = $2`, account, folder).
		Scan(&validity, &f.Messages, &f.Synced)
	if errors.Is(err, pgx.ErrNoRows) {
		return f, false, nil
	}
	f.UIDValidity = uint32(validity)
	return f, err == nil, err
}

func (s *PgStore) SetFolder(ctx context.Context, f Folder) error {
	_, err := s.pool.Exec(ctx, `insert into folder (account, name, uidvalidity, messages, synced) values ($1, $2, $3, $4, $5)
		on conflict (account, name) do update set uidvalidity = excluded.uidvalidity, messages = excluded.messages, synced = excluded.synced`,
		f.Account, f.Name, int64(f.UIDValidity), f.Messages, f.Synced)
	return err
}

func (s *PgStore) DropFolder(ctx context.Context, account, folder string) error {
	return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
		if _, err := tx.Exec(ctx, `delete from message where account = $1 and folder = $2`, account, folder); err != nil {
			return err
		}
		_, err := tx.Exec(ctx, `delete from folder where account = $1 and name = $2`, account, folder)
		return err
	})
}

func (s *PgStore) Flags(ctx context.Context, account, folder string) (map[uint32][]string, error) {
	rows, err := s.pool.Query(ctx, `select uid, flags from message where account = $1 and folder = $2`, account, folder)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[uint32][]string{}
	for rows.Next() {
		var uid int64
		var flags []string
		if err := rows.Scan(&uid, &flags); err != nil {
			return nil, err
		}
		out[uint32(uid)] = flags
	}
	return out, rows.Err()
}

func keepEvent(ctx context.Context, tx pgx.Tx, e *Event) error {
	if e == nil {
		return nil
	}
	body, err := json.Marshal(e.Body)
	if err != nil {
		return err
	}
	_, err = tx.Exec(ctx, `insert into outbox (key, body) values ($1, $2)`, e.Key, body)
	return err
}

func (s *PgStore) Add(ctx context.Context, m Message, e *Event) error {
	from, _ := json.Marshal(m.From)
	to, _ := json.Marshal(m.To)
	cc, _ := json.Marshal(orNone(m.Cc))
	atts, _ := json.Marshal(m.Attachments)
	flags := m.Flags
	if flags == nil {
		flags = []string{}
	}
	people := strings.Join([]string{peopleText(m.From), peopleText(m.To), peopleText(m.Cc)}, " ")
	// An address is one word to the search, so its parts are kept beside it: "bob" finds bob@example.org.
	people += " " + strings.NewReplacer("@", " ", ".", " ", ",", " ").Replace(people)
	return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
		if _, err := tx.Exec(ctx, `insert into message (account, folder, uid, uidvalidity, message_id, in_reply_to, date, sender,
			recipients, copied, people, subject, size, flags, attachments, body)
			values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
			on conflict (account, folder, uid) do update set uidvalidity = excluded.uidvalidity, message_id = excluded.message_id,
			in_reply_to = excluded.in_reply_to, date = excluded.date, sender = excluded.sender, recipients = excluded.recipients,
			copied = excluded.copied, people = excluded.people, subject = excluded.subject, size = excluded.size,
			flags = excluded.flags, attachments = excluded.attachments, body = excluded.body`,
			m.Account, m.Folder, int64(m.UID), int64(m.UIDValidity), m.MessageID, m.InReplyTo, m.Date, from, to, cc, people,
			strings.ToValidUTF8(strings.ReplaceAll(m.Subject, "\x00", ""), ""), m.Size, flags, atts,
			strings.ReplaceAll(m.Body, "\x00", "")); err != nil {
			return err
		}
		return keepEvent(ctx, tx, e)
	})
}

func orNone(ps []Person) []Person {
	if ps == nil {
		return []Person{}
	}
	return ps
}

func (s *PgStore) SetFlags(ctx context.Context, account, folder string, uid uint32, flags []string, e *Event) error {
	return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
		if _, err := tx.Exec(ctx, `update message set flags = $4 where account = $1 and folder = $2 and uid = $3`,
			account, folder, int64(uid), flags); err != nil {
			return err
		}
		return keepEvent(ctx, tx, e)
	})
}

func (s *PgStore) Remove(ctx context.Context, account, folder string, uid uint32, e *Event) error {
	return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
		if _, err := tx.Exec(ctx, `delete from message where account = $1 and folder = $2 and uid = $3`, account, folder, int64(uid)); err != nil {
			return err
		}
		return keepEvent(ctx, tx, e)
	})
}

const messageColumns = `account, folder, uid, uidvalidity, message_id, in_reply_to, date, sender, recipients, copied, subject,
	size, flags, attachments`

func scanMessage(row pgx.Row, withBody bool) (Message, error) {
	var m Message
	var uid, validity int64
	var from, to, cc, atts []byte
	dest := []any{&m.Account, &m.Folder, &uid, &validity, &m.MessageID, &m.InReplyTo, &m.Date, &from, &to, &cc, &m.Subject,
		&m.Size, &m.Flags, &atts}
	if withBody {
		dest = append(dest, &m.Body)
	}
	if err := row.Scan(dest...); err != nil {
		return m, err
	}
	m.UID, m.UIDValidity, m.Date = uint32(uid), uint32(validity), m.Date.UTC()
	for _, f := range []struct {
		raw  []byte
		into any
	}{{from, &m.From}, {to, &m.To}, {cc, &m.Cc}, {atts, &m.Attachments}} {
		if err := json.Unmarshal(f.raw, f.into); err != nil {
			return m, err
		}
	}
	return m, nil
}

func (s *PgStore) Get(ctx context.Context, account, folder string, uid uint32) (Message, bool, error) {
	m, err := scanMessage(s.pool.QueryRow(ctx, `select `+messageColumns+`, body from message where account = $1 and folder = $2 and uid = $3`,
		account, folder, int64(uid)), true)
	if errors.Is(err, pgx.ErrNoRows) {
		return Message{}, false, nil
	}
	return m, err == nil, err
}

func (s *PgStore) Search(ctx context.Context, q Query) ([]Message, error) {
	limit := q.Limit
	if limit <= 0 {
		limit = 20
	}
	rows, err := s.pool.Query(ctx, `select `+messageColumns+` from message
		where account = any($1) and ($2 = '' or folder = $2) and search @@ plainto_tsquery('simple'::regconfig, $3)
		order by date desc limit $4`, q.Accounts, q.Folder, q.Text, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Message
	for rows.Next() {
		m, err := scanMessage(rows, false)
		if err != nil {
			return nil, err
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

func (s *PgStore) Folders(ctx context.Context, account string) ([]Folder, error) {
	rows, err := s.pool.Query(ctx, `select name, uidvalidity, messages, synced from folder where account = $1 order by name`, account)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Folder
	for rows.Next() {
		f := Folder{Account: account}
		var validity int64
		if err := rows.Scan(&f.Name, &validity, &f.Messages, &f.Synced); err != nil {
			return nil, err
		}
		f.UIDValidity, f.Synced = uint32(validity), f.Synced.UTC()
		out = append(out, f)
	}
	return out, rows.Err()
}

func (s *PgStore) Pending(ctx context.Context, limit int) ([]Event, error) {
	rows, err := s.pool.Query(ctx, `select id, key, body from outbox order by id limit $1`, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Event
	for rows.Next() {
		var e Event
		var body []byte
		if err := rows.Scan(&e.ID, &e.Key, &body); err != nil {
			return nil, err
		}
		if err := json.Unmarshal(body, &e.Body); err != nil {
			return nil, err
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

func (s *PgStore) Said(ctx context.Context, ids []int64) error {
	_, err := s.pool.Exec(ctx, `delete from outbox where id = any($1)`, ids)
	return err
}

func (s *PgStore) SetState(ctx context.Context, st AccountState) error {
	_, err := s.pool.Exec(ctx, `insert into account_state (account, state, why, heard, synced, said) values ($1, $2, $3, $4, $5, $6)
		on conflict (account) do update set state = excluded.state, why = excluded.why, heard = excluded.heard,
		synced = excluded.synced, said = excluded.said`,
		st.Account, st.State, st.Why, nullTime(st.Heard), nullTime(st.Synced), st.Said)
	return err
}

func nullTime(t time.Time) *time.Time {
	if t.IsZero() {
		return nil
	}
	return &t
}

func (s *PgStore) States(ctx context.Context) (map[string]AccountState, error) {
	rows, err := s.pool.Query(ctx, `select account, state, why, heard, synced, said from account_state`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[string]AccountState{}
	for rows.Next() {
		var st AccountState
		var heard, synced *time.Time
		if err := rows.Scan(&st.Account, &st.State, &st.Why, &heard, &synced, &st.Said); err != nil {
			return nil, err
		}
		if heard != nil {
			st.Heard = heard.UTC()
		}
		if synced != nil {
			st.Synced = synced.UTC()
		}
		st.Said = st.Said.UTC()
		out[st.Account] = st
	}
	return out, rows.Err()
}