#!/bin/bash
set -e

# Provision Telegraf state inside the so-postgres container.
# Usage: so-telegraf-postgres <subcommand>
#   create_db    Ensure the so_telegraf database exists.
#   group_role   Provision the so_telegraf group role, the so_maint_telegraf
#                maintenance login role, telegraf/partman schemas, pg_partman,
#                the so_admin maintenance routines, and the hourly pg_cron
#                maintenance job.
#   user         Create or update a per-minion login role granted to so_telegraf.
#                Env: ROLE_USER, ROLE_PASS.
#   retention    Reconcile partman retention and premake on telegraf parents.
#                Env: RETENTION_DAYS.
#   maintenance  Drain default partitions and run partman maintenance, as
#                so_maint_telegraf.
#   check        Report partition health, as so_maint_telegraf. Non-zero if any
#                parent is unhealthy.
#
# A default partition holding rows for a day blocks creating that day's child,
# so maintenance drains defaults before calling partman. Use so-telegraf-repair
# on a grid already stuck in that state.

cmd="${1:?subcommand required}"

# Long enough for an hourly maintenance run to release its locks, short enough
# that soup and admins are not left hanging.
LOCK_TIMEOUT=60s

case "$cmd" in
  create_db)
    if ! docker exec so-postgres psql -U postgres -tAc \
        "SELECT 1 FROM pg_database WHERE datname='so_telegraf'" | grep -q 1; then
      docker exec so-postgres psql -v ON_ERROR_STOP=1 -U postgres \
        -c "CREATE DATABASE so_telegraf"
    fi
    ;;

  group_role)
    docker exec -i so-postgres psql -v ON_ERROR_STOP=1 -v lock_timeout="$LOCK_TIMEOUT" \
      -U postgres -d so_telegraf <<'EOSQL'
SET lock_timeout = :'lock_timeout';
DO $$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'so_telegraf') THEN
        CREATE ROLE so_telegraf NOLOGIN;
    END IF;
END
$$;
DO $$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = 'so_maint_telegraf') THEN
        CREATE ROLE so_maint_telegraf;
    END IF;
END
$$;
ALTER ROLE so_maint_telegraf LOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION NOBYPASSRLS PASSWORD NULL;
GRANT so_telegraf TO so_maint_telegraf;
ALTER ROLE so_maint_telegraf SET role = so_telegraf;
ALTER ROLE so_maint_telegraf SET search_path = pg_catalog;
GRANT CONNECT ON DATABASE so_telegraf TO so_telegraf;
CREATE SCHEMA IF NOT EXISTS telegraf AUTHORIZATION so_telegraf;
GRANT USAGE, CREATE ON SCHEMA telegraf TO so_telegraf;
CREATE SCHEMA IF NOT EXISTS partman;
CREATE EXTENSION IF NOT EXISTS pg_partman SCHEMA partman;
-- Telegraf (running as so_telegraf) calls partman.create_parent()
-- on first write of each metric, which needs USAGE on the partman
-- schema, EXECUTE on its functions/procedures, and write access to
-- partman.part_config so it can register new partitioned parents.
GRANT USAGE, CREATE ON SCHEMA partman TO so_telegraf;
GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA partman TO so_telegraf;
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA partman TO so_telegraf;
GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA partman TO so_telegraf;
-- partman creates per-parent template tables (partman.template_*) at
-- runtime; default privileges extend DML/sequence access to them.
ALTER DEFAULT PRIVILEGES IN SCHEMA partman
    GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO so_telegraf;
ALTER DEFAULT PRIVILEGES IN SCHEMA partman
    GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO so_telegraf;

-- Owned by postgres so no Telegraf role can create objects here; the routines
-- run with the caller's privileges.
CREATE SCHEMA IF NOT EXISTS so_admin AUTHORIZATION postgres;
REVOKE ALL ON SCHEMA so_admin FROM PUBLIC;

-- Resolved through pg_inherits, never by name, so a view or a table that only
-- carries the _default name is never read or written.
CREATE OR REPLACE FUNCTION so_admin.telegraf_defaults()
RETURNS TABLE (
    parent_table  regclass,
    default_part  regclass,
    has_rows      boolean,
    default_size  bigint
)
LANGUAGE plpgsql
STABLE
SECURITY INVOKER
AS $func$
DECLARE
    r  record;
