Postgres: pg_cron over the local socket, so_maint_telegraf role, repair and soup 3.5.0

salt/manager/tools/sbin/soup

@@ -546,6 +546,7 @@ postupgrade_changes() {
     [[ "$POSTVERSION" == "3.1.0" ]] && post_to_3.2.0
     [[ "$POSTVERSION" == "3.2.0" ]] && post_to_3.3.0
     [[ "$POSTVERSION" == "3.3.0" ]] && post_to_3.4.0
+    [[ "$POSTVERSION" == "3.4.0" ]] && post_to_3.5.0
     # All applicable post-upgrade steps completed; clear the resume marker.
     rm -f "$POSTVERSION_FILE"
     true
@@ -1065,17 +1066,31 @@ telegraf_repair() {
   # when there is something to repair, so everyone else is left alone.
   local repair=/usr/sbin/so-telegraf-repair
   [[ -x "$repair" ]] || return 0
-  docker ps --format '{{.Names}}' | grep -qx so-postgres || return 0
+  local state
+  state=$(docker inspect -f '{{.State.Running}} {{.State.Restarting}}' so-postgres 2>/dev/null) || return 0
+  if [[ "$state" != "true false" ]]; then
+    echo "  so-postgres is not running; skipping the Telegraf partition check."
+    FINAL_MESSAGE_QUEUE+=("WARNING: so-postgres is not running, so Telegraf partitions were not checked. If /nsm is full, follow \"Telegraf Partition Repair\" at:")
+    FINAL_MESSAGE_QUEUE+=("    - https://securityonion.net/docs/postgresql")
+    return 0
+  fi
 
   echo "Checking Telegraf metric partitions."
-  local status=0
-  "$repair" --check >> "$SOUP_LOG" 2>&1 || status=$?
+  local out status=0
+  out=$("$repair" --check 2>&1) || status=$?
+  echo "$out" >> "$SOUP_LOG"
   case "$status" in
     0) echo "  Telegraf partitions are healthy; nothing to repair." ;;
     1) echo "  Repairing stalled Telegraf partitions."
-       "$repair" --yes \
-         || echo "  warning: so-telegraf-repair failed; run it manually" >&2 ;;
-    *) echo "  Skipping; Telegraf is not storing metrics in Postgres on this host." ;;
+       if ! "$repair" --yes; then
+         echo "  warning: so-telegraf-repair failed; run it manually" >&2
+         FINAL_MESSAGE_QUEUE+=("WARNING: Telegraf partition repair failed; Postgres metrics will keep filling /nsm until it is repaired:")
+         FINAL_MESSAGE_QUEUE+=("    - sudo so-telegraf-repair")
+       fi ;;
+    2) echo "  Skipping: $(tail -n 1 <<<"$out")" ;;
+    *) echo "  warning: could not check Telegraf partitions" >&2
+       FINAL_MESSAGE_QUEUE+=("WARNING: Telegraf partitions could not be checked during the upgrade; run:")
+       FINAL_MESSAGE_QUEUE+=("    - sudo so-telegraf-repair --check") ;;
   esac
 }
 
@@ -1202,6 +1217,15 @@ add_detection_runtime_fields() {
 }
 ### 3.4.0 End ###
 
