chore(tests): add advisory locks and statement parser for test infrastructure

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
2026-06-11 22:08:23 +01:00
parent 6de0371e1d
commit 56d91ba4ef
1 changed files with 123 additions and 51 deletions
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"

 	"github.com/jackc/pgx/v5"
@@ -57,16 +58,25 @@ func NewPool(dsn string) (*pgxpool.Pool, error) {

 func Migrate(t *testing.T, pool *pgxpool.Pool) {
 	t.Helper()
-
 	ctx := context.Background()

+	// Acquire a dedicated connection so advisory lock, migration, and unlock all use the same session
+	conn, err := pool.Acquire(ctx)
+	if err != nil {
+		t.Fatalf("Failed to acquire connection for migration: %v", err)
+	}
+	defer conn.Release()
+
+	_, err = conn.Exec(ctx, "SELECT pg_advisory_lock(1337)")
+	if err != nil {
+		t.Fatalf("Failed to acquire migration advisory lock: %v", err)
+	}
+	defer conn.Exec(ctx, "SELECT pg_advisory_unlock(1337)")
+
 	// Check if database already has tables or types
-	var err error
 	var typeCount int
-	err = pool.QueryRow(ctx, "SELECT COUNT(*) FROM pg_type WHERE typnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')").Scan(&typeCount)
+	err = conn.QueryRow(ctx, "SELECT COUNT(*) FROM pg_type WHERE typnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')").Scan(&typeCount)
 	if err == nil && typeCount > 0 {
-		t.Log("Database already has types, dropping and recreating for clean state...")
-		
 		// Drop TYPES FIRST (they have CASCADE dependencies on tables)
 		typeDrops := []string{
 			"DROP TYPE IF EXISTS account_role CASCADE",
@@ -82,15 +92,19 @@ func Migrate(t *testing.T, pool *pgxpool.Pool) {
 			"DROP TYPE IF EXISTS milestone_unit CASCADE",
 			"DROP TYPE IF EXISTS discount_campaign_scope CASCADE",
 			"DROP TYPE IF EXISTS discount_campaign_status CASCADE",
+			"DROP TYPE IF EXISTS till_item_type CASCADE",
 		}
 		for _, stmt := range typeDrops {
-			if _, err := pool.Exec(ctx, stmt); err != nil {
+			if _, err := conn.Exec(ctx, stmt); err != nil {
 				t.Logf("Warning dropping type: %v (expected if using IF EXISTS)", err)
 			}
 		}
 		
 		// Drop all tables, sequences, and views in correct order
 		dropOrder := []string{
+			"till_sales",
+			"gift_card_transactions",
+			"gift_card_expired_balances",
 			"booking_discounts",
 			"loyalty_redemptions",
 			"discount_campaigns",
@@ -127,7 +141,7 @@ func Migrate(t *testing.T, pool *pgxpool.Pool) {
 		
 		for _, table := range dropOrder {
 			stmt := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE", table)
-			if _, err := pool.Exec(ctx, stmt); err != nil {
+			if _, err := conn.Exec(ctx, stmt); err != nil {
 				t.Logf("Warning dropping table %s: %v (expected if using IF EXISTS)", table, err)
 			}
 		}
@@ -140,7 +154,7 @@ func Migrate(t *testing.T, pool *pgxpool.Pool) {
 			"DROP SEQUENCE IF EXISTS exceptional_working_hours_groups_id_seq",
 		}
 		for _, stmt := range seqDrops {
-			if _, err := pool.Exec(ctx, stmt); err != nil {
+			if _, err := conn.Exec(ctx, stmt); err != nil {
 				t.Logf("Warning dropping sequence: %v", err)
 			}
 		}
@@ -165,18 +179,60 @@ func Migrate(t *testing.T, pool *pgxpool.Pool) {
 		t.Fatal("Could not find init-script.sql in any expected location")
 	}

-	// Simple migration: just create tables that don't exist
-	// Note: This doesn't handle stored procedures properly, but the database
-	// should already be set up with the correct schema
-	t.Log("Running migration...")
+	// Split and execute the schema SQL statement-by-statement for robust execution and diagnostic visibility
+	statements := splitSQLStatements(schemaSQL)
+	for i, stmt := range statements {
+		stmt = strings.TrimSpace(stmt)
+		if stmt == "" {
+			continue
+		}
+		firstLine := strings.Split(stmt, "\n")[0]
+		if len(firstLine) > 80 {
+			firstLine = firstLine[:80] + "..."
+		}

-	// Execute the schema SQL
-	_, err = pool.Exec(ctx, schemaSQL)
-	if err != nil {
-		t.Fatalf("Failed to execute migration: %v", err)
+		_, err = conn.Exec(ctx, stmt)
+		if err != nil {
+			t.Fatalf("Failed to execute migration statement [%d]: %s\nError: %v", i+1, firstLine, err)
+		}
 	}
 }

+func splitSQLStatements(sql string) []string {
+	var statements []string
+	var currentStatement strings.Builder
+	inDollarBlock := false
+
+	lines := strings.Split(sql, "\n")
+	for _, line := range lines {
+		trimmed := strings.TrimSpace(line)
+		if trimmed == "" {
+			continue
+		}
+		if strings.HasPrefix(trimmed, "--") {
+			continue
+		}
+
+		currentStatement.WriteString(line)
+		currentStatement.WriteString("\n")
+
+		// Count occurrences of $$ in the line (handles single-line $$ blocks too)
+		count := strings.Count(line, "$$")
+		if count%2 == 1 {
+			inDollarBlock = !inDollarBlock
+		}
+
+		if !inDollarBlock && strings.HasSuffix(trimmed, ";") {
+			statements = append(statements, currentStatement.String())
+			currentStatement.Reset()
+		}
+	}
+	if currentStatement.Len() > 0 {
+		statements = append(statements, currentStatement.String())
+	}
+	return statements
+}
+
 func Tx(t *testing.T, pool *pgxpool.Pool) pgx.Tx {
 	t.Helper()

@@ -201,43 +257,59 @@ func TruncateTables(t *testing.T, pool *pgxpool.Pool) {

 	ctx := context.Background()

-    tables := []string{
-        "booking_discounts",
-        "loyalty_redemptions",
-        "discount_campaigns",
-        "forgiven_no_shows",
-        "user_social_logins",
-        "verification_codes",
-        "booking_services",
-        "refunds",
-        "payments",
-        "user_saved_cards",
-        "square_deposits",
-        "affiliate_payouts",
-        "financial_aggregates",
-        "gift_cards",
-        "user_giftcard_balances",
-        "bookings",
-        "booking_edit_requests",
-        "user_patch_tests",
-        "patch_tests",
-        "services",
-        "admin_notifications",
-        "user_referrals",
-        "user_notification_preferences",
-        "time_blockers",
-        "working_hours",
-        "exceptional_group_applications",
-        "exceptional_working_hours",
-        "exceptional_working_hours_groups",
-        "business_settings",
-        "users",
-        "images",
-        "tags",
-    }
+	// Acquire a dedicated connection so advisory lock, truncations, and unlock all use the same session
+	conn, err := pool.Acquire(ctx)
+	if err != nil {
+		t.Fatalf("Failed to acquire connection for truncation: %v", err)
+	}
+	defer conn.Release()
+
+	_, err = conn.Exec(ctx, "SELECT pg_advisory_lock(1338)")
+	if err != nil {
+		t.Fatalf("Failed to acquire truncate advisory lock: %v", err)
+	}
+	defer conn.Exec(ctx, "SELECT pg_advisory_unlock(1338)")
+
+	tables := []string{
+		"till_sales",
+		"gift_card_transactions",
+		"gift_card_expired_balances",
+		"booking_discounts",
+		"loyalty_redemptions",
+		"discount_campaigns",
+		"forgiven_no_shows",
+		"user_social_logins",
+		"verification_codes",
+		"booking_services",
+		"refunds",
+		"payments",
+		"user_saved_cards",
+		"square_deposits",
+		"affiliate_payouts",
+		"financial_aggregates",
+		"gift_cards",
+		"user_giftcard_balances",
+		"bookings",
+		"booking_edit_requests",
+		"user_patch_tests",
+		"patch_tests",
+		"services",
+		"admin_notifications",
+		"user_referrals",
+		"user_notification_preferences",
+		"time_blockers",
+		"working_hours",
+		"exceptional_group_applications",
+		"exceptional_working_hours",
+		"exceptional_working_hours_groups",
+		"business_settings",
+		"users",
+		"images",
+		"tags",
+	}

 	for _, table := range tables {
-		_, err := pool.Exec(ctx, fmt.Sprintf("TRUNCATE TABLE %s CASCADE", table))
+		_, err := conn.Exec(ctx, fmt.Sprintf("TRUNCATE TABLE %s CASCADE", table))
 		if err != nil {
 			t.Logf("Warning: could not truncate %s: %v", table, err)
 		}