postgresql-setup.sh

#!/usr/bin/env bash

# Unoffiical bash strict mode.
# See: http://redsymbol.net/articles/unofficial-bash-strict-mode/
set -u
set -o pipefail
IFS=$'\n\t'

progname="$0"

case "$(uname)" in
  Linux)
    # Linux tools have long names for these options.
    recursive="--recursive"
    directory="--directory"
    force="--force"
    test="--test"
    ;;
  *)
    # These shoud work on any POSIX system
    recursive="-r"
    directory="-d"
    force="-f"
    test="-t"
  esac

# This should work on all Linux variants as well as FreeBSD.
numcores=$(getconf _NPROCESSORS_ONLN)
if test "${numcores}" -le 2 ; then
  numcores=1
else
  numcores=$(( numcores - 1 ))
  fi

function die {
	echo "$1"
	exit 1
}

function check_pgpass_file {
  if test -z ${PGPASSFILE+x} ; then
	echo "Error: The PGPASSFILE env var should be set to the location of the pgpass file."
	echo
	echo "Eg for mainnet:"
	echo "export PGPASSFILE=$(pwd)/config/pgpass"
	echo
	exit 1
	fi

  if test ! -f "${PGPASSFILE}" ; then
    echo "Error: PostgreSQL password file ${PGPASSFILE} does not exist."
    exit 1
    fi

  PGHOST=$(cut -d ":" -f 1 "${PGPASSFILE}")
  PGPORT=$(cut -d ":" -f 2 "${PGPASSFILE}")
  PGDATABASE=$(cut -d ":" -f 3 "${PGPASSFILE}")
  user=$(cut -d ":" -f 4 "${PGPASSFILE}")
  if [ "$user" != "*" ]; then
      PGUSER=$user
      export PGUSER
  fi;

  export PGHOST
  export PGPORT
  export PGDATABASE
}

function check_for_psql {
	# Make sure we have the psql executable.
	psql -V > /dev/null 2>&1 || die "Error : Missing 'psql' executable!"
}

function check_psql_superuser {
	set +e
	psql -l "${PGDATABASE}" > /dev/null 2>&1 || psql -l > /dev/null 2>&1
	if test $? -ne 0 ; then
		echo
		echo "Error : User '${PGUSER:-$(whoami)}' can't access postgres."
		echo
		echo "To fix this, log into the postgres account and run:"
		echo "    createuser --createdb --superuser ${PGUSER:-$(whoami)}"
		echo
		exit 1
		fi
	set -e
}

function check_connect_as_user {
	psql  "${PGDATABASE}" --no-password --command='\dt' > /dev/null
	if test $? -ne 0 ; then
		echo
		echo "Error : Not able to connect as '${PGUSER:-$(whoami)}' user."
		echo
		exit 1
		fi
}

function check_db_exists {
	set +e
	count=$(psql -t -l --csv "${PGDATABASE}" | grep -c "^${PGDATABASE},")
	if test "${count}" -lt 1 ; then
		echo
		echo "Error : No '${PGDATABASE}' database."
		echo
		echo "To create one run:"
		echo "    $progname --createdb"
		echo
		exit 1
		fi
	count=$(psql -t -l --csv "${PGDATABASE}" | grep "^${PGDATABASE}," | grep -c UTF8)
	if test "${count}" -ne 1 ; then
		echo
		echo "Error : '${PGDATABASE}' database exists, but is not UTF8."
		echo
		echo "To fix this you should drop the current one and create a new one using:"
		echo "    $progname --dropdb"
		echo "    $progname --createdb"
		echo
		exit 1
		fi
	set -e
}

function create_db {
	if test "$( psql "${PGDATABASE}" -tAc "SELECT 1 FROM pg_database WHERE datname='${PGDATABASE}'" )" != '1' ; then
		createdb -T template0 --owner="${PGUSER:-$(whoami)}" --encoding=UTF8 "${PGDATABASE}"
		fi
}



function drop_db {
	if test "$( psql "${PGDATABASE}" -tAc "SELECT 1 FROM pg_database WHERE datname='${PGDATABASE}'" )" = '1' ; then
		psql "${PGDATABASE}" --command="DROP OWNED BY CURRENT_USER cascade;"
	fi
}

function list_views {
	psql "${PGDATABASE}" \
		--command="select table_name from information_schema.views where table_catalog = '${PGDATABASE}' and table_schema = 'public' ;"
}

function create_migration {
	echo "To create a migration:"
	echo "cabal run cardano-db-tool -- create-migration --mdir schema/"
	exit 0
}

function run_migrations {
	echo "To run migrations:"
	echo "cabal run cardano-db-tool -- run-migrations --mdir schema/ --ldir ."
	echo "You probably do not need to do this."
	exit 0
}

function dump_schema {
	pg_dump -s --schema=public "${PGDATABASE}"
}