+### 3.5.0 Scripts ###
+
+post_to_3.5.0() {
+  telegraf_repair
+
+  set_postversion 3.5.0
+}
+### 3.5.0 End ###
+
 
 repo_sync() {
   echo "Sync the local repo."

salt/postgres/defaults.yaml

@@ -17,3 +17,4 @@ postgres:
     log_min_messages: 'warning'
     shared_preload_libraries: pg_cron
     cron.database_name: postgres
+    cron.host: /run/postgresql

salt/postgres/soc_postgres.yaml

@@ -87,3 +87,9 @@ postgres:
       global: True
       advanced: True
       helpLink: postgres
+    cron.host:
+      description: Where pg_cron connects to run scheduled jobs such as Telegraf partition maintenance. Must be the container's PostgreSQL socket directory, /run/postgresql, because pg_cron connects without a password and only local socket connections are trusted. Any other value makes every scheduled job fail. Changing it restarts PostgreSQL.
+      readonly: True
+      global: True
+      advanced: True
+      helpLink: postgres

salt/postgres/tools/sbin/so-stats-show

@@ -52,7 +52,7 @@ if [ -n "$FILTER_HOST" ] && ! [[ "$FILTER_HOST" =~ $HOST_RE ]]; then
 fi
 
 so_psql() {
-  docker exec so-postgres psql -U postgres -d so_telegraf -At -F $'\t' "$@"
+  docker exec -e PGOPTIONS='-c application_name=so-stats-show' so-postgres psql -U so_maint_telegraf -d so_telegraf -At -F $'\t' "$@"
 }
 
 if ! docker exec so-postgres psql -U postgres -lqt 2>/dev/null | cut -d\| -f1 | grep -qw so_telegraf; then
@@ -60,6 +60,11 @@ if ! docker exec so-postgres psql -U postgres -lqt 2>/dev/null | cut -d\| -f1 |
   exit 2
 fi
 
+if ! docker exec so-postgres psql -U postgres -tAc "SELECT 1 FROM pg_roles WHERE rolname = 'so_maint_telegraf'" 2>/dev/null | grep -q 1; then
+  echo "The so_maint_telegraf role is missing; run a highstate."
+  exit 1
+fi
+
 table_exists() {
   local table="$1"
   [ -n "$(so_psql -c "SELECT 1 FROM information_schema.tables WHERE table_schema='${SCHEMA}' AND table_name='${table}' LIMIT 1;")" ]

salt/postgres/tools/sbin/so-telegraf-postgres

@@ -4,15 +4,18 @@ 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, telegraf/partman schemas,
-#                pg_partman, the so_admin maintenance routines, and the hourly
-#                pg_cron maintenance job.
+#   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.
-#   check        Report partition health. Non-zero if any parent is unhealthy.
+#   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
@@ -20,6 +23,10 @@ set -e
 
 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 \
@@ -30,7 +37,9 @@ case "$cmd" in
     ;;
 
   group_role)
-    docker exec -i so-postgres psql -v ON_ERROR_STOP=1 -U postgres -d so_telegraf <<'EOSQL'
+    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
@@ -38,6 +47,17 @@ BEGIN
     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;
@@ -58,44 +78,82 @@ ALTER DEFAULT PRIVILEGES IN SCHEMA partman
 ALTER DEFAULT PRIVILEGES IN SCHEMA partman
     GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO so_telegraf;
 
--- pg_cron runs these as postgres, so they must not sit in a schema any
--- Telegraf role can create objects in.
+-- 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 to_regclass(pc.parent_table) AS pt
+        FROM partman.part_config pc
+        WHERE pc.parent_table LIKE 'telegraf.%'
+        ORDER BY pc.parent_table
+    LOOP
+        CONTINUE WHEN r.pt IS NULL;
+        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;
-    v_default  text;
-    v_rows     bigint;
+    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 parent_table, retention
-        FROM partman.part_config
-        WHERE parent_table LIKE 'telegraf.%'
-        ORDER BY parent_table
+        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
-        v_default := format('%I.%I',
-                            split_part(r.parent_table, '.', 1),
-                            split_part(r.parent_table, '.', 2) || '_default');
-
-        CONTINUE WHEN to_regclass(v_default) IS NULL;
-
-        EXECUTE format('SELECT count(*) FROM %s', v_default) INTO v_rows;
-        CONTINUE WHEN v_rows = 0;
-
-        RAISE WARNING 'so_admin.telegraf_maintenance: % rows stranded in %, draining',
-                      v_rows, v_default;
+        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',
-                           v_default, r.retention);
+                           r.default_part, r.retention);
             COMMIT;
         END IF;
 
