#!/bin/bash -e

. "$(dirname "${BASH_SOURCE[0]}")"/util.sh

.rel sync.sh

if self_not_included; then

# using prefixed connection vars

alias ssh2local='declare prefix=ssh_; import_vars; unset prefix

declare $(prefix=ssh_ get_prefix_vars); unset $(prefix=ssh_ get_prefix_vars)'

alias use_local='declare prefix=local_; import_vars; unset prefix'

alias use_remote='declare prefix=remote_; import_vars; unset prefix'

alias use_local_remote='{ use_remote; use_local; }'

	# *must* be run inside a function

alias use_root='declare prefix=root_; import_vars; unset prefix'

quote='"'

esc_name() { echo "$quote${1//$quote/$quote$quote}$quote"; }

mk_esc_name_alias() # usage: mk_esc_name_alias schema_esc

{ alias mk_"$1"='declare '"$1"'="${'"$1"':-$(esc_name "$'"${1%_esc}"'")}"; '\

'echo_vars '"$1"; }

mk_esc_name_alias schema_esc

mk_esc_name_alias table_esc

fi # load new aliases

if self_being_included; then

skip_table() # usage: create_table || table=... skip_table || return 0

{ what="table $table" already_exists_msg; }

: "${test=$(isset limit; exit2bool)}" # test mode when using limited # rows

# usage: (set_large_table; to_file db_cmd...) || return

alias set_large_table='declare del="${del-$(test "$limit"; exit2bool)}"

echo_vars del'

limit() # usage: "query... $([prefix=$'| |\n'] [limit|n=#] limit)"

{

	echo_func; kw_params prefix; local prefix="${prefix-

}"

	if isset n; then local limit="${limit-$n}"; fi

	echo -n "${limit:+${prefix}LIMIT $limit}"

}

cols2list() # usage: cols2list col...

{ echo_func; cmd=esc_name foreach_arg; delim=', ' join "$@"; }

mk_select() # usage: {query=... | table=... [cols=...] [filter=...] \

# [order_by=...]} [limit|n=#] mk_select

{

	echo_func; kw_params query table cols filter; mk_table_esc

	if is_array cols    ; then cols="$(    cols2list "${cols[@]}"    )"; fi

	if is_array order_by; then order_by="$(cols2list "${order_by[@]}")"; fi

	echo "$(rtrim "${query:-SELECT ${cols:-*} ${cols:+

}FROM $table_esc

${filter:+WHERE $filter

}${order_by:+ORDER BY $order_by

}}")\

$(limit)"

}

# export func usage: export_func() { ...; mk_select_var; ... }

# caller usage: {query=... | table=... [cols=...] [filter=...]} export_func

# cmd line usage: [limit=...] caller

alias mk_select_var='declare query="$(mk_select)"'

fi # load new aliases

if self_being_included; then

mk_drop() # usage: table=... mk_drop|*sql_ANSI

{

	log_local; log++; echo_func; kw_params table; : "${table?}"; mk_table_esc

	echo "DROP TABLE IF EXISTS $table_esc"

}

mk_truncate() # usage: table=... mk_truncate|*sql_ANSI

{

	log_local; log++; echo_func; kw_params table; : "${table?}"; mk_table_esc

	echo "TRUNCATE $table_esc"

}

### MySQL

alias set_database=\

'if test "$schema"; then declare database="${database-$schema}"; fi'

fi # load new aliases

if self_being_included; then

# usage: mysql*() { ...; mysql_cmd "$@"; } (with alias)

# caller usage: [ssh_server=...] [server=...] [user=...] [password=...] mysql*()

function mysql_cmd() # usage: mysql*() { ...; mysql_cmd cmd "$@"; }

# auto-adds connection/login opts when specified