BEGIN
    FOR r IN
        SELECT c.oid::regclass AS pt
        FROM partman.part_config pc
        JOIN pg_catalog.pg_class c ON c.oid = to_regclass(pc.parent_table)
        WHERE pc.parent_table LIKE 'telegraf.%'
          AND c.relkind = 'p'
        ORDER BY pc.parent_table
    LOOP
        parent_table := r.pt;

        SELECT c.oid::regclass INTO default_part
        FROM pg_catalog.pg_inherits i
        JOIN pg_catalog.pg_class c ON c.oid = i.inhrelid
        WHERE i.inhparent = r.pt
          AND c.relkind = 'r' AND c.relispartition
          AND pg_catalog.pg_get_expr(c.relpartbound, c.oid) = 'DEFAULT';

        IF default_part IS NULL THEN
            has_rows     := false;
            default_size := NULL;
        ELSE
            EXECUTE format('SELECT EXISTS (SELECT 1 FROM %s LIMIT 1)', default_part)
                INTO has_rows;
            default_size := pg_catalog.pg_total_relation_size(default_part);
        END IF;

        RETURN NEXT;
    END LOOP;
END;
$func$;

CREATE OR REPLACE PROCEDURE so_admin.telegraf_maintenance()
LANGUAGE plpgsql
SECURITY INVOKER
-- No SET clause: it would forbid the COMMITs below, so so_maint_telegraf pins search_path instead.
AS $proc$
DECLARE
    r  record;
BEGIN
    -- No per-parent EXCEPTION handler: partition_data_proc commits internally,
    -- and COMMIT is illegal while a subtransaction is active. A failing parent
    -- aborts the run and the next pass retries.
    FOR r IN
        SELECT pc.parent_table, pc.retention, d.default_part, d.default_size
        FROM so_admin.telegraf_defaults() d
        JOIN partman.part_config pc ON to_regclass(pc.parent_table) = d.parent_table
        WHERE d.default_part IS NOT NULL AND d.has_rows
        ORDER BY pc.parent_table
    LOOP
        RAISE WARNING 'so_admin.telegraf_maintenance: % holds rows (%), draining',
                      r.default_part, pg_size_pretty(r.default_size);

        -- Cheaper to delete expired rows than to repartition and then drop them.
        IF r.retention IS NOT NULL THEN
            EXECUTE format('DELETE FROM %s WHERE "time" < now() - %L::interval',
                           r.default_part, r.retention);
            COMMIT;
        END IF;

        -- Bounded so a large backlog drains across several runs.
        CALL partman.partition_data_proc(
            p_parent_table := r.parent_table,
            p_loop_count   := 200
        );
        COMMIT;
    END LOOP;

    CALL partman.run_maintenance_proc();
END;
$proc$;

DROP FUNCTION IF EXISTS so_admin.telegraf_partition_status();
CREATE FUNCTION so_admin.telegraf_partition_status()
RETURNS TABLE (
    parent_table      text,
    oldest_child      date,
    newest_child      date,
    days_ahead        int,
    retention         text,
    default_has_rows  boolean,
    default_size      text
)
LANGUAGE plpgsql
STABLE
AS $func$
DECLARE
    r  record;
BEGIN
    FOR r IN
        SELECT d.parent_table AS pt, d.has_rows, d.default_size AS size, pc.retention AS ret
        FROM so_admin.telegraf_defaults() d
        JOIN partman.part_config pc ON to_regclass(pc.parent_table) = d.parent_table
        ORDER BY d.parent_table::text
    LOOP
        parent_table := r.pt::text;
        retention    := r.ret;

        SELECT min(x), max(x) INTO oldest_child, newest_child
        FROM (
            SELECT to_date(substring(c.relname FROM '_p(\d{8})$'), 'YYYYMMDD') AS x
            FROM pg_catalog.pg_inherits i
            JOIN pg_catalog.pg_class c ON c.oid = i.inhrelid
            WHERE i.inhparent = r.pt
              AND pg_catalog.pg_get_expr(c.relpartbound, c.oid) <> 'DEFAULT'
        ) s;

        days_ahead       := newest_child - current_date;
        default_has_rows := r.has_rows;
        default_size     := pg_size_pretty(r.size);

        RETURN NEXT;
    END LOOP;
END;
$func$;

GRANT USAGE ON SCHEMA so_admin TO so_telegraf;
GRANT EXECUTE ON FUNCTION so_admin.telegraf_defaults() TO so_telegraf;
GRANT EXECUTE ON PROCEDURE so_admin.telegraf_maintenance() TO so_telegraf;
GRANT EXECUTE ON FUNCTION so_admin.telegraf_partition_status() TO so_telegraf;