@@ -103,7 +161,7 @@ BEGIN
         CALL partman.partition_data_proc(
             p_parent_table := r.parent_table,
             p_loop_count   := 200,
-            p_source_table := v_default
+            p_source_table := r.default_part::text
         );
         COMMIT;
     END LOOP;
@@ -112,58 +170,75 @@ BEGIN
 END;
 $proc$;
 
-CREATE OR REPLACE FUNCTION so_admin.telegraf_partition_status()
+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_rows  bigint,
-    default_size  text
+    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;
-    v_default  regclass;
+    r  record;
 BEGIN
     FOR r IN
-        SELECT pc.parent_table AS pt, pc.retention AS ret
-        FROM partman.part_config pc
-        WHERE pc.parent_table LIKE 'telegraf.%'
-        ORDER BY pc.parent_table
+        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;
+        parent_table := r.pt::text;
         retention    := r.ret;
 
-        SELECT min(d), max(d) INTO oldest_child, newest_child
+        SELECT min(x), max(x) INTO oldest_child, newest_child
         FROM (
-            SELECT to_date(substring(c.relname FROM '_p(\d{8})$'), 'YYYYMMDD') AS d
-            FROM pg_inherits i
-            JOIN pg_class c ON c.oid = i.inhrelid
-            WHERE i.inhparent = r.pt::regclass
-              AND pg_get_expr(c.relpartbound, c.oid) <> 'DEFAULT'
+            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;
-
-        v_default := to_regclass(format('%I.%I',
-                                        split_part(r.pt, '.', 1),
-                                        split_part(r.pt, '.', 2) || '_default'));
-        IF v_default IS NULL THEN
-            default_rows := NULL;
-            default_size := NULL;
-        ELSE
-            EXECUTE format('SELECT count(*) FROM %s', v_default::text) INTO default_rows;
-            default_size := pg_size_pretty(pg_total_relation_size(v_default));
-        END IF;
+        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()
@@ -175,20 +250,69 @@ 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 -U postgres -d postgres <<'EOSQL'
+    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;
--- cron.schedule_in_database is idempotent by jobname.
+-- 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_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
@@ -247,7 +371,7 @@ EOSQL
     ;;
 
   maintenance)
-    docker exec -i so-postgres psql -v ON_ERROR_STOP=1 -U postgres -d so_telegraf <<'EOSQL'
+    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
@@ -259,7 +383,7 @@ EOSQL
     ;;
 
   check)
-    docker exec -i so-postgres psql -U postgres -d so_telegraf <<'EOSQL'
+    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
@@ -278,9 +402,9 @@ LIMIT 5;
 \echo 'pg_cron is not installed in this database.'
 \endif
 EOSQL
