#!/bin/bash

## Should be executable N time in a row with same result.
##
## Ensure the outline application role owns every object of its
## database before the container boots and runs its migrations.
##
## Historical provisioning or restores executed as the "postgres"
## superuser leave objects owned by "postgres", which makes any
## later ALTER on these objects fail with "must be owner of ..."
## and puts outline in a crash-loop at migration time.  See
## README.org, section "Database ownership alignment".

. lib/common

set -e

relation="postgres-database"

db_role=$(outline:named-relation-get "$relation" user) || {
    err "Couldn't get ${WHITE}user${NORMAL} value" \
        "in ${DARKCYAN}$relation${NORMAL} relation's data."
    exit 1
}

## List every object of schema "public" not owned by the application
## role: tables, sequences, views, materialized views, standalone
## types and functions.  Extensions are excluded (managed by the
## postgres charm).
audit_query="SET app.role = '$db_role';
SELECT obj FROM (
    SELECT CASE c.relkind
             WHEN 'S' THEN 'sequence '
             WHEN 'v' THEN 'view     '
             WHEN 'm' THEN 'matview  '
             ELSE 'table    '
           END || c.relname AS obj
      FROM pg_class c
      JOIN pg_namespace n ON n.oid = c.relnamespace
     WHERE n.nspname = 'public'
       AND c.relkind IN ('r','p','S','v','m')
       AND pg_get_userbyid(c.relowner) <> current_setting('app.role')
    UNION ALL
    SELECT 'type     ' || t.typname AS obj
      FROM pg_type t
      JOIN pg_namespace n ON n.oid = t.typnamespace
     WHERE n.nspname = 'public'
       AND t.typtype IN ('e','d','c')
       AND NOT EXISTS (SELECT 1 FROM pg_class c WHERE c.reltype = t.oid)
       AND pg_get_userbyid(t.typowner) <> current_setting('app.role')
    UNION ALL
    SELECT CASE p.prokind
             WHEN 'p' THEN 'procedure '
             ELSE 'function '
           END || p.proname AS obj
      FROM pg_proc p
      JOIN pg_namespace n ON n.oid = p.pronamespace
     WHERE n.nspname = 'public'
       AND p.prokind IN ('f','p')
       AND pg_get_userbyid(p.proowner) <> current_setting('app.role')
) drift
ORDER BY obj;"

drift=$(sql < <(e "$audit_query")) || {
    err "Failed to audit database ownership for ${WHITE}$db_role${NORMAL}."
    exit 1
}

if [ -z "$drift" ]; then
    ## fast-path: nothing to do
    exit 0
fi

info "Found database objects not owned by '${db_role}', reassigning ownership:"
e "$drift" | prefix "  ${GRAY}|${NORMAL} " >&2

dbname=$(outline:named-relation-get "$relation" dbname) || {
    err "Couldn't get ${WHITE}dbname${NORMAL} value" \
        "in ${DARKCYAN}$relation${NORMAL} relation's data."
    exit 1
}

## Reassign ownership of every drifted object, then the database
## itself.  Role name is passed through a session GUC and quoted
## with format('%I') in every generated statement.
repair_query="SET app.role = '$db_role';
DO \$\$
DECLARE
    r record;
    app_role text := current_setting('app.role');
BEGIN
    FOR r IN
        SELECT CASE c.relkind
                 WHEN 'S' THEN format('ALTER SEQUENCE public.%I OWNER TO %I', c.relname, app_role)
                 WHEN 'v' THEN format('ALTER VIEW public.%I OWNER TO %I', c.relname, app_role)
                 WHEN 'm' THEN format('ALTER MATERIALIZED VIEW public.%I OWNER TO %I', c.relname, app_role)
                 ELSE format('ALTER TABLE public.%I OWNER TO %I', c.relname, app_role)
               END AS stmt
          FROM pg_class c
          JOIN pg_namespace n ON n.oid = c.relnamespace
         WHERE n.nspname = 'public'
           AND c.relkind IN ('r','p','S','v','m')
           AND pg_get_userbyid(c.relowner) <> app_role
        UNION ALL
        SELECT format('ALTER TYPE public.%I OWNER TO %I', t.typname, app_role)
          FROM pg_type t
          JOIN pg_namespace n ON n.oid = t.typnamespace
         WHERE n.nspname = 'public'
           AND t.typtype IN ('e','d','c')
           AND NOT EXISTS (SELECT 1 FROM pg_class c WHERE c.reltype = t.oid)
           AND pg_get_userbyid(t.typowner) <> app_role
        UNION ALL
        SELECT CASE p.prokind
                 WHEN 'p' THEN format('ALTER PROCEDURE public.%I(%s) OWNER TO %I',
                                      p.proname, pg_get_function_identity_arguments(p.oid), app_role)
                 ELSE format('ALTER FUNCTION public.%I(%s) OWNER TO %I',
                                      p.proname, pg_get_function_identity_arguments(p.oid), app_role)
               END AS stmt
          FROM pg_proc p
          JOIN pg_namespace n ON n.oid = p.pronamespace
         WHERE n.nspname = 'public'
           AND p.prokind IN ('f','p')
           AND pg_get_userbyid(p.proowner) <> app_role
    LOOP
        EXECUTE r.stmt;
    END LOOP;
END \$\$;
ALTER DATABASE \"$dbname\" OWNER TO \"$db_role\";"

sql < <(e "$repair_query") || {
    err "Failed to reassign ownership to '${db_role}'."
    exit 1
}

## Fail-hard: verify the realignment actually worked.
remaining_drift=$(sql < <(e "$audit_query")) || {
    err "Failed to re-audit database ownership for ${WHITE}$db_role${NORMAL}."
    exit 1
}

if [ -n "$remaining_drift" ]; then
    err "Some database objects are still not owned by '${db_role}':"
    e "$remaining_drift" | prefix "  ${GRAY}|${NORMAL} " >&2
    err "Deployment halted.  Fix the ownership manually and retry, e.g.:"
    err "  docker exec <postgres-container> psql -U postgres -d \"$dbname\" \\" >&2
    err "      -c 'ALTER TYPE public.<name> OWNER TO \"$db_role\"'" >&2
    exit 1
fi

info "Database ownership aligned on '${db_role}'."