-- Partitions created by a superuser session (an older repair) belong to
-- postgres; partman acting as so_telegraf could not drop them.
DO $$
DECLARE
    r  record;
BEGIN
    FOR r IN
        SELECT n.nspname, c.relname
        FROM pg_catalog.pg_class c
        JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
        WHERE c.relkind IN ('r', 'p')
          AND (n.nspname = 'telegraf'
               OR (n.nspname = 'partman' AND c.relname LIKE 'template\_telegraf\_%'))
          AND c.relowner <> 'so_telegraf'::regrole
    LOOP
        EXECUTE format('ALTER TABLE %I.%I OWNER TO so_telegraf', r.nspname, r.relname);
    END LOOP;
END
$$;

-- Drop the registration older releases left in so_telegraf.
SELECT CASE
         WHEN current_setting('cron.database_name', true) IS DISTINCT FROM current_database()
              AND EXISTS (SELECT 1 FROM pg_catalog.pg_extension WHERE extname = 'pg_cron')
         THEN 'true' ELSE 'false'
       END AS drop_stale_cron \gset
\if :drop_stale_cron
DROP EXTENSION pg_cron CASCADE;
\endif
EOSQL

    stalled=false
    probe_failed=false
    if out=$(docker exec -e PGOPTIONS="-c lock_timeout=$LOCK_TIMEOUT" so-postgres \
        psql -U so_maint_telegraf -d so_telegraf -tAc \
        "SELECT coalesce(bool_or(has_rows), false) FROM so_admin.telegraf_defaults()"); then
      case "$out" in
        t) stalled=true ;;
        f) stalled=false ;;
        *) probe_failed=true ;;
      esac
    else
      probe_failed=true
    fi
    if $probe_failed; then
      echo "ERROR: could not check Telegraf default partitions; the maintenance job was not changed." >&2
      exit 1
    fi

    # Guarded on the live GUC so applying this before the postgresql.conf change
    # has restarted the container skips rather than failing.
    docker exec -i so-postgres psql -v ON_ERROR_STOP=1 -v stalled="$stalled" \
      -U postgres -d postgres <<'EOSQL'
SELECT CASE WHEN current_setting('cron.database_name', true) = current_database()
            THEN 'true' ELSE 'false' END AS cron_here \gset
\if :cron_here
CREATE EXTENSION IF NOT EXISTS pg_cron;
-- A highstate reconciles the job's definition but never its active flag; only so-telegraf-repair turns it on.
BEGIN;
SELECT coalesce(max(jobid)::text, '') AS jobid,
       count(*) > 0 AS job_exists,
       coalesce(bool_and(active), true) AS job_active
FROM cron.job
WHERE jobname = 'telegraf-partman-maintenance' AND username = 'so_maint_telegraf' \gset
SELECT cron.unschedule(jobid) FROM cron.job
WHERE jobname = 'telegraf-partman-maintenance' AND username <> 'so_maint_telegraf';
\if :job_exists
SELECT cron.alter_job(:jobid,
  schedule := '17 * * * *',
  command  := 'CALL so_admin.telegraf_maintenance()',
  database := 'so_telegraf');
\else
-- Registered inactive on a stalled grid so nothing drains the default
-- partitions before so-telegraf-repair truncates them.
SELECT cron.schedule_in_database(
  'telegraf-partman-maintenance',
  '17 * * * *',
  'CALL so_admin.telegraf_maintenance()',
  'so_telegraf',
  'so_maint_telegraf',
  NOT :'stalled'::boolean
);
\endif
UPDATE cron.job SET nodename = current_setting('cron.host')
WHERE jobname = 'telegraf-partman-maintenance'
  AND nodename IS DISTINCT FROM current_setting('cron.host');