-    unhealthy=$(docker exec so-postgres psql -U postgres -d so_telegraf -tAc \
+    unhealthy=$(docker exec so-postgres psql -U so_maint_telegraf -d so_telegraf -tAc \
       "SELECT count(*) FROM so_admin.telegraf_partition_status()
-       WHERE coalesce(default_rows, 0) > 0 OR coalesce(days_ahead, -1) < 1")
+       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

salt/postgres/tools/sbin/so-telegraf-repair

@@ -11,26 +11,38 @@
 # maintenance once. Healthy grids are reported and left alone.
 #
 # Usage: so-telegraf-repair [--check] [--yes] [--no-restart]
-#   --check       Report health and change nothing.
+#   --check       Report health and change nothing. Also reports a missing
+#                 so_maint_telegraf role and a maintenance job that is missing
+#                 or inactive.
 #   --yes         Skip the confirmation prompt (for soup and other automation).
 #   --no-restart  Never restart so-postgres, even if pg_cron's launcher is dead.
 #
+# Repair mode (without --check) registers the so_maint_telegraf role and the
+# maintenance job if they are missing, pauses the job while it repairs, and
+# always leaves it active afterwards. Data work runs as so_maint_telegraf.
+#
 # Exit status:
 #   0  healthy, or repair completed
-#   1  repair is needed (--check only)
-#   2  cannot run here: so-postgres, so_telegraf or pg_partman is missing
+#   1  repair is needed (--check), or the repair failed after it started
+#   2  skipped: Telegraf output is not Postgres, or so-postgres, so_telegraf or pg_partman is missing
+#   3  could not check, for example a connection error
 
-set -e
+set -eE
 
 # Matches p_premake in telegraf.conf's create_parent template.
 PREMAKE=7
 JOB_NAME=telegraf-partman-maintenance
+MAINT_ROLE=so_maint_telegraf
+# Long enough for an hourly maintenance run to release its locks, short enough
+# that soup and admins are not left hanging.
+LOCK_TIMEOUT=60s
+PG_OPTS="-c lock_timeout=$LOCK_TIMEOUT -c application_name=so-telegraf-repair"
 
 CHECK_ONLY=false
 ASSUME_YES=false
 NO_RESTART=false
 
-usage() { sed -n '/^# Usage:/,/^#   2 /p' "$0" | sed 's/^# \?//'; }
+usage() { sed -n '/^# Usage:/,/^#   3 /p' "$0" | sed 's/^# \?//'; }
 
 while [[ $# -gt 0 ]]; do
   case "$1" in
@@ -38,70 +50,153 @@ while [[ $# -gt 0 ]]; do
     --yes|-y)          ASSUME_YES=true ;;
     --no-restart)      NO_RESTART=true ;;
     -h|--help)         usage; exit 0 ;;
-    *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;;
+    *) echo "Unknown option: $1" >&2; usage >&2; exit 3 ;;
   esac
   shift
 done
 
-skip() { echo "$*"; exit 2; }
+# Strips C0 (except tab and newline), DEL and C1 characters, so names and
+# messages written with Telegraf credentials can't drive the terminal.
+sanitise() { LC_ALL=C tr -d '\000-\010\013-\037\177' | LC_ALL=C sed 's/\xc2[\x80-\x9f]//g'; }
 
-psql_tg() { docker exec -i so-postgres psql -U postgres -d so_telegraf "$@"; }
-psql_pg() { docker exec -i so-postgres psql -U postgres -d postgres "$@"; }
+ERR_FILE=$(mktemp)
+ERR_STATUS=3
+ERR_PREFIX="Could not check Telegraf partitions"
+job_id=
+restore_active=
+
+on_err() {
+  # Command substitutions inherit this trap; only the top level reports.
+  [[ $BASH_SUBSHELL -eq 0 ]] || exit "$ERR_STATUS"
+  local line
+  line=$(tail -n 1 "$ERR_FILE" | sanitise)
+  echo "$ERR_PREFIX: ${line:-command failed}" >&2
+  exit "$ERR_STATUS"
+}
+
+on_exit() {
+  if [[ -n "$restore_active" ]]; then
+    psql_cron -tAc "SELECT cron.alter_job($job_id, active := $restore_active)" >/dev/null || true
+  fi
+  rm -f "$ERR_FILE"
+}
+
+trap on_err ERR
+trap on_exit EXIT
+
+finish() { trap - ERR; exit "$1"; }
+skip() { echo "$*"; finish 2; }
+
+# Keeps the last stderr line for on_err while still showing all of it.
+capture() {
+  local rc=0
+  : >"$ERR_FILE"
+  "$@" 2>"$ERR_FILE" || rc=$?
+  sanitise <"$ERR_FILE" >&2
+  return "$rc"
+}
+
+psql_tg() {
+  capture docker exec -i -e PGOPTIONS="$PG_OPTS" so-postgres psql -U "$MAINT_ROLE" -d so_telegraf "$@"
+}
+psql_pg() {
+  capture docker exec -i -e PGOPTIONS="$PG_OPTS" so-postgres psql -U postgres -d postgres "$@"
+}
+psql_cron() {
+  capture docker exec -i -e PGOPTIONS="$PG_OPTS" so-postgres psql -U postgres -d "$cron_db" "$@"
+}
 
-# query_to_xml so the per-table row counts need no helper function installed.
 REPORT="
-WITH parents AS (
-    SELECT pc.parent_table,
-           pc.premake,
-           split_part(pc.parent_table, '.', 1) AS sch,
-           split_part(pc.parent_table, '.', 2) AS tbl
-    FROM partman.part_config pc
-    WHERE pc.parent_table LIKE 'telegraf.%'
-), children AS (
-    SELECT p.parent_table,
-           max(to_date(substring(c.relname FROM '_p(\d{8})\$'), 'YYYYMMDD')) AS newest_child
-    FROM parents p
-    JOIN pg_class pt ON pt.oid = p.parent_table::regclass
-    JOIN pg_inherits i ON i.inhparent = pt.oid
-    JOIN pg_class c ON c.oid = i.inhrelid
-    WHERE pg_get_expr(c.relpartbound, c.oid) <> 'DEFAULT'
-    GROUP BY p.parent_table
-), defaults AS (
-    SELECT p.parent_table,
-           p.premake,
-           format('%I.%I', p.sch, p.tbl || '_default') AS default_table,
-           to_regclass(format('%I.%I', p.sch, p.tbl || '_default')) AS default_oid
-    FROM parents p
-)
 SELECT d.parent_table,
        c.newest_child,
        (c.newest_child - current_date) AS days_ahead,
-       d.premake,
-       CASE WHEN d.default_oid IS NULL THEN NULL ELSE
-            (xpath('/row/cnt/text()',
-                   query_to_xml(format('SELECT count(*) AS cnt FROM %s', d.default_table),
-                                false, true, '')))[1]::text::bigint
-       END AS default_rows,
-       CASE WHEN d.default_oid IS NULL THEN NULL
-            ELSE pg_size_pretty(pg_total_relation_size(d.default_oid)) END AS default_size
-FROM defaults d
-LEFT JOIN children c ON c.parent_table = d.parent_table
-ORDER BY 1
+       pc.premake,
+       d.has_rows AS default_has_rows,
+       pg_size_pretty(d.default_size) AS default_size
+FROM so_admin.telegraf_defaults() d
+JOIN partman.part_config pc ON to_regclass(pc.parent_table) = d.parent_table
+LEFT JOIN LATERAL (
+    SELECT max(to_date(substring(ch.relname FROM '_p(\d{8})\$'), 'YYYYMMDD')) AS newest_child
+    FROM pg_inherits i
+    JOIN pg_class ch ON ch.oid = i.inhrelid
+    WHERE i.inhparent = d.parent_table
+      AND pg_get_expr(ch.relpartbound, ch.oid) <> 'DEFAULT'
+) c ON true
+ORDER BY d.parent_table::text
 "
 
-docker ps --format '{{.Names}}' | grep -qx so-postgres \
+print_report() {
+  local report
+  report=$(psql_tg -c "$REPORT")
+  sanitise <<<"$report"
+}
+
+has_cron() {
+  [[ -n "$cron_db" ]] || return 0
+  psql_cron -tAc "SELECT 1 FROM pg_extension WHERE extname = 'pg_cron'"
+}
+
+lookup_job() {
+  local cron
+  cron=$(has_cron)
+  [[ -n "$cron" ]] || return 0
+  psql_cron -tA -F '|' -c \
+    "SELECT jobid, active FROM cron.job
+      WHERE jobname = '$JOB_NAME' AND username = '$MAINT_ROLE'
+      ORDER BY jobid LIMIT 1"
+}
+
+output=$(capture salt-call pillar.get telegraf:output --out=newline_values_only)
+output=$(tr -d '[:space:]' <<<"$output" | tr '[:lower:]' '[:upper:]')
+output=${output:-INFLUXDB}
+case "$output" in
+  POSTGRES|BOTH) ;;
+  *) skip "Telegraf output is $(sanitise <<<"$output"); not checking Postgres partitions." ;;
+esac
+
+containers=$(capture docker ps --format '{{.Names}}')
+grep -qx so-postgres <<<"$containers" \
   || skip "so-postgres is not running; nothing to repair."
