#!/bin/bash
# =============================================================================
# Load master (seed) data from the central repository.
#
# Shipped in muli-neo-database, installed to /usr/local/muli/bin.
#
#   muli-db-seed                        apply to every database on this host
#   muli-db-seed <database>             apply to one database, port derived
#   muli-db-seed <database> <port>      apply to one database
#   muli-db-seed --status               report only, change nothing
#   muli-db-seed --repair [<database>]  re-apply even if already at this revision
#   muli-db-seed --repair=<table> <db>  re-apply one table and its FK group
#
# --repair exists for the login-time self-heal in startm: a lookup fails, the
# row is restored, the lookup is retried. A normal run would do nothing there,
# because the container is already recorded at this revision - the only thing
# --repair changes is that muli_seed_state stops being a reason to skip.
#
# It relaxes no safety rule. Restoring a deleted row needs none relaxed: the
# delta already inserts any key the base is missing, unconditionally. What it
# will NOT do is overwrite a row that exists but has been edited locally to
# something wrong - that is a deliberate local edit as far as seeding can tell,
# and clause 2 of the delta leaves it alone. Repair restores what is missing;
# it does not overrule the site.
#
# The data itself is NOT in this package. It is pulled from the seed repository
# (see muli-db-seed-pull) into $MULI_SEED_ROOT, so a tax-table change is a
# commit there, not a release here.
#
# Three rules govern every load, and the third is the important one:
#
#   * `_u` tables are never written. They are the customer's.
#   * `_m` tables may be replaced wholesale. They are Muli's.
#   * provisioning tables use mode `create-only`: the row is established if it
#     is absent and never touched again. These are tables the site owns once
#     created - `orgs` holds a contract's own organisations alongside the dozen
#     Muli provisions, so a full-sync would delete every customer record in it.
#   * base tables are NEVER rebuilt - only a delta derived from the change to
#     `_m` is applied. A live `allocs` carries rows in neither `_m` nor `_u`,
#     inserted at runtime by rtbc_fns.inc as "Missing Reinserted"/seq 998 when
#     the trial balance meets a missing allocation code. Those rows are
#     referenced by financial data. Rebuilding the base table would delete them.
# =============================================================================
set -euo pipefail

SEED_ROOT="${MULI_SEED_ROOT:-/usr/local/muli/dbmaster}"
PG_VERSION="${MULI_PG_VERSION:-14}"
PG_CLUSTER_ROOT="${MULI_PG_CLUSTER_ROOT:-/usr/local/mulipg/${PG_VERSION}}"
SUPERUSER="${MULI_DB_SUPERUSER:-postgres}"
REPAIR=false
REPAIR_TABLE=""
STATUS_ONLY=false

log()  { echo "[muli-db-seed] $*"; }
warn() { echo "[muli-db-seed] WARNING: $*" >&2; }
die()  { echo "[muli-db-seed] ERROR: $*" >&2; exit 1; }

while [ $# -gt 0 ]; do
    case "$1" in
        --status)    STATUS_ONLY=true; shift ;;
        --repair)    REPAIR=true; shift ;;
        --repair=*)  REPAIR=true; REPAIR_TABLE="${1#--repair=}"; shift ;;
        -h|--help)   sed -n '3,40p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
        --*)         die "unknown option: $1" ;;
        *)           break ;;
    esac
done

MANIFEST="${SEED_ROOT}/manifest.tsv"
[ -f "$MANIFEST" ] || die "no manifest at ${MANIFEST}; has the seed repository been pulled?"

# Revision of the data being applied. Recorded per table so a container that has
# stopped pulling is visible.
REVISION=$(git -C "$SEED_ROOT" rev-parse --short HEAD 2>/dev/null || echo "unversioned")

