How to conditionally stop a psql script (based on a variable value)?

Question

Let's consider the following example (from the start of a psql script):

\c :db_to_run_on

TRUNCATE the_most_important_table;
-- tried to avoid similarities to anything that exists out there

Now if it is run this by the command

psql [connection details] -v db_to_run_on=\'dev_database\'

then it just runs and the user is happy. But what if (s)he decides to specify -v db_to_run_on=production_database? (Let's assume that this can happen, just like people run rm -rf / # don't try this at home!!! ocassionally.) Hopefully there is a fresh backup of that table...

So the question arises: how to check the variables passed to a script and stop further processing based on their value?

Answer 1

PostgreSQL 10

PostgreSQL 10 brings conditionals to psql. This is no longer an issue.

\if :db_to_run_on = 'dev_database'
  TRUNCATE the_most_important_table;
\endif

I guess you could also use DO..

\if :db_to_run_on != 'dev_database'
do $$
  BEGIN
    RAISE 'Meaningful error message here';
  END;
$$ LANGUAGE plpgsql;
\endif

Answer 2

There is an option in psql which stops executing commands on error, this is ON_ERROR_STOP. If we could raise an error somehow, this would do what we want.

The problem is that we have to test the variable and produce an error somehow. Since one can't use control structures in psql (because there are none)*, my only idea was to use SQL for testing. Well, producing an error conditionally is something which pl/pgsql is quite good at, so I wrote a function which would generate an error. I can now call this function from a simple CASE structure. A simple example:

-- let's assume for clarity that there is no function with this name in the database
CREATE OR REPLACE FUNCTION error_generator()
RETURNS boolean AS
$body$
BEGIN
    RAISE 'Meaningful error message here';
    RETURN FALSE; -- just for aesthetical purposes
END;
$body$
LANGUAGE plpgsql;

\set ON_ERROR_STOP on

BEGIN;

-- test for the variable value
-- notice that if :var is not set, it fails as well (with a syntax error)
SELECT CASE WHEN 1 = :var THEN error_generator() ELSE TRUE END;

INSERT INTO test_table (integer_value, text_value)
VALUES (:var, 'something');

COMMIT;

*: You can use any shell commands after \! and conditionals of the shell, but since \! opens a new shell, executing anything there does not have any effect for the current psql script.

Answer 3

A more concise version of dezso's answer:

CREATE OR REPLACE FUNCTION pg_temp.err(msg varchar) RETURNS boolean     
AS $$ BEGIN RAISE '%',msg; END; $$ LANGUAGE plpgsql;

You can then call this like:

\set ON_ERROR_STOP on

SELECT CASE WHEN (
  SELECT COUNT(*) FROM mytable
) > 0 THEN pg_temp.err('Already loaded') END;

Answer 4

To quit psql from a script, you can simply use the \quit meta command or generate an error while having ON_ERROR_STOP active. The differences are:

• \quit will do an clean exit with status 0 or terminate just the running script if using nested scripts or \i
• producing an error while having ON_ERROR_STOP turned on will exit with status 3 and propagate to the top-level script in case of nested scripts.

An easier alternative for generating an error then performing RAISE from a function or anonymous code block is just an invalid SQL command like: INVALID;.

To output the error message you can use the \warn meta-command which sends it to stderr (starting pg13) or simply \echo to stdout.

Combined with the \if (starting pg10) my attempt to the problem of the OP would be:

SELECT :'db_to_run_on' = 'production_database' AS forbidden \gset
\if :forbidden
    \warn 'You better not do this!'
    \set ON_ERROR_STOP on
    INVALID COMMAND;
\endif
\unset forbidden

\c :db_to_run_on
TRUNCATE the_most_important_table;

Unfortunately, I'm not aware of a way to suppress the 'syntax error' output generated by the invalid command. So depending on your use-case combining \warn 'You better not do this!' and \quit or \set ON_ERROR_STOP on and DO $$BEGIN RAISE 'You better not do this!'; END$$ LANGUAGE plpgsql; may be more suitable.

Answer 5

What I found works very well for me is to use a scripting language to generate a SQL file which I then pipe into psql, something like this:

#!/usr/bin/env ruby

raise "Not a good database name: #{ARGV.first.inspect}" unless ARGV.first =~ /^(dev|test)/

puts "\\timing off"
puts "set client_min_messages='warning';"
puts
puts "TRUNCATE the_most_important_table;"
puts "-- more commands"

Then, I call this from a driver script:

#!/bin/bash
/usr/bin/ruby generator ${1} | /usr/bin/psql --dbname=${1} --file=- --single-transaction

My driver script is usually a Rake file, but you get the idea.

Answer 6

In a query script (which I guess is somehow different than a psql script), \if was considered a syntax error but code like this worked for me:

DO $$ BEGIN
    IF problematic_value = 'undesired thing' THEN
        RAISE EXCEPTION 'Aborting: % is no good', problematic_value;
    END IF;
END $$;

Tested in PostgreSQL 13.2. The outer DO block is not needed inside a stored procedure.