-docker exec so-postgres psql -U postgres -tAc \
-  "SELECT 1 FROM pg_database WHERE datname='so_telegraf'" | grep -q 1 \
+has_db=$(psql_pg -tAc "SELECT 1 FROM pg_database WHERE datname = 'so_telegraf'")
+[[ -n "$has_db" ]] \
   || skip "The so_telegraf database does not exist; Telegraf is not writing to Postgres."
-psql_tg -tAc "SELECT 1 FROM pg_extension WHERE extname='pg_partman'" | grep -q 1 \
+
+# so_telegraf before the postgres state lands, postgres after.
+cron_db=$(psql_pg -tAc "SELECT current_setting('cron.database_name', true)")
+cron_db=$(tr -d '[:space:]' <<<"$cron_db")
+
+has_role=$(psql_pg -tAc "SELECT 1 FROM pg_roles WHERE rolname = '$MAINT_ROLE'")
+job=
+[[ "$cron_db" == postgres ]] && job=$(lookup_job)
+
+if [[ -z "$has_role" ]] && $CHECK_ONLY; then
+  echo "Repair is needed:"
+  echo "  * maintenance role $MAINT_ROLE is missing"
+  echo
+  echo "Re-run without --check to repair."
+  finish 1
+fi
+
+if ! $CHECK_ONLY && [[ -z "$has_role" || ( "$cron_db" == postgres && -z "$job" ) ]]; then
+  echo "Registering the maintenance role and job."
+  if ! /usr/sbin/so-telegraf-postgres group_role; then
+    echo "Could not register the maintenance job; run a highstate and try again." >&2
+    finish 3
+  fi
+  [[ "$cron_db" == postgres ]] && job=$(lookup_job)
+fi
+
+has_partman=$(psql_tg -tAc "SELECT 1 FROM pg_extension WHERE extname = 'pg_partman'")
+[[ -n "$has_partman" ]] \
   || skip "pg_partman is not installed in so_telegraf; nothing to repair."
 
 parents=$(psql_tg -tAc \
   "SELECT count(*) FROM partman.part_config WHERE parent_table LIKE 'telegraf.%'")
