databuild/tests/end_to_end/lib/db_utils.sh

#!/bin/bash

# Database utilities for DataBuild end-to-end tests

set -euo pipefail

# Source test utilities
source "$(dirname "${BASH_SOURCE[0]}")/test_utils.sh"

# Create isolated test database
create_test_database() {
    local db_name="$1"
    local test_dir="${DATABUILD_TEST_DIR:-/tmp}"
    local db_path="$test_dir/${db_name}.db"
    
    # Remove existing database if it exists
    rm -f "$db_path"
    
    # Create directory if needed
    mkdir -p "$(dirname "$db_path")"
    
    log_info "Created test database: $db_path" >&2
    echo "$db_path"
}

# Extract build events from database
extract_build_events() {
    local db_path="$1"
    local output_file="$2"
    
    if [[ ! -f "$db_path" ]]; then
        log_error "Database not found: $db_path"
        return 1
    fi
    
    # Extract events to JSON format
    sqlite3 "$db_path" -json \
        "SELECT event_type, partition_ref, job_name, timestamp, metadata 
         FROM build_events 
         ORDER BY timestamp, partition_ref, job_name;" > "$output_file"
    
    log_info "Extracted build events to: $output_file"
}

# Compare two databases' build events
compare_build_events() {
    local db1="$1"
    local db2="$2"
    local temp_dir="${DATABUILD_TEST_DIR:-/tmp}"
    
    local events1="$temp_dir/events1.json"
    local events2="$temp_dir/events2.json"
    
    # Extract events from both databases
    extract_build_events "$db1" "$events1"
    extract_build_events "$db2" "$events2"
    
    # Compare events (ignoring timestamps)
    if compare_files_flexible "$events1" "$events2"; then
        log_info "Build events match between databases"
        return 0
    else
        log_error "Build events differ between databases"
        log_info "Events from $db1:"
        cat "$events1"
        log_info "Events from $db2:"
        cat "$events2"
        return 1
    fi
}

# Compare core build events (job and partition only, ignoring request orchestration)
compare_core_build_events() {
    local db1="$1"
    local db2="$2"
    
    log_info "Comparing core build events (job and partition) between databases" >&2
    
    # Count core events by type
    local db1_job=$(count_build_events "$db1" "job")
    local db1_partition=$(count_build_events "$db1" "partition")
    local db2_job=$(count_build_events "$db2" "job") 
    local db2_partition=$(count_build_events "$db2" "partition")
    
    log_info "Database 1: $db1_job job events, $db1_partition partition events" >&2
    log_info "Database 2: $db2_job job events, $db2_partition partition events" >&2
    
    # Check if core events match
    if [[ "$db1_job" -eq "$db2_job" ]] && [[ "$db1_partition" -eq "$db2_partition" ]]; then
        log_info "Core build events match" >&2
        return 0
    else
        log_error "Core build events differ" >&2
        return 1
    fi
}

# Get detailed event breakdown
get_event_breakdown() {
    local db_path="$1"
    
    if [[ ! -f "$db_path" ]]; then
        echo "Database not found: $db_path" >&2
        return 1
    fi
    
    echo "Event breakdown for $db_path:" >&2
    sqlite3 "$db_path" \
        "SELECT event_type, COUNT(*) FROM build_events GROUP BY event_type ORDER BY event_type;" >&2
}

# Count events in database
count_build_events() {
    local db_path="$1"
    local event_type="${2:-}"
    
    if [[ ! -f "$db_path" ]]; then
        echo "0"
        return
    fi
    
    local query="SELECT COUNT(*) FROM build_events"
    if [[ -n "$event_type" ]]; then
        query="$query WHERE event_type = '$event_type'"
    fi
    
    sqlite3 "$db_path" "$query"
}

# Get partition status from database
get_partition_status() {
    local db_path="$1"
    local partition_ref="$2"
    
    if [[ ! -f "$db_path" ]]; then
        echo "NOT_FOUND"
        return
    fi
    
    # Get the latest event for this partition
    local status=$(sqlite3 "$db_path" \
        "SELECT event_type FROM build_events 
         WHERE partition_ref = '$partition_ref' 
         ORDER BY timestamp DESC 
         LIMIT 1;")
    
    echo "${status:-NOT_FOUND}"
}

# Wait for partition to reach expected status
wait_for_partition_status() {
    local db_path="$1"
    local partition_ref="$2"
    local expected_status="$3"
    local timeout="${4:-30}"
    
    local count=0
    while [[ $count -lt $timeout ]]; do
        local status=$(get_partition_status "$db_path" "$partition_ref")
        if [[ "$status" == "$expected_status" ]]; then
            log_info "Partition $partition_ref reached status: $expected_status"
            return 0
        fi
        
        sleep 1
        ((count++))
    done
    
    log_error "Partition $partition_ref did not reach status $expected_status within $timeout seconds"
    return 1
}

# Get all partitions built in database
get_built_partitions() {
    local db_path="$1"
    
    if [[ ! -f "$db_path" ]]; then
        return
    fi
    
    sqlite3 "$db_path" \
        "SELECT DISTINCT partition_ref FROM build_events 
         WHERE event_type = 'COMPLETED' 
         ORDER BY partition_ref;"
}

# Check if partition was built
is_partition_built() {
    local db_path="$1"
    local partition_ref="$2"
    
    local status=$(get_partition_status "$db_path" "$partition_ref")
    [[ "$status" == "COMPLETED" ]]
}

# Get build events for specific partition
get_partition_events() {
    local db_path="$1"
    local partition_ref="$2"
    local output_file="$3"
    
    if [[ ! -f "$db_path" ]]; then
        echo "[]" > "$output_file"
        return
    fi
    
    sqlite3 "$db_path" -json \
        "SELECT event_type, partition_ref, job_name, timestamp, metadata 
         FROM build_events 
         WHERE partition_ref = '$partition_ref' 
         ORDER BY timestamp;" > "$output_file"
}

# Validate database schema
validate_database_schema() {
    local db_path="$1"
    
    if [[ ! -f "$db_path" ]]; then
        log_error "Database not found: $db_path"
        return 1
    fi
    
    # Check if build_events table exists
    local table_exists=$(sqlite3 "$db_path" \
        "SELECT name FROM sqlite_master 
         WHERE type='table' AND name='build_events';")
    
    if [[ -z "$table_exists" ]]; then
        log_error "build_events table not found in database"
        return 1
    fi
    
    # Check required columns
    local columns=$(sqlite3 "$db_path" "PRAGMA table_info(build_events);" | cut -d'|' -f2)
    local required_columns=("event_type" "partition_ref" "job_name" "timestamp" "metadata")
    
    for col in "${required_columns[@]}"; do
        if ! echo "$columns" | grep -q "$col"; then
            log_error "Required column '$col' not found in build_events table"
            return 1
        fi
    done
    
    log_info "Database schema validation passed"
    return 0
}

# Clear all events from database
clear_build_events() {
    local db_path="$1"
    
    if [[ -f "$db_path" ]]; then
        sqlite3 "$db_path" "DELETE FROM build_events;"
        log_info "Cleared all build events from database"
    fi
}

# Get database size
get_database_size() {
    local db_path="$1"
    
    if [[ -f "$db_path" ]]; then
        stat -f%z "$db_path"
    else
        echo "0"
    fi
}

# Export database to SQL dump
export_database() {
    local db_path="$1"
    local output_file="$2"
    
    if [[ ! -f "$db_path" ]]; then
        log_error "Database not found: $db_path"
        return 1
    fi
    
    sqlite3 "$db_path" .dump > "$output_file"
    log_info "Exported database to: $output_file"
}