Refactor db pkg

2026-10-06 12:07:31 +00:00
parent bf6971e470
commit bbfac06732
5 changed files with 375 additions and 188 deletions
@@ -2,8 +2,11 @@ package db

 import (
 	"context"
+	"errors"
 	"log/slog"
 	"testing"
+
+	"github.com/mattn/go-sqlite3"
 )

 func TestQueries(t *testing.T) {
@@ -29,8 +32,8 @@ func TestQueries(t *testing.T) {
 	t.Run("ConversationTimestamp", func(t *testing.T) {
 		testConversationTimestamp(t, ctx, db)
 	})
-	t.Run("Approvals", func(t *testing.T) {
-		testApprovals(t, ctx, db)
+	t.Run("Executions", func(t *testing.T) {
+		testExecutions(t, ctx, db)
 	})
 	t.Run("Delete", func(t *testing.T) {
 		testDelete(t, ctx, db)
@@ -380,7 +383,7 @@ func testDelete(t *testing.T, ctx context.Context, db *SQLite) {
 	}
 }

-func testApprovals(t *testing.T, ctx context.Context, db *SQLite) {
+func testExecutions(t *testing.T, ctx context.Context, db *SQLite) {
 	var (
 		qr = New(db.Reader)
 		qw = New(db.Writer)
@@ -412,112 +415,228 @@ func testApprovals(t *testing.T, ctx context.Context, db *SQLite) {
 		t.Fatalf("InsertMessage: %v", err)
 	}

-	// Insert an approval.
-	approval, err := qw.InsertApproval(ctx, InsertApprovalParams{
+	// Insert an execution.
+	execution, err := qw.InsertExecution(ctx, InsertExecutionParams{
 		MessageId:  msg.Id,
+		Status:     ExecutionStatusPending,
 		ToolCallId: "call_abc",
 	})
 	if err != nil {
-		t.Fatalf("InsertApproval: %v", err)
+		t.Fatalf("InsertExecution: %v", err)
 	}
-	if approval.Status != ApprovalStatusPending {
-		t.Fatalf("expected pending, got %s", approval.Status)
+	if execution.Status != ExecutionStatusPending {
+		t.Fatalf("expected pending, got %s", execution.Status)
 	}
-	if approval.ToolCallId != "call_abc" {
-		t.Fatalf("expected tool_call_id call_abc, got %s", approval.ToolCallId)
+	if execution.ToolCallId != "call_abc" {
+		t.Fatalf("expected tool_call_id call_abc, got %s", execution.ToolCallId)
 	}

 	// Select by ID.
-	res, err := qr.SelectApprovalId(ctx, approval.Id)
+	res, err := qr.SelectExecutionId(ctx, execution.Id)
 	if err != nil {
-		t.Fatalf("SelectApprovalId: %v", err)
+		t.Fatalf("SelectExecutionId: %v", err)
 	}
 	if res == nil {
-		t.Fatal("SelectApprovalId returned nil")
+		t.Fatal("SelectExecutionId returned nil")
 	}
 	if res.MessageId != msg.Id {
 		t.Errorf("expected message_id %s, got %s", msg.Id, res.MessageId)
 	}

 	// Select by message ID.
-	approvals, err := qr.SelectApprovalsMessageId(ctx, msg.Id)
+	executions, err := qr.SelectExecutionsMessageId(ctx, msg.Id)
 	if err != nil {
-		t.Fatalf("SelectApprovalsMessageId: %v", err)
+		t.Fatalf("SelectExecutionsMessageId: %v", err)
 	}
-	if len(approvals) != 1 {
-		t.Fatalf("expected 1 approval, got %d", len(approvals))
+	if len(executions) != 1 {
+		t.Fatalf("expected 1 execution, got %d", len(executions))
 	}

 	// Select pending.
-	pending, err := qr.SelectPendingApprovals(ctx)
+	pending, err := qr.SelectPendingExecutions(ctx)
 	if err != nil {
-		t.Fatalf("SelectPendingApprovals: %v", err)
+		t.Fatalf("SelectPendingExecutions: %v", err)
 	}
 	if len(pending) != 1 {
-		t.Fatalf("expected 1 pending approval, got %d", len(pending))
+		t.Fatalf("expected 1 pending execution, got %d", len(pending))
 	}

 	// Transition to approved.
-	if err := qw.UpdateApprovalStatus(ctx, UpdateApprovalStatusParams{
-		Id:     approval.Id,
-		Status: ApprovalStatusApproved,
+	if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{
+		Id:     execution.Id,
+		Status: ExecutionStatusApproved,
 	}); err != nil {
-		t.Fatalf("UpdateApprovalStatus: %v", err)
+		t.Fatalf("UpdateExecutionStatus: %v", err)
 	}

-	res, err = qr.SelectApprovalId(ctx, approval.Id)
+	res, err = qr.SelectExecutionId(ctx, execution.Id)
 	if err != nil {
-		t.Fatalf("SelectApprovalId after update: %v", err)
+		t.Fatalf("SelectExecutionId after update: %v", err)
 	}
-	if res.Status != ApprovalStatusApproved {
+	if res.Status != ExecutionStatusApproved {
 		t.Fatalf("expected approved, got %s", res.Status)
 	}

 	// Pending should now be empty.
-	pending, err = qr.SelectPendingApprovals(ctx)
+	pending, err = qr.SelectPendingExecutions(ctx)
 	if err != nil {
-		t.Fatalf("SelectPendingApprovals after update: %v", err)
+		t.Fatalf("SelectPendingExecutions after update: %v", err)
 	}
 	if len(pending) != 0 {
 		t.Fatalf("expected 0 pending, got %d", len(pending))
 	}

-	// Insert a second approval and deny it.
-	approval2, err := qw.InsertApproval(ctx, InsertApprovalParams{
+	// Insert a second execution and deny it.
+	execution2, err := qw.InsertExecution(ctx, InsertExecutionParams{
 		MessageId:  msg.Id,
+		Status:     ExecutionStatusPending,
 		ToolCallId: "call_def",
 	})
 	if err != nil {
-		t.Fatalf("InsertApproval second: %v", err)
+		t.Fatalf("InsertExecution second: %v", err)
 	}

-	if err := qw.UpdateApprovalStatus(ctx, UpdateApprovalStatusParams{
-		Id:     approval2.Id,
-		Status: ApprovalStatusDenied,
+	if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{
+		Id:     execution2.Id,
+		Status: ExecutionStatusDenied,
 	}); err != nil {
-		t.Fatalf("UpdateApprovalStatus denied: %v", err)
+		t.Fatalf("UpdateExecutionStatus denied: %v", err)
 	}

-	res2, err := qr.SelectApprovalId(ctx, approval2.Id)
+	res2, err := qr.SelectExecutionId(ctx, execution2.Id)
 	if err != nil {
-		t.Fatalf("SelectApprovalId after deny: %v", err)
+		t.Fatalf("SelectExecutionId after deny: %v", err)
 	}
-	if res2.Status != ApprovalStatusDenied {
+	if res2.Status != ExecutionStatusDenied {
 		t.Fatalf("expected denied, got %s", res2.Status)
 	}

-	// Both approvals for the message, ordered by created_at.
-	all, err := qr.SelectApprovalsMessageId(ctx, msg.Id)
+	// Transition to running, then completed.
+	if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{
+		Id:     execution.Id,
+		Status: ExecutionStatusRunning,
+	}); err != nil {
+		t.Fatalf("UpdateExecutionStatus running: %v", err)
+	}
+	if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{
+		Id:     execution.Id,
+		Status: ExecutionStatusCompleted,
+	}); err != nil {
+		t.Fatalf("UpdateExecutionStatus completed: %v", err)
+	}
+
+	// Both executions for the message, ordered by created_at.
+	all, err := qr.SelectExecutionsMessageId(ctx, msg.Id)
 	if err != nil {
-		t.Fatalf("SelectApprovalsMessageId final: %v", err)
+		t.Fatalf("SelectExecutionsMessageId final: %v", err)
 	}
 	if len(all) != 2 {
-		t.Fatalf("expected 2 approvals, got %d", len(all))
+		t.Fatalf("expected 2 executions, got %d", len(all))
 	}
-	if all[0].Status != ApprovalStatusApproved {
-		t.Fatalf("expected first approval approved, got %s", all[0].Status)
+	if all[0].Status != ExecutionStatusCompleted {
+		t.Fatalf("expected first execution completed, got %s", all[0].Status)
 	}
-	if all[1].Status != ApprovalStatusDenied {
-		t.Fatalf("expected second approval denied, got %s", all[1].Status)
+	if all[1].Status != ExecutionStatusDenied {
+		t.Fatalf("expected second execution denied, got %s", all[1].Status)
+	}
+
+	// Parameter values and the schema default, on a second message.
+	msg2, err := qw.InsertMessage(ctx, InsertMessageParams{
+		ConversationId: c.Id,
+		Sequence:       2,
+		Role:           MessageRoleAssistant,
+		Content:        `{"role":"assistant","content":"second"}`,
+	})
+	if err != nil {
+		t.Fatalf("InsertMessage second: %v", err)
+	}
+
+	execRunning, err := qw.InsertExecution(ctx, InsertExecutionParams{
+		MessageId:  msg2.Id,
+		Status:     ExecutionStatusRunning,
+		ToolCallId: "call_running",
+	})
+	if err != nil {
+		t.Fatalf("InsertExecution running: %v", err)
+	}
+	if execRunning.Status != ExecutionStatusRunning {
+		t.Fatalf("expected running, got %s", execRunning.Status)
+	}
+
+	// A status other than running is stored as passed.
+	execPending, err := qw.InsertExecution(ctx, InsertExecutionParams{
+		MessageId:  msg2.Id,
+		Status:     ExecutionStatusPending,
+		ToolCallId: "call_pending",
+	})
+	if err != nil {
+		t.Fatalf("InsertExecution pending: %v", err)
+	}
+	if execPending.Status != ExecutionStatusPending {
+		t.Fatalf("expected pending, got %s", execPending.Status)
+	}
+
+	// The zero value is rejected by the table's status check.
+	_, err = qw.InsertExecution(ctx, InsertExecutionParams{
+		MessageId:  msg2.Id,
+		ToolCallId: "call_empty",
+	})
+	var sqliteErr sqlite3.Error
+	if !errors.As(err, &sqliteErr) {
+		t.Fatalf("InsertExecution with zero status: expected sqlite3.Error, got %v", err)
+	}
+	if sqliteErr.ExtendedCode != sqlite3.ErrConstraintCheck {
+		t.Fatalf("InsertExecution with zero status: expected ErrConstraintCheck, got %v", err)
+	}
+
+	// A second row for the same message and tool call is rejected.
+	_, err = qw.InsertExecution(ctx, InsertExecutionParams{
+		MessageId:  msg2.Id,
+		Status:     ExecutionStatusPending,
+		ToolCallId: "call_running",
+	})
+	if !errors.As(err, &sqliteErr) {
+		t.Fatalf("InsertExecution duplicate: expected sqlite3.Error, got %v", err)
+	}
+	if sqliteErr.ExtendedCode != sqlite3.ErrConstraintUnique {
+		t.Fatalf("InsertExecution duplicate: expected ErrConstraintUnique, got %v", err)
+	}
+
+	// A row inserted without a status takes the schema default, running.
+	if _, err := db.Writer.ExecContext(
+		ctx,
+		`INSERT INTO executions (id, message_id, tool_call_id) VALUES (?, ?, ?)`,
+		"exec_default", msg2.Id, "call_default",
+	); err != nil {
+		t.Fatalf("insert execution without status: %v", err)
+	}
+	res, err = qr.SelectExecutionId(ctx, "exec_default")
+	if err != nil {
+		t.Fatalf("SelectExecutionId default: %v", err)
+	}
+	if res == nil {
+		t.Fatal("SelectExecutionId default returned nil")
+	}
+	if res.Status != ExecutionStatusRunning {
+		t.Fatalf("expected default running, got %s", res.Status)
+	}
+
+	// The rejected insert stored no row: the message's rows are exactly the
+	// three successful inserts.
+	all2, err := qr.SelectExecutionsMessageId(ctx, msg2.Id)
+	if err != nil {
+		t.Fatalf("SelectExecutionsMessageId second: %v", err)
+	}
+	seen := map[string]bool{}
+	for _, e := range all2 {
+		seen[e.ToolCallId] = true
+	}
+	for _, id := range []string{"call_default", "call_pending", "call_running"} {
+		if !seen[id] {
+			t.Fatalf("execution %s not found", id)
+		}
+	}
+	if len(seen) != 3 {
+		t.Fatalf("expected 3 executions, got %d", len(seen))
 	}
 }