tier_of() {
    case "$1" in
        login_db) echo login ;;
        cdb*)     echo common ;;
        [0-9][0-9][0-9][0-9][0-9][A-I]) echo contract ;;
        *)        echo "" ;;
    esac
}
port_of() {
    local db="$1" c l i
    case "$(tier_of "$db")" in login) echo 59999; return ;; common) echo 60000; return ;; esac
    c="${db:0:5}"; l="${db:5:1}"
    i=$(awk -v x="$l" 'BEGIN{print index("ABCDEFGHI", toupper(x))}')
    [ "$i" -gt 0 ] || return 1
    echo $(( 50000 + 10#"${c:0:1}${c:2:3}${i}" ))
}
discover() {
    [ -d "$PG_CLUSTER_ROOT" ] || return 0
    local d n
    for d in "${PG_CLUSTER_ROOT}"/*pgsql; do
        [ -d "$d" ] || continue
        n=$(basename "$d"); n="${n%pgsql}"
        [ -n "$(tier_of "$n")" ] && echo "$n"
    done
}

q() { psql -tAq -v ON_ERROR_STOP=1 -U "$SUPERUSER" -p "$2" -d "$1" -c "$3"; }

# Column list of a table, comma separated, in ordinal order.
columns_of() { q "$1" "$2" "SELECT string_agg(quote_ident(attname), ',' ORDER BY attnum)
                           FROM pg_attribute WHERE attrelid='$3'::regclass
                            AND attnum>0 AND NOT attisdropped"; }

ensure_ledger() {
    q "$1" "$2" "SET client_min_messages=warning;
                 CREATE TABLE IF NOT EXISTS muli_seed_state (
                     table_name text PRIMARY KEY,
                     revision   text NOT NULL,
                     row_count  integer,
                     applied_at timestamptz NOT NULL DEFAULT now());
                 GRANT SELECT ON muli_seed_state TO muli;" >/dev/null
}

schema_at_least() { # schema_at_least <db> <port> <version>
    [ "$3" = "-" ] && return 0
    local have
    have=$(q "$1" "$2" "SELECT count(*) FROM schema_migrations WHERE version >= '$3'" 2>/dev/null || echo 0)
    [ "${have:-0}" -gt 0 ]
}

requirements_met() { # requirements_met <db> <port> <comma-separated tables>
    [ "$3" = "-" ] || [ -z "$3" ] && return 0
    local spec t n rest col val
    for spec in $(echo "$3" | tr ',' ' '); do
        case "$spec" in
            # `table:column=value` - the specific parent row must exist, not
            # merely some row. `digital` was excluded from the feed after
            # `requires=project` passed on every contract (they all have
            # projects) and the load then failed on digital_proj_fkey, because
            # the projects it needed belonged to the contract it was exported
            # from. A non-empty parent table is not the same as the parent a
            # child row actually references.
            *:*=*)
                t="${spec%%:*}"; rest="${spec#*:}"; col="${rest%%=*}"; val="${rest#*=}"
                n=$(q "$1" "$2" "SELECT count(*) FROM ${t} WHERE trim(${col}) = '${val}'" 2>/dev/null || echo 0)
                ;;
            *)
                t="$spec"
                n=$(q "$1" "$2" "SELECT count(*) FROM ${t}" 2>/dev/null || echo 0)
                ;;
        esac
        [ "${n:-0}" -gt 0 ] || return 1
    done
    return 0
}

# --- foreign key grouping -----------------------------------------------------
# Tables joined by a foreign key cannot be loaded independently. A full-sync
# DELETE on the parent is refused while children survive, and an INSERT on the
# child needs its parent already there. Loading each table in its own
# transaction has no way to order that, which is how opendoc_tpl and
# opendocrev_tpl failed on a contract whose parent table had drifted.
#
# The edges are read from the live catalogue rather than declared in the
# manifest, for the same reason the schema guard derives its tier instead of
# hardcoding it: a constraint added later then orders itself, instead of
# silently breaking the load until somebody notices.
sql_list() { # sql_list <space separated>  -> 'a','b','c'
    local out="" t
    for t in $1; do out="${out}${out:+,}'${t}'"; done
    [ -n "$out" ] || out="''"
    echo "$out"
}

fk_edges() { # fk_edges <db> <port> <tables...>  -> lines of "child parent"
    local list; list=$(sql_list "$3")
    q "$1" "$2" "
      SELECT DISTINCT c.conrelid::regclass::text || ' ' || c.confrelid::regclass::text
        FROM pg_constraint c
       WHERE c.contype = 'f'
         AND c.conrelid <> c.confrelid
         AND c.conrelid::regclass::text  IN (${list})
         AND c.confrelid::regclass::text IN (${list})" 2>/dev/null || true
}

# Kahn: emit a table once every table it references has been emitted, so the
# result is parents first. Any table caught in a cycle is emitted at the end -
# it still loads, it just loses the ordering guarantee, which beats refusing to
# load it at all.
topo_order() { # topo_order <tables> <edges "child parent">
    local tables="$1" edges="$2"
    # Declared separately: under `set -u` a later assignment in the same
    # `local` cannot read an earlier one from that same statement.
    local remaining="$tables" out="" progress t p ok
    while [ -n "$remaining" ]; do
        progress=""
        for t in $remaining; do
            ok=true
            while read -r c p; do
                [ "$c" = "$t" ] || continue
                case " $out " in *" $p "*) ;; *) case " $remaining " in *" $p "*) ok=false ;; esac ;; esac
            done <<< "$edges"
            if [ "$ok" = true ]; then out="${out}${out:+ }${t}"; progress=yes; fi
        done
        if [ -z "$progress" ]; then
            warn "foreign key cycle among: ${remaining}; loading them unordered"
            out="${out}${out:+ }${remaining}"; break
        fi
        local next=""
        for t in $remaining; do case " $out " in *" $t "*) ;; *) next="${next}${next:+ }${t}" ;; esac; done
        remaining="$next"
    done
    echo "$out"
}

# Connected components over the undirected edge set: everything that must share
# a transaction.
group_of() { # group_of <table> <edges>  -> space separated members
    local seed="$1" edges="$2"
    local grp="$seed" added=yes c p
    while [ "$added" = yes ]; do
        added=no
        while read -r c p; do
            [ -n "$c" ] || continue
            case " $grp " in
                *" $c "*) case " $grp " in *" $p "*) ;; *) grp="$grp $p"; added=yes ;; esac ;;
            esac
            case " $grp " in
                *" $p "*) case " $grp " in *" $c "*) ;; *) grp="$grp $c"; added=yes ;; esac ;;
            esac
        done <<< "$edges"
    done
    echo "$grp"
}

# Populated per database in pass 1 and read by apply_group; global because bash
# cannot pass an associative array to a function.
declare -A SEED_SPEC

seed_one_db() { # seed_one_db <db> <port> <tier>
    local db="$1" port="$2" tier="$3" applied=0 skipped=0 failed=0
    ensure_ledger "$db" "$port"

    # Pass 1 - work out what this database can take. Nothing is written yet:
    # the foreign key grouping needs the whole eligible set before the first
    # table is touched.
    SEED_SPEC=()
    local -A due=()
    local eligible="" file mtier table keycols mode base minschema requires current
    while IFS=$'\t' read -r file mtier table keycols mode base minschema requires; do
        case "$file" in ''|\#*) continue ;; esac
        [ "$mtier" = "$tier" ] || continue
        [ -f "${SEED_ROOT}/${file}" ] || { warn "${file} listed but not present"; continue; }

        if ! q "$db" "$port" "SELECT to_regclass('${table}') IS NOT NULL" | grep -qx t; then
            warn "${db}: table ${table} does not exist, skipping ${file}"; skipped=$((skipped+1)); continue
        fi
        if ! schema_at_least "$db" "$port" "$minschema"; then
            warn "${db}: ${table} needs schema >= ${minschema}, not yet applied; skipping"
            skipped=$((skipped+1)); continue
        fi
        # Some tables reference per-contract provisioning records rather than
        # other seed data - opendoc_tpl has a foreign key to project,
        # rpctemplate_m to rpcs. On an unprovisioned database those loads fail
        # with a foreign key violation, which is noise rather than a fault.
        if ! requirements_met "$db" "$port" "${requires:--}"; then
            warn "${db}: ${table} needs provisioning data (${requires}); skipping"
            skipped=$((skipped+1)); continue
        fi

        SEED_SPEC[$table]="${file}"$'\t'"${keycols}"$'\t'"${mode}"$'\t'"${base}"
        eligible="${eligible}${eligible:+ }${table}"
        current=$(q "$db" "$port" "SELECT revision FROM muli_seed_state WHERE table_name='${table}'" || true)
        if [ "$REPAIR" = true ] && { [ -z "$REPAIR_TABLE" ] || [ "$REPAIR_TABLE" = "$table" ]; }; then
            # Repair ignores the ledger. Naming one table still pulls in its
            # foreign key group below, because a group cannot be loaded in part.
            due[$table]=1
        elif [ "$current" != "$REVISION" ]; then
            due[$table]=1
        fi
    done < "$MANIFEST"

    if [ -z "$eligible" ]; then
        log "${db} (${tier}): 0 applied, ${skipped} skipped"
        return 0
    fi

    local edges; edges=$(fk_edges "$db" "$port" "$eligible")

    # Pass 2 - apply, one foreign key group at a time.
    local seen="" t m grp members ordered any_due n
    for t in $eligible; do
        case " $seen " in *" $t "*) continue ;; esac
        grp=$(group_of "$t" "$edges")
        members=""
        for m in $grp; do
            case " $eligible " in *" $m "*) members="${members}${members:+ }${m}" ;; esac
        done
        ordered=$(topo_order "$members" "$edges")
        seen="${seen} ${ordered}"

        # A group moves together. If any member is behind, every member is
        # reloaded - re-applying a current table is an upsert to the data it
        # already holds, and loading only part of a group is exactly what
        # cannot be ordered safely.
        any_due=no
        for m in $ordered; do [ -n "${due[$m]:-}" ] && any_due=yes; done
        n=$(echo "$ordered" | wc -w | tr -d ' ')

        if [ "$any_due" = no ]; then
            if [ "$STATUS_ONLY" = true ]; then
                for m in $ordered; do log "  ${m}: up to date (${REVISION})"; done
            fi
            continue
        fi
        if [ "$STATUS_ONLY" = true ]; then
            for m in $ordered; do
                log "  ${m}: would apply ${REVISION} (currently ${due[$m]:+behind}${due[$m]:-current})"
                applied=$((applied+1))
            done
            continue
        fi

        if [ "$n" -eq 1 ]; then
            IFS=$'\t' read -r file keycols mode base <<< "${SEED_SPEC[$ordered]}"
            if apply_table "$db" "$port" "$file" "$ordered" "$keycols" "$mode" "$base"; then
                applied=$((applied+1))
            else
                warn "${db}: ${ordered} failed; continuing with the remaining tables"
                failed=$((failed+1))
            fi
        else
            log "  foreign key group: ${ordered}"
            if apply_group "$db" "$port" "$ordered"; then
                applied=$((applied+n))
            else
                warn "${db}: group [${ordered}] failed; continuing with the remaining tables"
                failed=$((failed+n))
            fi
        fi
    done

    if [ "$failed" -gt 0 ]; then
        log "${db} (${tier}): ${applied} applied, ${skipped} skipped, ${failed} FAILED"
        return 1
    fi
    log "${db} (${tier}): ${applied} applied, ${skipped} skipped"
}

# Load a set of foreign-key-related tables in one transaction, so the ordering
# a single-table load cannot express becomes possible:
#
#   stage every table, then DELETE children before parents, then INSERT parents
#   before children.
#
# Same guarantees as apply_table otherwise - one transaction, `_u` never
# written, base tables converged by the delta rather than rebuilt.
apply_group() { # apply_group <db> <port> <ordered parents-first tables>
    local db="$1" port="$2" ordered="$3"
    local sql="BEGIN;" i=0 t file keycols mode base cols keyjoin setlist idx
    local -A IDX=()

    for t in $ordered; do
        i=$((i+1)); IDX[$t]=$i
        IFS=$'\t' read -r file keycols mode base <<< "${SEED_SPEC[$t]}"
        cols=$(columns_of "$db" "$port" "$t")
        sql="${sql}
CREATE TEMP TABLE _seed_stage_${i} (LIKE ${t} INCLUDING DEFAULTS) ON COMMIT DROP;
\\copy _seed_stage_${i} (${cols}) FROM '${SEED_ROOT}/${file}' WITH (FORMAT csv, HEADER true)
CREATE INDEX ON _seed_stage_${i} (${keycols});
ANALYZE _seed_stage_${i};"
        [ "$base" != "-" ] && sql="${sql}
CREATE TEMP TABLE _seed_old_${i} AS SELECT * FROM ${t};"
        log "  ${t} <- ${file} (${mode})"
    done

    # Deletes run children first: reverse of the parents-first order.
    local reversed=""
    for t in $ordered; do reversed="${t}${reversed:+ }${reversed}"; done
    for t in $reversed; do
        IFS=$'\t' read -r file keycols mode base <<< "${SEED_SPEC[$t]}"
        [ "$mode" = "full-sync" ] || continue
        idx=${IDX[$t]}
        keyjoin=$(echo "$keycols" | tr ',' '\n' | awk '{printf "%s a.%s = b.%s", (NR>1?" AND":""), $1, $1}')
        sql="${sql}
DELETE FROM ${t} a WHERE NOT EXISTS (SELECT 1 FROM _seed_stage_${idx} b WHERE ${keyjoin});"
    done

    # Inserts, then each table's delta, parents first.
    for t in $ordered; do
        IFS=$'\t' read -r file keycols mode base <<< "${SEED_SPEC[$t]}"
        idx=${IDX[$t]}
        cols=$(columns_of "$db" "$port" "$t")
        setlist=$(echo "$cols" | tr ',' '\n' | grep -vxF -f <(echo "$keycols" | tr ',' '\n') \
                  | awk '{printf "%s%s=EXCLUDED.%s", (NR>1?", ":""), $1, $1}')
        if [ "$mode" = "create-only" ]; then
            sql="${sql}
INSERT INTO ${t} (${cols}) SELECT ${cols} FROM _seed_stage_${idx}
  ON CONFLICT (${keycols}) DO NOTHING;"
        else
            sql="${sql}
INSERT INTO ${t} (${cols}) SELECT ${cols} FROM _seed_stage_${idx}
  ON CONFLICT (${keycols}) DO UPDATE SET ${setlist};"
        fi
        if [ "$base" != "-" ]; then
            sql="${sql}
$("$(dirname "$0")/muli-db-seed-delta" "${t}" "${base}" "${keycols}" "${cols}" "_seed_old_${idx}")"
        fi
    done

    sql="${sql}
COMMIT;"

    printf '%s\n' "$sql" | psql -v ON_ERROR_STOP=1 -q -U "$SUPERUSER" -p "$port" -d "$db" || return 1

    local rowcount
    for t in $ordered; do
        rowcount=$(q "$db" "$port" "SELECT count(*) FROM ${t}")
        q "$db" "$port" "INSERT INTO muli_seed_state (table_name, revision, row_count)
                         VALUES ('${t}', '${REVISION}', ${rowcount})
                         ON CONFLICT (table_name) DO UPDATE
                            SET revision=EXCLUDED.revision, row_count=EXCLUDED.row_count, applied_at=now()" >/dev/null
        log "    ${t}: ${rowcount} rows"
    done
}

apply_table() { # apply_table <db> <port> <file> <table> <keys> <mode> <base>
    local db="$1" port="$2" file="$3" table="$4" keycols="$5" mode="$6" base="$7"
    local cols keyjoin setlist rowcount
    cols=$(columns_of "$db" "$port" "$table")

    # key equality between two aliases, and the SET list for an upsert
    # Plain `=` for the same reason as in muli-db-seed-delta: key columns are
    # PRIMARY KEY columns and cannot be NULL, and IS NOT DISTINCT FROM cannot
    # use an index.
    keyjoin=$(echo "$keycols" | tr ',' '\n' | awk '{printf "%s a.%s = b.%s", (NR>1?" AND":""), $1, $1}')
    setlist=$(echo "$cols" | tr ',' '\n' | grep -vxF -f <(echo "$keycols" | tr ',' '\n') \
              | awk '{printf "%s%s=EXCLUDED.%s", (NR>1?", ":""), $1, $1}')

    log "  ${table} <- ${file} (${mode})"

    # Everything for one table in one transaction: a partial load of a tax table
    # is worse than none.
    if ! psql -v ON_ERROR_STOP=1 -q -U "$SUPERUSER" -p "$port" -d "$db" <<SQL
BEGIN;
CREATE TEMP TABLE _seed_stage (LIKE ${table} INCLUDING DEFAULTS) ON COMMIT DROP;
\copy _seed_stage (${cols}) FROM '${SEED_ROOT}/${file}' WITH (FORMAT csv, HEADER true)
-- Indexed before anything joins to it: the full-sync DELETE below and the
-- delta both match on the key columns, and without this they degrade to a
-- sequential scan per row - painful on postcd_aus (16k), unspsc (18k) and
-- calendar (93k).
CREATE INDEX ON _seed_stage (${keycols});
ANALYZE _seed_stage;
$( [ "$base" != "-" ] && echo "CREATE TEMP TABLE _seed_old AS SELECT * FROM ${table};" )

INSERT INTO ${table} (${cols}) SELECT ${cols} FROM _seed_stage
  ON CONFLICT (${keycols}) $( [ "$mode" = "create-only" ] && echo "DO NOTHING" || echo "DO UPDATE SET ${setlist}" );

$( [ "$mode" = "full-sync" ] && cat <<DEL
DELETE FROM ${table} a
 WHERE NOT EXISTS (SELECT 1 FROM _seed_stage b WHERE ${keyjoin});
DEL
)
$( [ "$base" != "-" ] && "$(dirname "$0")/muli-db-seed-delta" "${table}" "${base}" "${keycols}" "${cols}" )
COMMIT;
SQL
    then
        return 1
    fi

    rowcount=$(q "$db" "$port" "SELECT count(*) FROM ${table}")
    q "$db" "$port" "INSERT INTO muli_seed_state (table_name, revision, row_count)
                     VALUES ('${table}', '${REVISION}', ${rowcount})
                     ON CONFLICT (table_name) DO UPDATE
                        SET revision=EXCLUDED.revision, row_count=EXCLUDED.row_count, applied_at=now()" >/dev/null
    log "    ${rowcount} rows"
}

case $# in
    2) TARGETS="$1"; ONE_PORT="$2" ;;
    # One argument is a database name; the port comes from the same rule
    # muli-db-migrate uses, so a caller that only knows which database it is
    # connected to - startm - does not have to work the port out.
    1) TARGETS="$1"; ONE_PORT="" ;;
    0) TARGETS=$(discover); ONE_PORT="" ;;
    *) die "usage: muli-db-seed [--status|--repair[=table]] [<database> [<port>]]" ;;
esac
[ -n "$TARGETS" ] || { log "no databases found"; exit 0; }

log "seed revision ${REVISION} from ${SEED_ROOT}"
[ "$REPAIR" = true ] && log "repair mode${REPAIR_TABLE:+ (${REPAIR_TABLE})}: ignoring recorded revisions" || true
rc=0
for db in $TARGETS; do
    tier=$(tier_of "$db")
    port="${ONE_PORT:-$(port_of "$db")}" || { warn "no port for ${db}"; continue; }
    seed_one_db "$db" "$port" "$tier" || rc=1
done
exit $rc