{

	echo_func

	local cmd="$1"; shift

	local ssh_server="$(localize_url "$ssh_server")"

	local server="$(localize_url "$server")"

	if test "$ssh_server"; then

		local ssh_dest="${ssh_dest-${ssh_user:+$ssh_user@}$ssh_server}"

	fi

	local var=ssh_dest; local_inv

	command "time" ${ssh_dest:+ssh "$ssh_dest" }"$cmd" \

${server:+ --host="$server" }${user:+--user="$user" } --password\

${password+="$password"} --quick "$@" # --quick: don't buffer entire result

}

alias mysql_cmd='mysql_cmd "$FUNCNAME"'

fi # load new aliases

if self_being_included; then

function mysql() # usage: [database=...] [data_only=1] [log_queries=] mysql ...

{

	echo_func; kw_params output_data data_only query_verbosity

	if test "$data_only"; then local output_data="${output_data-1}"; fi

	local log_queries="${log_queries-1}"

	set_database

	set -- ${database:+--database="$database" }--local-infile=1 \

--${data_only:+skip-}column-names "$@"

	if test "$output_data"; then echo_stdin|mysql_cmd --batch "$@"

	else cmd_log_fd=1 mysql_cmd ${log_queries:+--verbose --verbose --verbose }\

"$@" # --verbose*3: echo runtimes

	fi

}

mysql_ANSI()

{

	echo_func

	(echo "SET sql_mode = 'ANSI';"; cat)|mysql "$@"

}

# always use ANSI mode, to support "" identifiers. `` *are* still supported in

# this mode, so it also works with SHOW CREATE TABLE output and dumpfiles.

alias mysql='mysql_ANSI'

fi # load new aliases

if self_being_included; then

mysql_root() { echo_func; use_root; mysql "$@"; }

mysql_truncate() { echo_func; mk_truncate|mysql; }

mysql_import() # usage: table=... [cols=...] [append=1] mysql_import <file

# without append=1, first ensures the table is empty

{

	echo_func

	mk_table_esc

	local mysql_load_data_format="${mysql_load_data_format-\

FIELDS TERMINATED BY ','

OPTIONALLY ENCLOSED BY '\"'

}"

	if test ! "$append"; then mysql_truncate; fi

	mysql_load_data_format="${mysql_load_data_format%

}"

	ssh2local # ssh does not tunnel nonstandard fds

	mysql "$@" 40<&0 <<EOF

LOAD DATA LOCAL INFILE '/dev/fd/40'

${append:+IGNORE }INTO TABLE $table_esc

$mysql_load_data_format

IGNORE 1 LINES

EOF

}

mysql_export() # does not support CSV

# caller usage: [database=...] {query=... | table=... [cols=...] [filter=...]} \

# mysql_export

# cmd line usage: [limit=...] caller

{

	echo_func

	mk_select_var

	output_data=1 mysql "$@" <<<"$query"

}

mysql_export_outfile() # supports CSV, but requires the FILE privilege

{

	echo_func

	: "${file:?}"

	mk_select_var

	local mysql_load_data_format="${mysql_load_data_format-\

FIELDS TERMINATED BY ','

OPTIONALLY ENCLOSED BY '\"'

}"

	local head="${query%%FROM*}" # includes trailing newline

	head="${head%

}"

	local tail="${query#$head}"

	mysql_load_data_format="${mysql_load_data_format%

}"

	mysql "$@" <<EOF

$head

INTO OUTFILE '$file'

$mysql_load_data_format

$tail

EOF

}

## backups

mysqldump() # usage: [schema=1 | data=1] [create_db=1] mysqldump db [table...]

{

	echo_func

	echo_vars schema data create_db

	local var=create_db; local_inv

	mysql_cmd ${database:+--databases "$database" ${no_create_db:+--tables }} \

--lock-tables=false --set-charset \

${postgres_compat:+${data:+--compatible=postgresql }--add-locks=false }\

${schema:+--no-data }${data:+--no-create-info }"$@"

}

mysqldump_diffable()

{

	echo_func

	mysqldump "$@"|{ pipe_delay; echo_run sed 's/^(-- Dump completed).*$/\1/'; }

}

mysql_rm_privileged_statements()

