begriffs open source - pg_scribe/blob - tests/test_restore.sh

   1 #!/usr/bin/env bash
   2 #
   3 # Test suite for pg_scribe --restore command
   4 #
   5 # This test suite:
   6 # - Creates temporary test databases
   7 # - Tests various --restore scenarios
   8 # - Verifies expected outcomes
   9 # - Cleans up all resources
  10 #
  11
  12 set -euo pipefail
  13
  14 # Colors for test output
  15 RED='\033[0;31m'
  16 GREEN='\033[0;32m'
  17 YELLOW='\033[0;33m'
  18 BLUE='\033[0;34m'
  19 NC='\033[0m' # No Color
  20
  21 # Test configuration
  22 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
  23 PG_SCRIBE="$SCRIPT_DIR/scripts/pg_scribe"
  24 TEST_DIR="/tmp/pg_scribe_test_restore_$$"
  25 TEST_DB_PREFIX="pg_scribe_restore_test_$$"
  26 PGUSER="${PGUSER:-postgres}"
  27
  28 # Test counters
  29 TESTS_RUN=0
  30 TESTS_PASSED=0
  31 TESTS_FAILED=0
  32
  33 # Cleanup tracking
  34 DATABASES_TO_CLEANUP=()
  35 SLOTS_TO_CLEANUP=()
  36
  37 #
  38 # Logging functions
  39 #
  40
  41 log_test() {
  42     echo -e "${BLUE}TEST:${NC} $*"
  43 }
  44
  45 log_pass() {
  46     echo -e "${GREEN}PASS:${NC} $*"
  47     ((TESTS_PASSED++))
  48 }
  49
  50 log_fail() {
  51     echo -e "${RED}FAIL:${NC} $*"
  52     ((TESTS_FAILED++))
  53 }
  54
  55 log_info() {
  56     echo -e "${YELLOW}INFO:${NC} $*"
  57 }
  58
  59 #
  60 # Helper functions
  61 #
  62
  63 run_psql() {
  64     local dbname="$1"
  65     shift
  66     psql -U "$PGUSER" -d "$dbname" -tAq "$@"
  67 }
  68
  69 query_db() {
  70     local dbname="$1"
  71     local query="$2"
  72     run_psql "$dbname" -c "$query" 2>/dev/null || true
  73 }
  74
  75 create_test_db() {
  76     local dbname="$1"
  77     log_info "Creating test database: $dbname"
  78
  79     # Drop if exists
  80     psql -U "$PGUSER" -d postgres -c "DROP DATABASE IF EXISTS $dbname;" &>/dev/null || true
  81
  82     # Create database
  83     psql -U "$PGUSER" -d postgres -c "CREATE DATABASE $dbname;" &>/dev/null
  84
  85     DATABASES_TO_CLEANUP+=("$dbname")
  86 }
  87
  88 # shellcheck disable=SC2317  # Function called from cleanup trap handler
  89 drop_test_db() {
  90     local dbname="$1"
  91     log_info "Dropping test database: $dbname"
  92
  93     # Terminate connections
  94     psql -U "$PGUSER" -d postgres -c "
  95         SELECT pg_terminate_backend(pid)
  96         FROM pg_stat_activity
  97         WHERE datname = '$dbname' AND pid <> pg_backend_pid();
  98     " &>/dev/null || true
  99
 100     # Drop database
 101     psql -U "$PGUSER" -d postgres -c "DROP DATABASE IF EXISTS $dbname;" &>/dev/null || true
 102 }
 103
 104 # shellcheck disable=SC2317  # Function called from cleanup trap handler
 105 drop_replication_slot() {
 106     local dbname="$1"
 107     local slot="$2"
 108     log_info "Dropping replication slot: $slot"
 109
 110     # Check if slot exists
 111     local exists
 112     exists=$(query_db "$dbname" "
 113         SELECT 1 FROM pg_replication_slots WHERE slot_name = '$slot';
 114     ")
 115
 116     if [[ -n "$exists" ]]; then
 117         # Drop slot
 118         query_db "$dbname" "SELECT pg_drop_replication_slot('$slot');" || true
 119     fi
 120 }
 121
 122 create_table_with_pk() {
 123     local dbname="$1"
 124     local table="$2"
 125     query_db "$dbname" "
 126         CREATE TABLE $table (
 127             id SERIAL PRIMARY KEY,
 128             name TEXT,
 129             created_at TIMESTAMP DEFAULT now()
 130         );
 131     "
 132 }
 133
 134 initialize_backup() {
 135     local dbname="$1"
 136     local backup_dir="$2"
 137     local slot="$3"
 138
 139     # Init backup system
 140     "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null
 141     SLOTS_TO_CLEANUP+=("$dbname:$slot")
 142 }
 143
 144 #
 145 # Test cases
 146 #
 147
 148 test_basic_restore_from_base_only() {
 149     ((TESTS_RUN++))
 150     log_test "Basic restore from base backup only"
 151
 152     local source_db="${TEST_DB_PREFIX}_source_basic"
 153     local restore_db="${TEST_DB_PREFIX}_restore_basic"
 154     local slot="test_slot_basic"
 155     local backup_dir="$TEST_DIR/basic"
 156
 157     # Setup source database
 158     create_test_db "$source_db"
 159     create_table_with_pk "$source_db" "users"
 160     query_db "$source_db" "INSERT INTO users (name) VALUES ('Alice'), ('Bob'), ('Charlie');"
 161
 162     # Initialize backup
 163     mkdir -p "$backup_dir"
 164     initialize_backup "$source_db" "$backup_dir" "$slot"
 165
 166     # Restore to new database with --create
 167     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 168         DATABASES_TO_CLEANUP+=("$restore_db")
 169
 170         # Verify data
 171         local count
 172         count=$(query_db "$restore_db" "SELECT COUNT(*) FROM users;")
 173         if [[ "$count" -ne 3 ]]; then
 174             log_fail "Expected 3 rows, got $count"
 175             return 1
 176         fi
 177
 178         # Verify specific data (psql returns one value per line)
 179         local names
 180         names=$(query_db "$restore_db" "SELECT name FROM users ORDER BY id;")
 181         local expected_names=$'Alice\nBob\nCharlie'
 182         if [[ "$names" != "$expected_names" ]]; then
 183             log_fail "Data mismatch: expected '$expected_names', got '$names'"
 184             return 1
 185         fi
 186
 187         log_pass "Basic restore successful"
 188         return 0
 189     else
 190         log_fail "Restore command failed"
 191         return 1
 192     fi
 193 }
 194
 195 test_restore_with_differentials() {
 196     ((TESTS_RUN++))
 197     log_test "Restore with base + sealed differentials"
 198
 199     local source_db="${TEST_DB_PREFIX}_source_diff"
 200     local restore_db="${TEST_DB_PREFIX}_restore_diff"
 201     local slot="test_slot_diff"
 202     local backup_dir="$TEST_DIR/differentials"
 203
 204     # Setup source database
 205     create_test_db "$source_db"
 206     create_table_with_pk "$source_db" "orders"
 207     query_db "$source_db" "INSERT INTO orders (name) VALUES ('Order1'), ('Order2');"
 208
 209     # Initialize backup (creates first chain)
 210     mkdir -p "$backup_dir"
 211     initialize_backup "$source_db" "$backup_dir" "$slot"
 212
 213     # Start streaming to latest chain
 214     "$PG_SCRIBE" --start -d "$source_db" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null &
 215     local pg_scribe_pid=$!
 216
 217     # Wait for streaming to start
 218     sleep 2
 219
 220     # Make some changes
 221     query_db "$source_db" "INSERT INTO orders (name) VALUES ('Order3');"
 222     query_db "$source_db" "UPDATE orders SET name = 'Order1_Updated' WHERE name = 'Order1';"
 223     query_db "$source_db" "DELETE FROM orders WHERE name = 'Order2';"
 224
 225     # Wait for changes to be captured
 226     sleep 2
 227
 228     # Rotate differential to seal it
 229     "$PG_SCRIBE" --rotate-diff -f "$backup_dir" -U "$PGUSER" &>/dev/null || true
 230     sleep 1
 231
 232     # Stop streaming
 233     kill -TERM "$pg_scribe_pid" 2>/dev/null || true
 234     wait "$pg_scribe_pid" 2>/dev/null || true
 235
 236     # Find latest chain
 237     local chain_dir
 238     chain_dir=$(find "$backup_dir" -maxdepth 1 -type d -name 'chain-*' 2>/dev/null | sort | tail -1)
 239
 240     # Verify sealed differential exists
 241     local diff_count
 242     diff_count=$(find "$chain_dir" -name 'diff-*.sql' 2>/dev/null | wc -l)
 243     if [[ "$diff_count" -eq 0 ]]; then
 244         log_fail "No sealed differentials found"
 245         return 1
 246     fi
 247
 248     # Restore to new database
 249     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 250         DATABASES_TO_CLEANUP+=("$restore_db")
 251
 252         # Verify final state (should have Order1_Updated and Order3, not Order2)
 253         local count
 254         count=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders;")
 255         if [[ "$count" -ne 2 ]]; then
 256             log_fail "Expected 2 rows after differential restore, got $count"
 257             return 1
 258         fi
 259
 260         # Verify Order1 was updated
 261         local order1_name
 262         order1_name=$(query_db "$restore_db" "SELECT name FROM orders WHERE id = 1;")
 263         if [[ "$order1_name" != "Order1_Updated" ]]; then
 264             log_fail "UPDATE not applied: expected 'Order1_Updated', got '$order1_name'"
 265             return 1
 266         fi
 267
 268         # Verify Order3 exists
 269         local order3_exists
 270         order3_exists=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders WHERE name = 'Order3';")
 271         if [[ "$order3_exists" -ne 1 ]]; then
 272             log_fail "INSERT not applied: Order3 not found"
 273             return 1
 274         fi
 275
 276         # Verify Order2 was deleted
 277         local order2_exists
 278         order2_exists=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders WHERE name = 'Order2';")
 279         if [[ "$order2_exists" -ne 0 ]]; then
 280             log_fail "DELETE not applied: Order2 still exists"
 281             return 1
 282         fi
 283
 284         log_pass "Restore with differentials successful"
 285         return 0
 286     else
 287         log_fail "Restore command failed"
 288         return 1
 289     fi
 290 }
 291
 292 test_restore_with_ddl_changes() {
 293     ((TESTS_RUN++))
 294     log_test "Restore with DDL changes in differential"
 295
 296     local source_db="${TEST_DB_PREFIX}_source_ddl"
 297     local restore_db="${TEST_DB_PREFIX}_restore_ddl"
 298     local slot="test_slot_ddl"
 299     local backup_dir="$TEST_DIR/ddl"
 300
 301     # Setup source database
 302     create_test_db "$source_db"
 303     create_table_with_pk "$source_db" "products"
 304     query_db "$source_db" "INSERT INTO products (name) VALUES ('Widget');"
 305
 306     # Initialize backup (creates chain)
 307     mkdir -p "$backup_dir"
 308     initialize_backup "$source_db" "$backup_dir" "$slot"
 309
 310     # Start streaming
 311     "$PG_SCRIBE" --start -d "$source_db" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null &
 312     local pg_scribe_pid=$!
 313
 314     # Wait for streaming to start
 315     sleep 2
 316
 317     # Add a column (DDL change)
 318     query_db "$source_db" "ALTER TABLE products ADD COLUMN price NUMERIC(10,2);"
 319
 320     # Insert data using new column
 321     query_db "$source_db" "INSERT INTO products (name, price) VALUES ('Gadget', 19.99);"
 322
 323     # Wait for changes to be captured
 324     sleep 2
 325
 326     # Rotate to seal differential
 327     "$PG_SCRIBE" --rotate-diff -f "$backup_dir" -U "$PGUSER" &>/dev/null || true
 328     sleep 1
 329
 330     # Stop streaming
 331     kill -TERM "$pg_scribe_pid" 2>/dev/null || true
 332     wait "$pg_scribe_pid" 2>/dev/null || true
 333
 334     # Restore to new database
 335     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 336         DATABASES_TO_CLEANUP+=("$restore_db")
 337
 338         # Verify new column exists
 339         local has_price_column
 340         has_price_column=$(query_db "$restore_db" "
 341             SELECT COUNT(*) FROM information_schema.columns
 342             WHERE table_name = 'products' AND column_name = 'price';
 343         ")
 344         if [[ "$has_price_column" -ne 1 ]]; then
 345             log_fail "DDL not applied: price column not found"
 346             return 1
 347         fi
 348
 349         # Verify data with new column
 350         local gadget_price
 351         gadget_price=$(query_db "$restore_db" "SELECT price FROM products WHERE name = 'Gadget';")
 352         if [[ "$gadget_price" != "19.99" ]]; then
 353             log_fail "Data with new column not correct: expected '19.99', got '$gadget_price'"
 354             return 1
 355         fi
 356
 357         log_pass "Restore with DDL changes successful"
 358         return 0
 359     else
 360         log_fail "Restore command failed"
 361         return 1
 362     fi
 363 }
 364
 365 test_restore_sequence_synchronization() {
 366     ((TESTS_RUN++))
 367     log_test "Restore with sequence synchronization"
 368
 369     local source_db="${TEST_DB_PREFIX}_source_seq"
 370     local restore_db="${TEST_DB_PREFIX}_restore_seq"
 371     local slot="test_slot_seq"
 372     local backup_dir="$TEST_DIR/sequence"
 373
 374     # Setup source database
 375     create_test_db "$source_db"
 376     create_table_with_pk "$source_db" "items"
 377     query_db "$source_db" "INSERT INTO items (name) VALUES ('Item1'), ('Item2'), ('Item3');"
 378
 379     # Initialize backup (creates first chain)
 380     mkdir -p "$backup_dir"
 381     initialize_backup "$source_db" "$backup_dir" "$slot"
 382
 383     # Add more data to advance sequence
 384     query_db "$source_db" "INSERT INTO items (name) VALUES ('Item4'), ('Item5');"
 385
 386     # Start streaming to capture incremental changes
 387     "$PG_SCRIBE" --start -d "$source_db" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null &
 388     local pg_scribe_pid=$!
 389     sleep 2
 390
 391     # Rotate differential to seal it
 392     "$PG_SCRIBE" --rotate-diff -f "$backup_dir" -U "$PGUSER" &>/dev/null || true
 393     sleep 1
 394
 395     # Stop streaming
 396     kill -TERM "$pg_scribe_pid" 2>/dev/null || true
 397     wait "$pg_scribe_pid" 2>/dev/null || true
 398
 399     # Restore to new database
 400     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 401         DATABASES_TO_CLEANUP+=("$restore_db")
 402
 403         # Get current sequence value
 404         local seq_val
 405         seq_val=$(query_db "$restore_db" "SELECT last_value FROM items_id_seq;")
 406
 407         # Should be at least 5 (we inserted 5 items total)
 408         if [[ "$seq_val" -lt 5 ]]; then
 409             log_fail "Sequence not synchronized: expected >= 5, got $seq_val"
 410             return 1
 411         fi
 412
 413         # Try inserting new row - should get ID 6
 414         query_db "$restore_db" "INSERT INTO items (name) VALUES ('Item6');"
 415         local new_id
 416         new_id=$(query_db "$restore_db" "SELECT id FROM items WHERE name = 'Item6';")
 417         if [[ "$new_id" -ne 6 ]]; then
 418             log_fail "Next sequence value incorrect: expected 6, got $new_id"
 419             return 1
 420         fi
 421
 422         log_pass "Sequence synchronization successful"
 423         return 0
 424     else
 425         log_fail "Restore command failed"
 426         return 1
 427     fi
 428 }
 429
 430 test_restore_no_sync_sequences() {
 431     ((TESTS_RUN++))
 432     log_test "Restore with --no-sync-sequences flag"
 433
 434     local source_db="${TEST_DB_PREFIX}_source_noseq"
 435     local restore_db="${TEST_DB_PREFIX}_restore_noseq"
 436     local slot="test_slot_noseq"
 437     local backup_dir="$TEST_DIR/noseq"
 438
 439     # Setup source database
 440     create_test_db "$source_db"
 441     create_table_with_pk "$source_db" "records"
 442     query_db "$source_db" "INSERT INTO records (name) VALUES ('Rec1'), ('Rec2');"
 443
 444     # Initialize backup (creates chain)
 445     mkdir -p "$backup_dir"
 446     initialize_backup "$source_db" "$backup_dir" "$slot"
 447
 448     # Restore with --no-sync-sequences
 449     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" --no-sync-sequences &>/dev/null; then
 450         DATABASES_TO_CLEANUP+=("$restore_db")
 451
 452         # Sequence should be at the value from pg_dump (2)
 453         local seq_val
 454         seq_val=$(query_db "$restore_db" "SELECT last_value FROM records_id_seq;")
 455
 456         # With --no-sync-sequences, sequence value is from pg_dump
 457         # It should match what's in the backup
 458         if [[ "$seq_val" -lt 1 ]]; then
 459             log_fail "Sequence value unexpectedly low: $seq_val"
 460             return 1
 461         fi
 462
 463         log_pass "Restore with --no-sync-sequences successful"
 464         return 0
 465     else
 466         log_fail "Restore command failed"
 467         return 1
 468     fi
 469 }
 470
 471 test_restore_to_existing_database() {
 472     ((TESTS_RUN++))
 473     log_test "Restore to existing database (without --create)"
 474
 475     local source_db="${TEST_DB_PREFIX}_source_exist"
 476     local restore_db="${TEST_DB_PREFIX}_restore_exist"
 477     local slot="test_slot_exist"
 478     local backup_dir="$TEST_DIR/existing"
 479
 480     # Setup source database
 481     create_test_db "$source_db"
 482     create_table_with_pk "$source_db" "data"
 483     query_db "$source_db" "INSERT INTO data (name) VALUES ('Test1');"
 484
 485     # Initialize backup
 486     mkdir -p "$backup_dir"
 487     initialize_backup "$source_db" "$backup_dir" "$slot"
 488
 489     # Pre-create target database
 490     create_test_db "$restore_db"
 491
 492     # Restore without --create
 493     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -U "$PGUSER" &>/dev/null; then
 494         # Verify data
 495         local count
 496         count=$(query_db "$restore_db" "SELECT COUNT(*) FROM data;")
 497         if [[ "$count" -ne 1 ]]; then
 498             log_fail "Expected 1 row, got $count"
 499             return 1
 500         fi
 501
 502         log_pass "Restore to existing database successful"
 503         return 0
 504     else
 505         log_fail "Restore command failed"
 506         return 1
 507     fi
 508 }
 509
 510 test_restore_fails_if_db_exists_with_create() {
 511     ((TESTS_RUN++))
 512     log_test "Restore fails if database exists with --create"
 513
 514     local source_db="${TEST_DB_PREFIX}_source_exists"
 515     local restore_db="${TEST_DB_PREFIX}_restore_exists"
 516     local slot="test_slot_exists"
 517     local backup_dir="$TEST_DIR/exists"
 518
 519     # Setup source database
 520     create_test_db "$source_db"
 521     create_table_with_pk "$source_db" "test"
 522     query_db "$source_db" "INSERT INTO test (name) VALUES ('Test');"
 523
 524     # Initialize backup
 525     mkdir -p "$backup_dir"
 526     initialize_backup "$source_db" "$backup_dir" "$slot"
 527
 528     # Pre-create target database
 529     create_test_db "$restore_db"
 530
 531     # Restore with --create should fail
 532     local exit_code=0
 533     "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null || exit_code=$?
 534
 535     if [[ $exit_code -eq 4 ]]; then
 536         log_pass "Correctly failed when database exists with --create"
 537         return 0
 538     else
 539         log_fail "Expected exit code 4, got $exit_code"
 540         return 1
 541     fi
 542 }
 543
 544 test_restore_missing_backup_directory() {
 545     ((TESTS_RUN++))
 546     log_test "Restore fails with missing backup directory"
 547
 548     local restore_db="${TEST_DB_PREFIX}_restore_missing"
 549     local backup_dir="$TEST_DIR/nonexistent"
 550
 551     # Try to restore from non-existent directory
 552     local exit_code=0
 553     "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -U "$PGUSER" &>/dev/null || exit_code=$?
 554
 555     if [[ $exit_code -eq 4 ]]; then
 556         log_pass "Correctly failed with missing backup directory"
 557         return 0
 558     else
 559         log_fail "Expected exit code 4, got $exit_code"
 560         return 1
 561     fi
 562 }
 563
 564 test_restore_specific_chain() {
 565     ((TESTS_RUN++))
 566     log_test "Restore specific chain with --base-backup (chain ID)"
 567
 568     local source_db="${TEST_DB_PREFIX}_source_specific"
 569     local restore_db="${TEST_DB_PREFIX}_restore_specific"
 570     local slot="test_slot_specific"
 571     local backup_dir="$TEST_DIR/specific"
 572
 573     # Setup source database
 574     create_test_db "$source_db"
 575     create_table_with_pk "$source_db" "entries"
 576     query_db "$source_db" "INSERT INTO entries (name) VALUES ('Entry1');"
 577
 578     # Initialize backup (creates first chain)
 579     mkdir -p "$backup_dir"
 580     initialize_backup "$source_db" "$backup_dir" "$slot"
 581
 582     # Find the chain ID
 583     local chain_dir
 584     chain_dir=$(find "$backup_dir" -maxdepth 1 -type d -name 'chain-*' 2>/dev/null | sort | tail -1)
 585
 586     if [[ -z "$chain_dir" ]]; then
 587         log_fail "Chain directory not found"
 588         return 1
 589     fi
 590
 591     local chain_id
 592     chain_id=$(basename "$chain_dir" | sed 's/^chain-//')
 593
 594     # Restore using specific chain ID
 595     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" --base-backup="$chain_id" -C -U "$PGUSER" &>/dev/null; then
 596         DATABASES_TO_CLEANUP+=("$restore_db")
 597
 598         # Verify data
 599         local count
 600         count=$(query_db "$restore_db" "SELECT COUNT(*) FROM entries;")
 601         if [[ "$count" -ne 1 ]]; then
 602             log_fail "Expected 1 row, got $count"
 603             return 1
 604         fi
 605
 606         log_pass "Restore with specific chain successful"
 607         return 0
 608     else
 609         log_fail "Restore command failed"
 610         return 1
 611     fi
 612 }
 613
 614 test_restore_multiple_tables() {
 615     ((TESTS_RUN++))
 616     log_test "Restore multiple tables with relationships"
 617
 618     local source_db="${TEST_DB_PREFIX}_source_multi"
 619     local restore_db="${TEST_DB_PREFIX}_restore_multi"
 620     local slot="test_slot_multi"
 621     local backup_dir="$TEST_DIR/multi"
 622
 623     # Setup source database with multiple related tables
 624     create_test_db "$source_db"
 625     query_db "$source_db" "
 626         CREATE TABLE customers (
 627             id SERIAL PRIMARY KEY,
 628             name TEXT NOT NULL
 629         );
 630         CREATE TABLE orders (
 631             id SERIAL PRIMARY KEY,
 632             customer_id INTEGER REFERENCES customers(id),
 633             product TEXT NOT NULL
 634         );
 635     "
 636     query_db "$source_db" "INSERT INTO customers (name) VALUES ('Alice'), ('Bob');"
 637     query_db "$source_db" "INSERT INTO orders (customer_id, product) VALUES (1, 'Widget'), (2, 'Gadget');"
 638
 639     # Initialize backup
 640     mkdir -p "$backup_dir"
 641     initialize_backup "$source_db" "$backup_dir" "$slot"
 642
 643     # Restore
 644     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 645         DATABASES_TO_CLEANUP+=("$restore_db")
 646
 647         # Verify customers
 648         local customer_count
 649         customer_count=$(query_db "$restore_db" "SELECT COUNT(*) FROM customers;")
 650         if [[ "$customer_count" -ne 2 ]]; then
 651             log_fail "Expected 2 customers, got $customer_count"
 652             return 1
 653         fi
 654
 655         # Verify orders
 656         local order_count
 657         order_count=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders;")
 658         if [[ "$order_count" -ne 2 ]]; then
 659             log_fail "Expected 2 orders, got $order_count"
 660             return 1
 661         fi
 662
 663         # Verify foreign key relationship
 664         local alice_orders
 665         alice_orders=$(query_db "$restore_db" "
 666             SELECT o.product FROM orders o
 667             JOIN customers c ON c.id = o.customer_id
 668             WHERE c.name = 'Alice';
 669         ")
 670         if [[ "$alice_orders" != "Widget" ]]; then
 671             log_fail "Foreign key relationship not preserved"
 672             return 1
 673         fi
 674
 675         log_pass "Multiple table restore successful"
 676         return 0
 677     else
 678         log_fail "Restore command failed"
 679         return 1
 680     fi
 681 }
 682
 683 #
 684 # Cleanup
 685 #
 686
 687 # shellcheck disable=SC2317  # Function called via trap handler
 688 cleanup() {
 689     log_info "Cleaning up test resources..."
 690
 691     # Stop any background pg_scribe processes
 692     pkill -f "pg_scribe.*--start" 2>/dev/null || true
 693     pkill -f "pg_recvlogical" 2>/dev/null || true
 694     sleep 1
 695
 696     # Drop replication slots
 697     for slot_info in "${SLOTS_TO_CLEANUP[@]}"; do
 698         IFS=':' read -r dbname slot <<< "$slot_info"
 699         drop_replication_slot "$dbname" "$slot" 2>/dev/null || true
 700     done
 701
 702     # Drop databases
 703     for dbname in "${DATABASES_TO_CLEANUP[@]}"; do
 704         drop_test_db "$dbname"
 705     done
 706
 707     # Remove test directory
 708     if [[ -d "$TEST_DIR" ]]; then
 709         rm -rf "$TEST_DIR"
 710     fi
 711
 712     log_info "Cleanup complete"
 713 }
 714
 715 #
 716 # Main test runner
 717 #
 718
 719 main() {
 720     echo "========================================"
 721     echo "pg_scribe --restore Test Suite"
 722     echo "========================================"
 723     echo ""
 724
 725     # Verify pg_scribe exists
 726     if [[ ! -x "$PG_SCRIBE" ]]; then
 727         echo "ERROR: pg_scribe not found or not executable: $PG_SCRIBE"
 728         exit 1
 729     fi
 730
 731     # Verify PostgreSQL is running
 732     if ! psql -U "$PGUSER" -d postgres -c "SELECT 1;" &>/dev/null; then
 733         echo "ERROR: Cannot connect to PostgreSQL"
 734         exit 1
 735     fi
 736
 737     # Verify wal_level is logical
 738     local wal_level
 739     wal_level=$(psql -U "$PGUSER" -d postgres -tAq -c "SHOW wal_level;")
 740     if [[ "$wal_level" != "logical" ]]; then
 741         echo "ERROR: wal_level must be 'logical', currently: $wal_level"
 742         echo "Update ~/.pgenv/pgsql/data/postgresql.conf and restart PostgreSQL"
 743         exit 1
 744     fi
 745
 746     # Verify wal2sql extension is available
 747     if ! psql -U "$PGUSER" -d postgres -c "CREATE EXTENSION IF NOT EXISTS wal2sql;" &>/dev/null; then
 748         echo "ERROR: wal2sql extension not available"
 749         echo "Build and install: cd wal2sql && make && make install"
 750         exit 1
 751     fi
 752
 753     # Create test directory
 754     mkdir -p "$TEST_DIR"
 755
 756     # Set up cleanup trap
 757     trap cleanup EXIT INT TERM
 758
 759     echo "Running tests..."
 760     echo ""
 761
 762     # Run all tests (use || true to prevent set -e from exiting)
 763     test_basic_restore_from_base_only || true
 764     test_restore_with_differentials || true
 765     test_restore_with_ddl_changes || true
 766     test_restore_specific_chain || true
 767     test_restore_to_existing_database || true
 768     test_restore_fails_if_db_exists_with_create || true
 769     test_restore_missing_backup_directory || true
 770     test_restore_multiple_tables || true
 771     test_restore_sequence_synchronization || true
 772     test_restore_no_sync_sequences || true
 773
 774     # Summary
 775     echo ""
 776     echo "========================================"
 777     echo "Test Results"
 778     echo "========================================"
 779     echo "Tests run:    $TESTS_RUN"
 780     echo -e "Tests passed: ${GREEN}$TESTS_PASSED${NC}"
 781     echo -e "Tests failed: ${RED}$TESTS_FAILED${NC}"
 782     echo ""
 783
 784     if [[ $TESTS_FAILED -eq 0 ]]; then
 785         echo -e "${GREEN}All tests passed!${NC}"
 786         exit 0
 787     else
 788         echo -e "${RED}Some tests failed!${NC}"
 789         exit 1
 790     fi
 791 }
 792
 793 main "$@"