plumber/internal/store/migrate.go

package store

import (
	"context"
	"database/sql"
	"fmt"
	"log"
)

const migrateLockKey int64 = 0x706c756d5f6d6967 // "plum_mig"

// migrateUserProfileColumns adds avatar_url and state when missing (existing DBs).
func migrateUserProfileColumns(ctx context.Context, exec execContext) error {
	cols := []string{"avatar_url", "state"}
	for _, col := range cols {
		stmt := fmt.Sprintf(`ALTER TABLE users ADD COLUMN IF NOT EXISTS %s TEXT NOT NULL DEFAULT ''`, col)
		if _, err := exec.ExecContext(ctx, stmt); err != nil {
			return fmt.Errorf("add column %s: %w", col, err)
		}
	}
	return nil
}

// migrateUserEmail adds email and a partial unique index on lower(email).
func migrateUserEmail(ctx context.Context, exec execContext) error {
	if _, err := exec.ExecContext(ctx, `ALTER TABLE users ADD COLUMN IF NOT EXISTS email TEXT NOT NULL DEFAULT ''`); err != nil {
		return fmt.Errorf("add column email: %w", err)
	}
	if _, err := exec.ExecContext(ctx, `
CREATE UNIQUE INDEX IF NOT EXISTS users_email_lower_uidx
  ON users (lower(email))
  WHERE email <> ''`); err != nil {
		return fmt.Errorf("users_email_lower_uidx: %w", err)
	}
	return nil
}

// migratePosts creates the unified post model.
func migratePosts(ctx context.Context, exec execContext) error {
	steps := []struct {
		name string
		sql  string
	}{
		{name: "create posts", sql: `
CREATE TABLE IF NOT EXISTS posts (
  id TEXT PRIMARY KEY,
  parent_id TEXT REFERENCES posts(id) ON DELETE CASCADE,
  author_id TEXT NOT NULL REFERENCES users(id),
  title TEXT NOT NULL DEFAULT '',
  body TEXT NOT NULL,
  city TEXT NOT NULL DEFAULT '',
  post_date TEXT NOT NULL DEFAULT '',
  post_state TEXT NOT NULL,
  created_at TEXT NOT NULL,
  updated_at TEXT NOT NULL,
  CONSTRAINT posts_shape_check CHECK (
    (parent_id IS NULL AND title <> '' AND post_date <> '')
    OR
    (parent_id IS NOT NULL AND title = '' AND city = '' AND post_date = '')
  )
)`},
		{name: "index post replies", sql: `
CREATE INDEX IF NOT EXISTS idx_posts_parent_created
  ON posts(parent_id, created_at, id)`},
		{name: "index post authors", sql: `
CREATE INDEX IF NOT EXISTS idx_posts_author_created
  ON posts(author_id, created_at DESC, id DESC)`},
		{name: "index root posts", sql: `
CREATE INDEX IF NOT EXISTS idx_posts_root_date
  ON posts(post_date, post_state)
  WHERE parent_id IS NULL`},
		{name: "create post votes", sql: `
CREATE TABLE IF NOT EXISTS post_votes (
  user_id TEXT NOT NULL REFERENCES users(id),
  post_id TEXT NOT NULL REFERENCES posts(id) ON DELETE CASCADE,
  value INTEGER NOT NULL CHECK (value IN (-1, 1)),
  PRIMARY KEY (user_id, post_id)
)`},
	}
	for _, step := range steps {
		if _, err := exec.ExecContext(ctx, step.sql); err != nil {
			return fmt.Errorf("%s: %w", step.name, err)
		}
	}
	return nil
}

func migratePostVoteIndex(ctx context.Context, exec execContext) error {
	if _, err := exec.ExecContext(ctx, `
CREATE INDEX IF NOT EXISTS idx_post_votes_post_id
  ON post_votes(post_id)`); err != nil {
		return fmt.Errorf("idx_post_votes_post_id: %w", err)
	}
	return nil
}

