kandev/apps/backend/internal/worktree/store_postgres_test.go

package worktree

import (
	"context"
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/jmoiron/sqlx"

	"github.com/kandev/kandev/internal/task/models"
	tasksqlite "github.com/kandev/kandev/internal/task/repository/sqlite"
	"github.com/kandev/kandev/internal/testutil"
)

func TestSQLiteStore_ReinitializesSchemaOnPostgres(t *testing.T) {
	db := testutil.OpenIsolatedPostgres(t, testutil.PostgresDSNFromEnv(t))

	if _, err := tasksqlite.NewWithDB(db, db, nil); err != nil {
		t.Fatalf("first task schema init: %v", err)
	}
	if _, err := NewSQLiteStore(db, db); err != nil {
		t.Fatalf("first worktree schema init: %v", err)
	}
	if _, err := tasksqlite.NewWithDB(db, db, nil); err != nil {
		t.Fatalf("second task schema init: %v", err)
	}
	if _, err := NewSQLiteStore(db, db); err != nil {
		t.Fatalf("second worktree schema init: %v", err)
	}
}

func TestSQLiteStore_CreateWorktreeSerializesAgainstPostgresSessionRebind(t *testing.T) {
	ctx := context.Background()
	dsn := testutil.PostgresDSNFromEnv(t)
	db := testutil.OpenIsolatedPostgres(t, dsn)
	if _, err := tasksqlite.NewWithDB(db, db, nil); err != nil {
		t.Fatalf("init task schema: %v", err)
	}
	holderDB := openWorktreePostgresConnection(t, dsn, db)
	moverDB := openWorktreePostgresConnection(t, dsn, db)
	observerDB := openWorktreePostgresConnection(t, dsn, db)
	store, err := NewSQLiteStore(moverDB, moverDB)
	if err != nil {
		t.Fatalf("new worktree store: %v", err)
	}
	seeds := []struct {
		label     string
		statement string
	}{
		{"workspace", `INSERT INTO workspaces (id, name, created_at, updated_at) VALUES ('workspace', 'Workspace', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`},
		{"original task", `INSERT INTO tasks (id, workspace_id, title, created_at, updated_at) VALUES ('task-original', 'workspace', 'Original', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`},
		{"rebound task", `INSERT INTO tasks (id, workspace_id, title, created_at, updated_at) VALUES ('task-rebound', 'workspace', 'Rebound', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`},
		{"original environment", `INSERT INTO task_environments (id, task_id, executor_type, status, workspace_path, created_at, updated_at)
			VALUES ('env-original', 'task-original', 'worktree', 'ready', '/tmp/original', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`},
		{"rebound environment", `INSERT INTO task_environments (id, task_id, executor_type, status, workspace_path, created_at, updated_at)
			VALUES ('env-rebound', 'task-rebound', 'worktree', 'ready', '/tmp/rebound', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`},
		{"session", `INSERT INTO task_sessions (id, task_id, state, task_environment_id, started_at, updated_at)
			VALUES ('session-rebind', 'task-original', 'WAITING_FOR_INPUT', 'env-original', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`},
	}
	for _, seed := range seeds {
		if _, err := db.ExecContext(ctx, seed.statement); err != nil {
			t.Fatalf("seed %s: %v", seed.label, err)
		}
	}

	rebindTx, err := holderDB.BeginTxx(ctx, nil)
	if err != nil {
		t.Fatalf("begin rebind: %v", err)
	}
	defer func() { _ = rebindTx.Rollback() }()
	if _, err := rebindTx.ExecContext(ctx, `UPDATE task_sessions SET task_environment_id = 'env-rebound' WHERE id = 'session-rebind'`); err != nil {
		t.Fatalf("stage session rebind: %v", err)
	}

	var moverPID int
	if err := moverDB.GetContext(ctx, &moverPID, `SELECT pg_backend_pid()`); err != nil {
		t.Fatalf("read mover backend pid: %v", err)
	}
	done := make(chan error, 1)
	finished := make(chan struct{})
	go func() {
		defer close(finished)
		done <- store.CreateWorktree(ctx, &Worktree{
			ID: "wt-rebind", SessionID: "session-rebind", TaskEnvironmentID: "env-original",
			RepositoryID: "repo-rebind", BranchSlug: "branch-rebind", Path: "/tmp/rebound/repo",
			Branch: "feature/rebind", Status: StatusActive,
		})
	}()
	if err := waitForWorktreePostgresLock(ctx, observerDB, moverPID, finished); err != nil {
		t.Fatal(err)
	}
	select {
	case <-finished:
		t.Fatal("CreateWorktree returned while the session rebind still held the row lock")
	default:
	}
	if err := rebindTx.Commit(); err != nil {
		t.Fatalf("commit session rebind: %v", err)
	}
	select {
	case createErr := <-done:
		if !errors.Is(createErr, models.ErrWorkspaceReuseUnsafe) {
			t.Fatalf("CreateWorktree error = %v, want ErrWorkspaceReuseUnsafe", createErr)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("CreateWorktree did not resume after the session rebind committed")
	}
	var count int
	if err := db.GetContext(ctx, &count, `SELECT COUNT(*) FROM task_environment_repos WHERE worktree_id = 'wt-rebind'`); err != nil {
		t.Fatalf("count rebound inventory: %v", err)
	}
	if count != 0 {
		t.Fatalf("rebound worktree inventory rows = %d, want 0", count)
	}
}

func openWorktreePostgresConnection(t *testing.T, dsn string, schemaDB *sqlx.DB) *sqlx.DB {
	t.Helper()
	var schema string
	if err := schemaDB.Get(&schema, `SELECT current_schema()`); err != nil {
		t.Fatalf("read postgres test schema: %v", err)
	}
	db, err := sqlx.Open("pgx", dsn)
	if err != nil {
		t.Fatalf("open postgres test connection: %v", err)
	}
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	t.Cleanup(func() { _ = db.Close() })
	if _, err := db.Exec(`SET search_path TO ` + schema); err != nil {
		t.Fatalf("set postgres test search path: %v", err)
	}
	return db
}

func waitForWorktreePostgresLock(ctx context.Context, observer *sqlx.DB, pid int, done <-chan struct{}) error {
	deadlineCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-done:
			return fmt.Errorf("CreateWorktree completed before backend %d entered a PostgreSQL lock wait", pid)
		default:
		}
		var waiting bool
		if err := observer.GetContext(deadlineCtx, &waiting, `
			SELECT EXISTS (
				SELECT 1 FROM pg_stat_activity
				WHERE pid = $1 AND wait_event_type = 'Lock'
			)
		`, pid); err != nil {
			return fmt.Errorf("query PostgreSQL lock state for backend %d: %w", pid, err)
		}
		if waiting {
			return nil
		}
		select {
		case <-done:
			return fmt.Errorf("CreateWorktree completed before backend %d entered a PostgreSQL lock wait", pid)
		case <-ticker.C:
		case <-deadlineCtx.Done():
			return fmt.Errorf("timed out waiting for backend %d to enter a PostgreSQL lock wait", pid)
		}
	}
}