run a remote postgresql query, compress, and ship csv back as zip

query-copy.sh

#!/bin/bash
# Retry function for SSH-related commands
retry_ssh_command() {
  local max_retries=5  # Maximum number of retries
  local sleep_time=2   # Time to wait between retries
  local attempt=1    # Attempt counter
  local cmd="$1"    # The command to execute
  local exit_code
  while [[ $attempt -le $max_retries ]]; do
    #echo "Attempt $attempt: $cmd"
    eval "$cmd"
    exit_code=$?
    # Check exit code
    if [[ $exit_code -eq 255 || $exit_code -eq 12 || $exit_code -eq 1 ]]; then
      echo "SSH connection error (exit code $exit_code). Retrying in $sleep_time seconds..."
    elif [[ $exit_code -ne 0 ]]; then
      echo "Remote command failed with exit code $exit_code. Not retrying."
      return $exit_code
    else
      return 0 # Command succeeded
    fi
    ((attempt++))
    sleep $sleep_time
  done
  echo "SSH command failed after $max_retries retries."
  return 255
}
# Main script starts here
# Variables
thost="$1"      # Target host
qryfn="$2"      # Query file name
output_file="$3"   # Output file name (optional)

if [[ -z "$thost" || -z "$qryfn" ]]; then
  echo "Usage: $0 <remote_host> <query_file> [output_file]"
  exit 1
fi

# Extract base name without .sql extension (needed for both output file and zip content)
base_name=$(basename "$qryfn" .sql)

# Get ISO timestamp without timezone
timestamp=$(date '+%Y-%m-%d_%H-%M-%S')

# Generate output filename if not provided
if [[ -z "$output_file" ]]; then
  output_file="${thost}-${base_name}-${timestamp}.zip"
fi

echo "* EXECUTING QUERY $qryfn on $thost AND COMPRESSING RESULTS"

# Execute query and pipe directly to zip
query_output=$(mktemp)
(echo -n '\copy ('; cat "$qryfn" | tr '\n' ' '; echo -n ') TO STDOUT CSV HEADER') | \
retry_ssh_command "ssh \"$thost\" \"./psql.sh\"" > "$query_output"
query_status=$?

if [[ $query_status -ne 0 ]]; then
    echo "Query execution failed (status: $query_status). Aborting."
    rm -f "$query_output"
    exit 1
fi

if [[ ! -s "$query_output" ]]; then
    echo "Query returned no data. Aborting."
    rm -f "$query_output"
    exit 1
fi

# Create CSV file first
csv_file="${thost}-${base_name}-${timestamp}.csv"
mv "$query_output" "$csv_file"

# Create zip file with CSV
zip -q "$output_file" "$csv_file"
zip_status=$?
rm -f "$csv_file"

if [[ $zip_status -ne 0 ]]; then
    echo "Compression failed (status: $zip_status). Aborting."
    rm -f "$output_file"
    exit 1
fi

# Verify zip file was created and is valid
if [[ ! -s "$output_file" ]]; then
    echo "Failed to create zip file. Aborting."
    rm -f "$output_file"
    exit 1
fi

# List the contents to verify
unzip -l "$output_file" || {
    echo "Created zip file is invalid. Aborting."
    rm -f "$output_file"
    exit 1
}
# Done
echo "* ALL DONE! Output written to $output_file"