{

	echo_func

	grep -vE -e '^/\*!40000 ALTER TABLE `.*` (DIS|EN)ABLE KEYS \*/;$' \

-e '^(UN)?LOCK TABLES( `.*` WRITE)?;$'

}

## permissions

mysql_seal_table() # usage: table=... user=... mysql_seal_table

# prevents further modifications to a table by a user

{

	echo_func; kw_params table user; : "${table:?}" "${user:?}"; mk_table_esc

	benign_error=1 mysql_root <<<"REVOKE ALL PRIVILEGES ON $table FROM '$user'@'%'" || true

	benign_error=1 mysql_root <<<"REVOKE GRANT OPTION   ON $table FROM '$user'@'%'" || true

}

## server admin

mysql_ctl() # usage: mysql_ctl {start|stop|restart|...}

{

	echo_func

	pattern='^stop: Unknown instance:' ignore_e=1 ignore_err_msg \

sudo service mysql "$@" # ignore errors if not running

}

mysql_snapshot() # usage: [live=] [from=...] [to=...] mysql_snapshot rsync_opts

{

	echo_func; kw_params from to; local from="${from:-/var/lib/mysql}"

	ctl=mysql_ctl db_snapshot "$@"

}

### PostgreSQL

alias use_pg='declare prefix=pg_; import_vars; unset prefix'

fi # load new aliases

if self_being_included; then

# usage: pg_*() { ...; pg_cmd "$@"; } (uses alias)

function pg_cmd() # usage for fn: pg_*() { ...; pg_cmd cmd "$@"; }

# auto-adds connection/login opts when specified

{

	echo_func; kw_params as_root

	use_pg

	log_local

	log--

	if test "$server";   then local PGHOST="$server";       export PGHOST; fi

	if test "$user";     then local PGUSER="$user";         export PGUSER; fi

	if test "$password"; then local PGPASSWORD="$password"; log++ export PGPASSWORD;fi

	if test "$database"; then local PGDATABASE="$database"; export PGDATABASE;fi

	log++

	time if test "$as_root"; then sudo -u postgres "$@"; else command "$@"; fi

}

alias pg_cmd='"pg_cmd" "${FUNCNAME%%__*}"'

fi # load new aliases

if self_being_included; then

psql() # usage: [stdin=copy_pstdin_file|<(pstdin_cmd)] [output_data=1] psql \

# <<<"cmds (eg. \copy from pstdin)"

# $as_root: when on, avoids redirections as these are not passed through by sudo

# SUDO_USER: when set, avoids outputting to /dev/fd/#, because this causes a

# "Permission denied" error when running as sudo on Linux

{

	echo_func; kw_params stdin output_data

	local can_redir="$(! test "$as_root" && ! isset SUDO_USER; exit2bool)"

		echo_vars can_redir

	if test "$stdin"; then

		test "$can_redir" || die 'custom $stdin not supported with sudo'

	fi

	local verbose_="$(test "$can_redir" && can_log; exit2bool)"

		echo_vars verbose_

	local data2stdout="$(test "$output_data" -a "$can_redir"; exit2bool)"

		echo_vars data2stdout

	if test "$verbose_"; then set -- --echo-all --echo-hidden "$@"; fi

	if test "$can_redir"; then

		local redirs=("${redirs[@]}" '40<&0' "0<${stdin:-&20}" '41>&1')

	fi

	(

		# hide stack traces/DETAIL sections of error messages at verbosity <2

		if ! log++ can_log; then echo '\set VERBOSITY terse'; fi

		if test "$verbose_"; then echo '\timing on'; fi

		echo "SET client_min_messages = \

$(if test "$verbose_"; then echo NOTICE; else echo WARNING; fi);"

		cat

	)|cmd_log_fd=${can_redir:+1} pg_cmd ${can_redir:+--file /dev/fd/40 }\

${data2stdout:+--output /dev/fd/41 }--set ON_ERROR_STOP=1 --quiet "$@"\

|| verbosity_min=2 die_error_hidden

		# --output is for query *results*, not echoed statements

}