function create_snapshot {
	tgz_file=$1.tgz
	ledger_path=$2
	tmp_dir=$(mktemp "${directory}" -t db-sync-snapshot-XXXXXXXXXX)
	echo $"Working directory: ${tmp_dir}"
	# Process ledger state first to fail fast on missing/corrupt snapshots
	if [ -n "${ledger_path}" ]; then
	  if [ -d "${ledger_path}" ]; then
	    # New consensus directory format: tar+gzip the snapshot directory
	    ledger_dir_name=$(basename "${ledger_path}")
	    state_dir=$(dirname "${ledger_path}")
	    lsm_snap_dir="${state_dir}/lsm/snapshots/${ledger_dir_name}"
	    if [ -d "${lsm_snap_dir}" ]; then
	      echo "Detected LSM backend snapshot at slot ${ledger_dir_name}"
	      # LSM backend: bundle ledger snapshot, LSM table snapshot, and LSM session metadata (salt).
	      # The active/ directory is not needed — it gets cleared on session restore.
	      tar czf "${tmp_dir}/${ledger_dir_name}.tar.gz" \
	        -C "${state_dir}" "${ledger_dir_name}" "lsm/snapshots/${ledger_dir_name}" "lsm/metadata"
	    elif [ -d "${state_dir}/lsm" ]; then
	      # State directory contains an lsm/ subtree but the expected snapshot path is missing.
	      # Refuse to silently fall back to the InMemory branch — the resulting tarball would
	      # exclude the LSM tree files and restoration would replay from genesis.
	      echo "ERROR: ${state_dir}/lsm exists but no LSM snapshot found at ${lsm_snap_dir}." >&2
	      echo "       The LSM session may be using a non-standard layout. Aborting to avoid" >&2
	      echo "       producing an incomplete snapshot." >&2
	      rm "${recursive}" "${force}" "${tmp_dir}"
	      exit 1
	    else
	      echo "Detected InMemory backend snapshot at slot ${ledger_dir_name}"
	      # InMemory backend: just the ledger snapshot directory
	      tar czf "${tmp_dir}/${ledger_dir_name}.tar.gz" -C "${state_dir}" "${ledger_dir_name}"
	    fi
	  elif [ -f "${ledger_path}" ]; then
	    echo "Detected legacy .lstate snapshot: $(basename "${ledger_path}")"
	    # Legacy .lstate file format
	    lstate_gz_file=$(basename "${ledger_path}").gz
	    gzip --to-stdout "${ledger_path}" > "${tmp_dir}/${lstate_gz_file}"
	  else
	    echo "Ledger state path '${ledger_path}' does not exist."
	    rm "${recursive}" "${force}" "${tmp_dir}"
	    exit 1
	  fi
	fi
	pg_dump --no-owner --schema=public --jobs="${numcores}" "${PGDATABASE}" --format=directory --file="${tmp_dir}/db/"
	# Use plain tar here because the database dump files and the ledger state file are already gzipped. Disable Shellcheck SC2046 to avoid empty '' getting added while quoting
	# shellcheck disable=SC2046
	tar cvf - --directory "${tmp_dir}" $(ls "${tmp_dir}") | tee "${tgz_file}.tmp" | \
	  sha256sum | head -c 64 | sed -e "s/$/  ${tgz_file}\n/" > "${tgz_file}.sha256sum"
	mv "${tgz_file}.tmp" "${tgz_file}"
	rm "${recursive}" "${force}" "${tmp_dir}"
	if test "$(tar "${test}" --file "${tgz_file}")" ; then
	  echo "Tar reports the snapshot file as being corrupt."
	  echo "It is not safe to drop the database and restore using this file."
	  exit 1
	fi
	echo "Created ${tgz_file} + .sha256sum"
}