COMMIT;
SELECT coalesce(bool_and(active), true) AS job_active
FROM cron.job
WHERE jobname = 'telegraf-partman-maintenance' AND username = 'so_maint_telegraf' \gset
\if :job_active
\else
\echo 'WARNING: the Telegraf maintenance job is inactive; run ''sudo so-telegraf-repair''.'
\endif
\else
\echo 'pg_cron metadata database is not `postgres` yet; skipping job registration.'
\endif
EOSQL
    ;;

  user)
    : "${ROLE_USER:?ROLE_USER is required}"
    : "${ROLE_PASS:?ROLE_PASS is required}"
    # psql does not substitute :vars inside dollar-quoted strings, so the
    # conditional CREATE/ALTER is built outside any DO block and dispatched
    # with \gexec. format() handles identifier/literal quoting.
    docker exec -i so-postgres psql \
      -v ON_ERROR_STOP=1 \
      -v role_user="$ROLE_USER" \
      -v role_pass="$ROLE_PASS" \
      -U postgres -d so_telegraf <<'EOSQL'
-- Keep the password out of postgres.log if this DDL errors.
SET log_min_error_statement = panic;
SELECT format(
  CASE WHEN EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = :'role_user')
       THEN 'ALTER ROLE %I WITH LOGIN PASSWORD %L'
       ELSE 'CREATE ROLE %I WITH LOGIN PASSWORD %L'
  END,
  :'role_user',
  :'role_pass'
) \gexec
GRANT CONNECT ON DATABASE so_telegraf TO :"role_user";
GRANT so_telegraf TO :"role_user";
EOSQL
    ;;

  retention)
    : "${RETENTION_DAYS:?RETENTION_DAYS is required}"
    # \gset + \if guards against a missing pg_partman without using a DO
    # block (psql :var substitution doesn't reach into dollar-quoted code).
    # premake is reconciled here because telegraf.conf only applies it to
    # parents created from now on.
    docker exec -i so-postgres psql \
      -v ON_ERROR_STOP=1 \
      -v retention_days="$RETENTION_DAYS" \
      -U postgres -d so_telegraf <<'EOSQL'
SELECT CASE WHEN EXISTS (SELECT 1 FROM pg_catalog.pg_extension WHERE extname = 'pg_partman')
            THEN 'true' ELSE 'false' END AS has_partman \gset
\if :has_partman
-- infinite_time_partitions so a gap in metrics does not stop partman from
-- premaking forward, which is what leaves everything in the default.
UPDATE partman.part_config
SET retention = :'retention_days' || ' days',
    retention_keep_table = false,
    premake = 7,
    infinite_time_partitions = true
WHERE parent_table LIKE 'telegraf.%'
  AND EXISTS (SELECT 1 FROM pg_catalog.pg_class c
              WHERE c.oid = to_regclass(parent_table) AND c.relkind = 'p');
\endif
EOSQL
    ;;

  maintenance)
    docker exec -i so-postgres psql -v ON_ERROR_STOP=1 -U so_maint_telegraf -d so_telegraf <<'EOSQL'
SELECT CASE WHEN to_regproc('so_admin.telegraf_maintenance') IS NOT NULL
            THEN 'true' ELSE 'false' END AS has_proc \gset
\if :has_proc
CALL so_admin.telegraf_maintenance();
\else
\echo 'so_admin.telegraf_maintenance() is missing; run so-telegraf-postgres group_role first.'
\endif
EOSQL
    ;;

  check)
    docker exec -i so-postgres psql -U so_maint_telegraf -d so_telegraf <<'EOSQL'
\pset border 2
SELECT * FROM so_admin.telegraf_partition_status();
EOSQL
    docker exec -i so-postgres psql -U postgres -d postgres <<'EOSQL'
\pset border 2
SELECT CASE WHEN to_regclass('cron.job_run_details') IS NOT NULL
            THEN 'true' ELSE 'false' END AS has_cron \gset
\if :has_cron
SELECT d.status, d.return_message, d.start_time
FROM cron.job_run_details d
JOIN cron.job j ON j.jobid = d.jobid
WHERE j.jobname = 'telegraf-partman-maintenance'
ORDER BY d.start_time DESC
LIMIT 5;
\else
\echo 'pg_cron is not installed in this database.'
\endif
EOSQL
    unhealthy=$(docker exec so-postgres psql -U so_maint_telegraf -d so_telegraf -tAc \
      "SELECT count(*) FROM so_admin.telegraf_partition_status()
       WHERE default_has_rows OR coalesce(days_ahead, -1) < 1")
    if [ "${unhealthy:-1}" != "0" ]; then
      echo "so-telegraf-postgres check: $unhealthy telegraf parent(s) unhealthy" >&2
      exit 1
    fi
    echo "so-telegraf-postgres check: all telegraf parents healthy"
    ;;

  *)
    echo "Unknown subcommand: $cmd" >&2
    exit 1
    ;;
esac