-stranded=$(psql_tg -tAc "SELECT coalesce(sum(default_rows), 0) FROM ( $REPORT ) t")
-behind=$(psql_tg -tAc \
-  "SELECT count(*) FROM ( $REPORT ) t WHERE coalesce(days_ahead, -1) < 1")
+counts=$(psql_tg -tAc \
+  "SELECT count(*) FILTER (WHERE default_has_rows) || '|' ||
+          count(*) FILTER (WHERE coalesce(days_ahead, -1) < 1)
+     FROM ( $REPORT ) t")
+IFS='|' read -r stranded behind <<<"$counts"
 # premake < 7, or infinite_time_partitions off: without the latter partman
 # refuses to premake forward across the gap the stall left behind.
 misconfigured=$(psql_tg -tAc \
@@ -115,16 +210,24 @@ launcher=$(psql_pg -tAc \
   "SELECT count(*) FROM pg_stat_activity
     WHERE backend_type ILIKE '%pg_cron%' OR application_name ILIKE '%pg_cron%'")
 
-# so_telegraf before the postgres state lands, postgres after.
-cron_db=$(docker exec so-postgres psql -U postgres -tAc \
-  "SELECT current_setting('cron.database_name', true)" | tr -d '[:space:]')
 last_run=never
-if [[ -n "$cron_db" ]]; then
-  last_run=$(docker exec so-postgres psql -U postgres -d "$cron_db" -tAc \
-    "SELECT coalesce(max(d.start_time)::text, 'never')
+cron_installed=$(has_cron)
+if [[ -n "$cron_installed" ]]; then
+  run=$(psql_cron -tAc \
+    "SELECT d.status || ' at ' || d.start_time
        FROM cron.job_run_details d JOIN cron.job j USING (jobid)
-      WHERE j.jobname = '$JOB_NAME'" 2>/dev/null | tr -d '[:space:]' || echo unknown)
-  [[ -n "$last_run" ]] || last_run=never
+      WHERE j.jobname = '$JOB_NAME'
+      ORDER BY d.start_time DESC LIMIT 1")
+  if [[ -n "$run" ]]; then
+    message=$(psql_cron -tAc \
+      "SELECT coalesce(d.return_message, '')
+         FROM cron.job_run_details d JOIN cron.job j USING (jobid)
+        WHERE j.jobname = '$JOB_NAME'
+        ORDER BY d.start_time DESC LIMIT 1")
+    message=$(sanitise <<<"$message")
+    message=${message//$'\n'/ }
+    last_run="$(sanitise <<<"$run"): ${message:0:200}"
+  fi
 fi
 
 # A grid that has never written a metric has nothing to recover, and an empty
@@ -132,43 +235,67 @@ fi
 restart_needed=false
 [[ "$launcher" -eq 0 && "$parents" -gt 0 && -n "$cron_db" ]] && restart_needed=true
 
+# group_role registers the job only in the postgres database.
+job_inactive=false
+if [[ "$cron_db" == postgres ]]; then
+  [[ -z "$job" || "${job##*|}" == f ]] && job_inactive=true
+fi
+
 repair_needed=false
 [[ "$stranded" -gt 0 ]] && repair_needed=true
 [[ "$behind" -gt 0 ]] && repair_needed=true
 [[ "$misconfigured" -gt 0 ]] && repair_needed=true
 $restart_needed && repair_needed=true
+$job_inactive && repair_needed=true
 
 echo "Telegraf partition status:"
-psql_tg -c "$REPORT"
-echo "Rows stranded in default partitions: $stranded"
-echo "pg_cron metadata database: ${cron_db:-unset}"
+print_report
+echo "Default partitions with rows: $stranded"
+echo "pg_cron metadata database: $(sanitise <<<"${cron_db:-unset}")"
+[[ "$cron_db" == postgres ]] \
+  || echo "pg_cron metadata database is $(sanitise <<<"${cron_db:-unset}"); not checking the maintenance job."
 echo "pg_cron launcher running: $([[ "$launcher" -gt 0 ]] && echo yes || echo no)"
 echo "Last $JOB_NAME run: $last_run"
 echo
 
 if ! $repair_needed; then
   echo "Telegraf partitions are healthy. Nothing to do."
-  exit 0
+  finish 0
 fi
 
 if $CHECK_ONLY; then
   echo "Repair is needed:"
-  [[ "$stranded" -gt 0 ]]    && echo "  * $stranded row(s) stranded in default partitions"
+  [[ "$stranded" -gt 0 ]]    && echo "  * $stranded default partition(s) hold rows"
   [[ "$behind" -gt 0 ]]      && echo "  * $behind parent(s) with no partition for the current window"
   [[ "$misconfigured" -gt 0 ]] && echo "  * $misconfigured parent(s) with stale partman settings"
   $restart_needed            && echo "  * pg_cron's launcher is dead; maintenance is not running at all"
+  $job_inactive              && echo "  * the maintenance job is missing or inactive"
   echo
   echo "Re-run without --check to repair."
-  exit 1
+  finish 1
 fi
 
 if [[ "$stranded" -gt 0 ]] && ! $ASSUME_YES; then
-  echo "This will permanently discard the $stranded stranded row(s) above."
+  echo "This will permanently discard the rows in the $stranded default partition(s) above."
   $restart_needed && ! $NO_RESTART && \
     echo "so-postgres will also be restarted, which briefly interrupts SOC."
-  [[ -t 0 ]] || { echo "Not a terminal; re-run with --yes to confirm." >&2; exit 2; }
+  [[ -t 0 ]] || { echo "Not a terminal; re-run with --yes to confirm." >&2; finish 1; }
   read -r -p "Continue? [y/N] " answer
-  [[ "$answer" =~ ^[Yy]$ ]] || { echo "Aborted."; exit 0; }
+  [[ "$answer" =~ ^[Yy]$ ]] || { echo "Aborted."; finish 0; }
+fi
+
+ERR_STATUS=1
+ERR_PREFIX="Telegraf partition repair failed"
+
+if [[ "$cron_db" == postgres ]]; then
+  if [[ -z "$job" ]]; then
+    echo "Could not register the maintenance job; run a highstate and try again." >&2
+    finish 3
+  fi
+  job_id=${job%%|*}
+  [[ "${job##*|}" == t ]] && restore_active=true || restore_active=false
+  echo "Pausing the $JOB_NAME job during the repair."
+  psql_cron -tAc "SELECT cron.alter_job($job_id, active := false)" >/dev/null
 fi
 
 if [[ "$misconfigured" -gt 0 ]]; then
@@ -186,27 +313,24 @@ if [[ "$stranded" -gt 0 ]]; then
   # One transaction: Telegraf is still writing, so a default emptied without
   # its partition in place is refilled before maintenance can attach one.
   psql_tg -v ON_ERROR_STOP=1 <<'EOSQL'
+BEGIN;
 DO $$
 DECLARE
     r record;
 BEGIN
     FOR r IN
-        SELECT pc.parent_table,
-               format('%I.%I', n.nspname, c.relname) AS default_table
-        FROM partman.part_config pc
-        JOIN pg_class p ON p.oid = pc.parent_table::regclass
-        JOIN pg_inherits i ON i.inhparent = p.oid
-        JOIN pg_class c ON c.oid = i.inhrelid
-        JOIN pg_namespace n ON n.oid = c.relnamespace
-        WHERE pc.parent_table LIKE 'telegraf.%'
-          AND pg_get_expr(c.relpartbound, c.oid) = 'DEFAULT'
+        SELECT pc.parent_table, d.default_part
+        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
     LOOP
-        EXECUTE format('TRUNCATE TABLE %s', r.default_table);
+        EXECUTE format('TRUNCATE TABLE %s', r.default_part);
         PERFORM partman.create_partition_time(
             r.parent_table, ARRAY[date_trunc('day', now())]::timestamptz[]);
     END LOOP;
 END
 $$;
+COMMIT;
 EOSQL
 fi
 
@@ -222,7 +346,7 @@ if $restart_needed; then
       sleep 2
     done
     docker exec so-postgres pg_isready -U postgres -q \
-      || { echo "so-postgres did not come back; check 'docker logs so-postgres'." >&2; exit 1; }
+      || { echo "so-postgres did not come back; check 'docker logs so-postgres'." >&2; finish 1; }
   fi
 fi
 
@@ -238,8 +362,14 @@ CALL partman.run_maintenance_proc();
 \endif
 EOSQL
 
+if [[ -n "$job_id" ]]; then
+  psql_cron -tAc "SELECT cron.alter_job($job_id, active := true)" >/dev/null
+  restore_active=
+  echo "The $JOB_NAME job is active."
+fi
+
 echo
 echo "Telegraf partition status after repair:"
-psql_tg -c "$REPORT"
+print_report
 echo "The $JOB_NAME job runs hourly at :17. Confirm it fired with:"
 echo "  so-telegraf-repair --check"