function restore_snapshot {
	# Create ledger state dir if it doesn't exist
	if ! test -d "$2" ; then
	  echo "Creating ledger state directory: $2"
	  mkdir -p "$2"
	fi
	# Check ledger state dir is empty (both old .lstate files and new snapshot dirs)
	file_count=$(find "$2" -maxdepth 1 -type f -name '*.lstate' -o -type d -name '[0-9]*' 2>/dev/null | wc -l)
	if test "${file_count}" -gt 0 ; then
	  echo "Ledger state directory ($2) is not empty. Please empty it and then retry."
	  exit 1
	fi
	tmp_dir=$(mktemp "${directory}" -t db-sync-snapshot-XXXXXXXXXX)
	tar xvf "$1" --directory "$tmp_dir"
	if test -d "${tmp_dir}/db/" ; then
	  # Check for new consensus snapshot directory format (tar.gz named after slot number)
	  snapshot_tgz=$(find "${tmp_dir}/" -maxdepth 1 -name '[0-9]*.tar.gz' | head -1)
	  if [ -n "${snapshot_tgz:-}" ] ; then
	    # New consensus directory format: extract snapshot tar.gz to ledger state dir
	    tar xzf "${snapshot_tgz}" -C "$2/"
	    # For LSM: create the active/ directory if lsm/ was restored (required by LSM session)
	    if [ -d "$2/lsm" ] && [ ! -d "$2/lsm/active" ]; then
	      echo "Restored LSM backend snapshot: $(basename "${snapshot_tgz}" .tar.gz)"
	      mkdir -p "$2/lsm/active"
	    else
	      echo "Restored InMemory backend snapshot: $(basename "${snapshot_tgz}" .tar.gz)"
	    fi
	  else
	    # Legacy .lstate file format
	    lstate_gz_file=$(find "${tmp_dir}/" -iname "*.lstate.gz")
	    if [ -n "${lstate_gz_file:-}" ] ; then
	      lstate_file=$(basename "${lstate_gz_file}" | sed 's/.gz$//')
	      echo "Restored legacy .lstate snapshot: ${lstate_file}"
	      gunzip --to-stdout "${lstate_gz_file}" > "$2/${lstate_file}"
	    fi
	  fi

	  # Important: specify --schema=public below to skip over `create schema public`
	  # statement generated by pg_dump
	  # shellcheck disable=SC2046
	  pg_restore \
	    --schema=public \
	    --jobs="${numcores}" \
	    --format=directory \
	    --dbname="${PGDATABASE}" \
	    --no-owner \
	    $( [ -z "${SKIP_RESTORE_ERROR:-}" ] && echo "--exit-on-error" ) \
	    "${tmp_dir}/db/"
	else
	  # Old snapshot format produced by this script
	  db_file=$(find "${tmp_dir}/" -iname "*.sql")
	  lstate_file=$(find "${tmp_dir}/" -iname "*.lstate")
	  mv "${lstate_file}" "$2"
	  psql --dbname="${PGDATABASE}" --file="${db_file}"
	fi
	rm "${recursive}" "${tmp_dir}"
}

function usage_exit {
	echo
	echo "Usage:"
	echo "    $progname --check             - Check database exists and is set up correctly."
	echo "    $progname --createdb          - Create database."
	echo "    $progname --dropdb            - Drop database."
	echo "    $progname --list-views        - List the currently definied views."
	echo "    $progname --recreatedb        - Drop and recreate database."
	echo "    $progname --create-user       - Create database user (from config/pgass file)."
	echo "    $progname --create-migration  - Create a migration (if one is needed)."
	echo "    $progname --run-migrations    - Run all migrations applying as needed."
	echo "    $progname --dump-schema       - Dump the schema of the database."
	echo
	echo "    - Create a db-sync state snapshot"
	echo "      $progname --create-snapshot <snapshot-file> [<ledger-state-file>]"
	echo
	echo "    - Restore a db-sync state snapshot."
	echo "      $progname --restore-snapshot <snapshot-file> <ledger-state-dir>"
	echo
	exit 0
}

# postgresql_version=$(psql -V | head -1 | sed -e "s/.* //;s/\.[0-9]*$//")

set -e

case "${1:-""}" in
	--check)
		check_pgpass_file
		check_for_psql
		check_psql_superuser
		check_db_exists
		check_connect_as_user
		;;
	--createdb)
		check_pgpass_file
		check_for_psql
		check_psql_superuser
		create_db
		;;
	--dropdb)
		check_pgpass_file
		check_for_psql
		check_psql_superuser
		drop_db
		;;
	--list-views)
		check_pgpass_file
		check_for_psql
		check_psql_superuser
		check_db_exists
		check_connect_as_user
		list_views
		;;
	--recreatedb)
		check_pgpass_file
		check_for_psql
		check_psql_superuser
		check_db_exists
		check_connect_as_user
		drop_db
		create_db
		echo "The database ${PGDATABASE} has been dropped and recreated."
		echo "The tables will be recreated when the application is run."
		exit 0
		;;
	--create-user)
		check_pgpass_file
		check_for_psql
		check_psql_superuser
		create_user
		;;
	--create-migration)
		create_migration
		;;
	--run-migrations)
		run_migrations
		;;
	--dump-schema)
		check_pgpass_file
		check_db_exists
		dump_schema
		;;
	--create-snapshot)
		check_pgpass_file
		check_db_exists
		if test $# -lt 2 ; then
		  echo "Expecting the snapshot file name (without extension) and optionally the ledger state file as arguments."
		  exit 1
		fi
		if test -z "$2" ; then
		  echo "Second argument should be the snapshot file name (without extension)."
		  exit 1
		fi
		# Third argument can be a file (.lstate) or directory (consensus snapshot format)
		create_snapshot "$2" "${3:-}"
		;;
	--restore-snapshot)
		check_pgpass_file
		check_for_psql
		check_psql_superuser
		if [ "${RESTORE_RECREATE_DB:-Y}" == "Y" ]; then
			drop_db
			create_db
		fi
		if test $# -ne 3 ; then
		  echo "Expecting exactly 2 more arguments, the snapshot file and the ledger state directory."
		  exit 1
		  fi
		if test -z "$2" ; then
		  echo "Second argument should be the snapshot file."
		  exit 1
		  fi
		if test -n "${3:-}" && test -f "${3}"; then
		  echo "Third argument is a file and expecting a directory."
		  exit 1
		  fi
		restore_snapshot "$2" "${3:-}"
		;;
	*)
		usage_exit
		;;
	esac

echo "All good!"
exit 0