pg_export() # usage: mk_select_opts... [pg_copy_format=CSV...] pg_export >out

{

	echo_func

	mk_select_var

	local pg_copy_format="${pg_copy_format-CSV HEADER}"

	output_data=1 psql "$@" <<<"COPY ($query) TO STDOUT $pg_copy_format;"

}

pg_header()

{

	echo_func

	local pg_copy_format="CSV HEADER" limit=0

	pg_export "$@"|echo_stdout

}

pg_export_table_no_header()

{

	echo_func

	local pg_copy_format="CSV"

	pg_export "$@"

}

pg_export_table_to_dir()

{

	echo_func

	local table="$1"; shift; mk_table_esc

	stdout="$exports_dir/$table.csv" to_file pg_export "$@"

}

pg_export_table_to_dir_no_header()

{

	echo_func

	local table="$1"; shift; mk_table_esc

	stdout="$exports_dir/$table.no_header.cols=$(pg_header).csv" to_file \

pg_export_table_no_header "$@"

}

## backups

pg_dump() # usage: database={...|} [schema=...] [struct=1 | data=1] [owners=1] \

# [compress=1] [create_schema=] [users=1] pg_dump [opts...] >out

# database='': entire cluster

{

	echo_func;kw_params database struct data owners compress create_schema users

	: "${database?}"

	local create_schema="${create_schema-1}"

	# subset

	if test "$schema"; then set -- --schema="\"$schema\"" "$@"; fi

	if test "$struct"; then set -- "$@" --schema-only; fi

	if test "$data"  ; then set -- "$@" --data-only  ; fi

	if test "$database"; then # just one DB

		# format

		if test ! "$owners"; then set -- "$@" --no-owner   ; fi

		if test "$compress"; then set -- "$@" --format=custom --compress=9

							 else set -- "$@" --format=plain --inserts

		fi

		(

			# filter output

			if test ! "$create_schema"; then

				fd=1 redirs= filter_fd sed 's/^CREATE SCHEMA/--&/'

			fi

			pattern='^pg_dump: No matching tables were found$' ignore_err_msg \

pg_cmd "$@"

		)

	else # entire cluster

		# subset

		if test "$users"; then set -- "$@" --globals-only; fi

		(cd /; as_root=1 "pg_cmd" pg_dumpall "$@")

	fi

}

## DB structure

pg_schema_exists() # usage: schema=__ [benign_error=1] pg_schema_exists

{

	log-- echo_func; : "${schema:?}"; mk_schema_esc

	pattern='cannot create temporary relation in non-temporary schema' \

ignore_e=3 stderr_matches psql <<<"CREATE TEMP TABLE $schema_esc.t ()"

}

pg_require_schema() # usage: schema=... pg_require_schema

{

	echo_func; : "${schema:?}"

	pg_schema_exists || die "schema $schema does not exist"

}

pg_table_exists() # usage: [schema=__] table=__ [benign_error=1] pg_table_exists

{

	echo_func; : "${table:?}"; mk_table_esc

	! pattern='relation .* does not exist' \

ignore_e=3 log++ stderr_matches psql <<<"SELECT NULL FROM $table_esc LIMIT 0"

}

## server admin

pg_ctl() # usage: pg_ctl {start|stop|restart|...}

{

	echo_func

	sudo service postgresql "$@"

}

pg_start_backup() # requires config `wal_level = archive|hot_standby`

{ echo_func; as_root=1 psql <<<"SELECT pg_start_backup('backup', true);"; }

pg_stop_backup() # requires pg_start_backup() to have run

{ echo_func; as_root=1 psql <<<"SELECT pg_stop_backup();"; }

pg_snapshot() # usage: [live=] [from=...] [to=...] pg_snapshot rsync_opts...

{

	echo_func; kw_params from to; local from="${from:-/var/lib/postgresql}"

	if pg_start_backup; then # perform online backup

		try db_copy "$@"

		pg_stop_backup

		end_try

	else ctl=pg_ctl db_snapshot "$@" # fall back to stopping server

	fi

}

fi