func migratePostAuthorIndex(ctx context.Context, exec execContext) error {
	if _, err := exec.ExecContext(ctx, `
CREATE INDEX IF NOT EXISTS idx_posts_author_created
  ON posts(author_id, created_at DESC, id DESC)`); err != nil {
		return fmt.Errorf("idx_posts_author_created: %w", err)
	}
	return nil
}

func migrateDropLegacyPostTables(ctx context.Context, exec execContext) error {
	steps := []struct {
		name string
		sql  string
	}{
		{"drop legacy answers", `DROP TABLE IF EXISTS answers`},
		{"drop legacy votes", `DROP TABLE IF EXISTS votes`},
		{"drop legacy questions", `DROP TABLE IF EXISTS questions`},
	}
	for _, step := range steps {
		if _, err := exec.ExecContext(ctx, step.sql); err != nil {
			return fmt.Errorf("%s: %w", step.name, err)
		}
	}
	return nil
}

func migratePostDate(ctx context.Context, exec execContext) error {
	steps := []struct {
		name string
		sql  string
	}{
		{"rename post date", `
DO $migration$
BEGIN
  IF EXISTS (
    SELECT 1
    FROM information_schema.columns
    WHERE table_schema = current_schema()
      AND table_name = 'posts'
      AND column_name = 'hunt_date'
  ) AND NOT EXISTS (
    SELECT 1
    FROM information_schema.columns
    WHERE table_schema = current_schema()
      AND table_name = 'posts'
      AND column_name = 'post_date'
  ) THEN
    ALTER TABLE posts RENAME COLUMN hunt_date TO post_date;
  END IF;
END
$migration$`},
		{"drop legacy root date index", `
DROP INDEX IF EXISTS idx_posts_root_hunt`},
		{"create root date index", `
DO $migration$
BEGIN
  IF EXISTS (
    SELECT 1
    FROM information_schema.columns
    WHERE table_schema = current_schema()
      AND table_name = 'posts'
      AND column_name = 'post_state'
  ) THEN
    CREATE INDEX IF NOT EXISTS idx_posts_root_date
      ON posts(post_date, post_state)
      WHERE parent_id IS NULL;
  ELSE
    CREATE INDEX IF NOT EXISTS idx_posts_root_date
      ON posts(post_date, hidden)
      WHERE parent_id IS NULL;
  END IF;
END
$migration$`},
	}
	for _, step := range steps {
		if _, err := exec.ExecContext(ctx, step.sql); err != nil {
			return fmt.Errorf("%s: %w", step.name, err)
		}
	}
	return nil
}

func migratePostState(ctx context.Context, exec execContext) error {
	if _, err := exec.ExecContext(ctx, `
ALTER TABLE posts
  ADD COLUMN IF NOT EXISTS post_state TEXT`); err != nil {
		return fmt.Errorf("add post state: %w", err)
	}
	hasHidden, err := migrationColumnExists(ctx, exec, "posts", "hidden")
	if err != nil {
		return fmt.Errorf("check hidden column: %w", err)
	}
	if hasHidden {
		if _, err := exec.ExecContext(ctx, `
UPDATE posts
SET post_state = CASE
  WHEN hidden = 0 THEN $1
  ELSE $2
END`, string(PostStateVisible), string(PostStateHidden)); err != nil {
			return fmt.Errorf("copy hidden state: %w", err)
		}
	}

	steps := []struct {
		name string
		sql  string
	}{
		{"drop root date index", `
DROP INDEX IF EXISTS idx_posts_root_date`},
		{"drop legacy post shape constraint", `
ALTER TABLE posts DROP CONSTRAINT IF EXISTS posts_check`},
		{"drop post shape constraint", `
ALTER TABLE posts DROP CONSTRAINT IF EXISTS posts_shape_check`},
		{"drop hidden", `
ALTER TABLE posts DROP COLUMN IF EXISTS hidden`},
		{"require post state", `
ALTER TABLE posts ALTER COLUMN post_state SET NOT NULL`},
		{"create post shape constraint", `
ALTER TABLE posts
  ADD CONSTRAINT posts_shape_check CHECK (
    (parent_id IS NULL AND title <> '' AND post_date <> '')
    OR
    (parent_id IS NOT NULL AND title = '' AND city = '' AND post_date = '')
  )`},
		{"create root date index", `
CREATE INDEX idx_posts_root_date
  ON posts(post_date, post_state)
  WHERE parent_id IS NULL`},
	}
	for _, step := range steps {
		if _, err := exec.ExecContext(ctx, step.sql); err != nil {
			return fmt.Errorf("%s: %w", step.name, err)
		}
	}
	return nil
}

func migrationColumnExists(
	ctx context.Context,
	exec execContext,
	tableName string,
	columnName string,
) (bool, error) {
	rows, err := exec.QueryContext(ctx, `
SELECT EXISTS (
  SELECT 1
  FROM information_schema.columns
  WHERE table_schema = current_schema()
    AND table_name = $1
    AND column_name = $2
)`, tableName, columnName)
	if err != nil {
		return false, err
	}
	defer rows.Close()
	if !rows.Next() {
		return false, rows.Err()
	}
	var exists bool
	if err := rows.Scan(&exists); err != nil {
		return false, err
	}
	return exists, rows.Err()
}

type execContext interface {
	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
}

// applyMigrations runs versioned migrations under a session-level advisory lock
// held for the entire process (check versions → apply → record).
func applyMigrations(db *sql.DB, schemaSQL string) error {
	ctx := context.Background()
	conn, err := db.Conn(ctx)
	if err != nil {
		return err
	}
	defer conn.Close()

	if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1)`, migrateLockKey); err != nil {
		return fmt.Errorf("migrate lock: %w", err)
	}
	defer func() {
		if _, unlockErr := conn.ExecContext(ctx, `SELECT pg_advisory_unlock($1)`, migrateLockKey); unlockErr != nil {
			log.Printf("migrate unlock: %v", unlockErr)
		}
	}()

	if _, err := conn.ExecContext(ctx, `
CREATE TABLE IF NOT EXISTS schema_migrations (
  version TEXT PRIMARY KEY,
  applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`); err != nil {
		return fmt.Errorf("schema_migrations: %w", err)
	}

	applied, err := appliedVersions(ctx, conn)
	if err != nil {
		return err
	}

	migrations := []struct {
		version string
		run     func(context.Context, execContext) error
	}{
		{"001_schema", func(ctx context.Context, exec execContext) error {
			return applySchema(ctx, exec, schemaSQL)
		}},
		{"002_user_profile_columns", migrateUserProfileColumns},
		{"003_user_email", migrateUserEmail},
		{"004_posts", migratePosts},
		{"005_post_vote_post_id_index", migratePostVoteIndex},
		{"006_post_date", migratePostDate},
		{"007_post_state", migratePostState},
		{"008_post_author_index", migratePostAuthorIndex},
		{"009_drop_legacy_post_tables", migrateDropLegacyPostTables},
	}
	for _, m := range migrations {
		if applied[m.version] {
			continue
		}
		log.Printf("migrate: applying %s", m.version)
		if err := m.run(ctx, conn); err != nil {
			return fmt.Errorf("migrate %s: %w", m.version, err)
		}
		if _, err := conn.ExecContext(ctx, `INSERT INTO schema_migrations (version) VALUES ($1)`, m.version); err != nil {
			return fmt.Errorf("record %s: %w", m.version, err)
		}
	}
	return nil
}

func appliedVersions(ctx context.Context, exec execContext) (map[string]bool, error) {
	rows, err := exec.QueryContext(ctx, `SELECT version FROM schema_migrations`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[string]bool{}
	for rows.Next() {
		var v string
		if err := rows.Scan(&v); err != nil {
			return nil, err
		}
		out[v] = true
	}
	return out, rows.Err()
}