PrivacyBudget, Differential Privacy Export Gate

id: privacybudget-differential-privacy-export-gate
namespace: company.data

description: |
  Reserve epsilon atomically in PostgreSQL, run bounded differential-privacy
  aggregates only for approved requests, publish the sanitized export to S3,
  and retain an auditable decision for approved and rejected requests.

concurrency:
  limit: 1
  behavior: QUEUE

inputs:
  - id: demo_mode
    type: BOOL
    defaults: true
    description: Start disposable PostgreSQL and LocalStack containers with
      deterministic fixtures. Set false for production services.

  - id: budget_subject
    type: STRING
    defaults: analytics-demo
    description: Stable subject whose cumulative epsilon budget is stored and locked
      in PostgreSQL.

  - id: requested_epsilon
    type: FLOAT
    defaults: 0.4
    description: Total epsilon charged for this export and split equally across its
      two aggregate mechanisms.

  - id: demo_budget_limit
    type: FLOAT
    defaults: 1.0
    description: Budget seeded only in demo mode. Production limits must already
      exist in PostgreSQL.

  - id: postgres_url
    type: STRING
    defaults: jdbc:postgresql://postgres.internal:5432/privacy
    description: Production JDBC URL containing the privacy budget, source table,
      and audit ledger. Ignored in demo mode.

  - id: s3_bucket
    type: STRING
    defaults: privacybudget-exports
    description: Existing production S3 bucket; the deterministic demo creates this
      bucket in LocalStack.

  - id: s3_region
    type: STRING
    defaults: us-east-1
    description: AWS region used by the S3 create, upload, and round-trip
      verification tasks.

  - id: s3_endpoint_override
    type: STRING
    defaults: https://s3.amazonaws.com
    description: Production S3-compatible endpoint. Ignored in demo mode.

  - id: s3_compatibility_mode
    type: BOOL
    defaults: false
    description: Enable compatibility and path-style addressing for a production
      S3-compatible service.

variables:
  object_key: "privacybudget/{{ inputs.budget_subject }}/{{ execution.id }}/aggregate.json"
  resolved_postgres_url: "{{ inputs.demo_mode ?
    'jdbc:postgresql://host.docker.internal:25432/privacybudget' :
    inputs.postgres_url }}"
  resolved_postgres_user: "{{ inputs.demo_mode ? 'privacybudget' : secret('POSTGRES_USERNAME') }}"
  resolved_postgres_password: "{{ inputs.demo_mode ? 'privacybudget-demo-password'
    : secret('POSTGRES_PASSWORD') }}"
  resolved_s3_endpoint: "{{ inputs.demo_mode ? 'http://host.docker.internal:29000'
    : inputs.s3_endpoint_override }}"
  resolved_s3_access_key: "{{ inputs.demo_mode ? 'privacybudget-demo' : secret('AWS_ACCESS_KEY_ID') }}"
  resolved_s3_secret_key: "{{ inputs.demo_mode ? 'privacybudget-demo-secret' :
    secret('AWS_SECRET_ACCESS_KEY') }}"
  resolved_s3_compatibility: "{{ inputs.demo_mode ? true : inputs.s3_compatibility_mode }}"
  demo_compose: |
    services:
      postgres:
        image: postgres:17.6-alpine
        environment:
          POSTGRES_DB: privacybudget
          POSTGRES_USER: privacybudget
          POSTGRES_PASSWORD: privacybudget-demo-password
        ports:
          - "127.0.0.1:25432:5432"
        healthcheck:
          test: ["CMD-SHELL", "pg_isready -U privacybudget -d privacybudget"]
          interval: 2s
          timeout: 2s
          retries: 30
      localstack:
        image: localstack/localstack:4.0.3
        environment:
          SERVICES: s3
          DEBUG: 0
        ports:
          - "127.0.0.1:29000:4566"

triggers:
  - id: daily_privacy_export
    type: io.kestra.plugin.core.trigger.Schedule
    description: Disabled-by-default production export; enable after the budget row,
      dataset, bucket, and secrets are configured.
    cron: "0 4 * * *"
    timezone: UTC
    disabled: true
    inputs:
      demo_mode: false
      requested_epsilon: 0.2

tasks:
  - id: validate_request
    type: io.kestra.plugin.core.flow.If
    description: Reject malformed identifiers, unsafe epsilon values, or missing
      destinations before touching external systems.
    condition: >-
      {{ inputs.budget_subject is empty
         or inputs.s3_bucket is empty
         or inputs.requested_epsilon <= 0
         or inputs.requested_epsilon > 10
         or (inputs.demo_mode and inputs.demo_budget_limit <= 0) }}
    then:
      - id: reject_invalid_request
        type: io.kestra.plugin.core.execution.Fail
        errorMessage: budget_subject and s3_bucket must be non-empty; epsilon values
          must be positive and requested_epsilon must not exceed 10.

  - id: provision_demo_services
    type: io.kestra.plugin.core.flow.If
    description: Start pinned PostgreSQL and LocalStack containers only for the
      deterministic local demonstration.
    condition: "{{ inputs.demo_mode }}"
    then:
      - id: compose_up
        type: io.kestra.plugin.docker.cli.Compose
        taskRunner:
          type: io.kestra.plugin.scripts.runner.docker.Docker
          user: "0"
          volumes:
            - /var/run/docker.sock:/var/run/docker.sock
        composeFile: "{{ render(vars.demo_compose) }}"
        composeArgs:
          - --project-name
          - "privacybudget-{{ execution.id | lower }}"
          - up
          - -d
          - --wait
        timeout: PT10M
        retry:
          type: exponential
          interval: PT2S
          maxInterval: PT10S
          maxAttempts: 3

  - id: prepare_privacy_store
    type: io.kestra.plugin.jdbc.postgresql.Queries
    description: Create the budget ledger, source-table contract, immutable request
      evidence, and atomic idempotent reservation function.
    url: "{{ render(vars.resolved_postgres_url) }}"
    username: "{{ render(vars.resolved_postgres_user) }}"
    password: "{{ render(vars.resolved_postgres_password) }}"
    fetchType: NONE
    sql: |
      CREATE SCHEMA IF NOT EXISTS privacybudget;

      CREATE TABLE IF NOT EXISTS privacybudget.budgets (
        subject_id text PRIMARY KEY,
        epsilon_limit numeric(12,6) NOT NULL CHECK (epsilon_limit > 0),
        epsilon_spent numeric(12,6) NOT NULL DEFAULT 0 CHECK (epsilon_spent >= 0 AND epsilon_spent <= epsilon_limit),
        reset_at timestamptz,
        updated_at timestamptz NOT NULL DEFAULT now()
      );

      CREATE TABLE IF NOT EXISTS privacybudget.daily_metrics (
        metric_date date PRIMARY KEY,
        order_value numeric(10,2) NOT NULL CHECK (order_value BETWEEN 0 AND 500)
      );

      CREATE TABLE IF NOT EXISTS privacybudget.export_audits (
        request_id text PRIMARY KEY,
        subject_id text NOT NULL,
        requested_epsilon numeric(12,6) NOT NULL,
        epsilon_before numeric(12,6),
        epsilon_after numeric(12,6),
        decision text NOT NULL CHECK (decision IN ('APPROVED', 'REJECTED')),
        reason text NOT NULL,
        status text NOT NULL CHECK (status IN ('RESERVED', 'REJECTED', 'PUBLISHED', 'FAILED')),
        s3_bucket text NOT NULL,
        s3_key text NOT NULL,
        algorithms jsonb,
        export_sha256 text,
        error_message text,
        created_at timestamptz NOT NULL DEFAULT now(),
        published_at timestamptz
      );

      CREATE OR REPLACE FUNCTION privacybudget.reserve_export(
        p_request_id text,
        p_subject_id text,
        p_requested_epsilon numeric,
        p_s3_bucket text,
        p_s3_key text
      ) RETURNS TABLE (
        decision text,
        reason text,
        epsilon_before numeric,
        epsilon_after numeric,
        status text
      ) LANGUAGE plpgsql AS $$
      DECLARE
        existing privacybudget.export_audits%ROWTYPE;
        budget privacybudget.budgets%ROWTYPE;
      BEGIN
        PERFORM pg_advisory_xact_lock(hashtextextended(p_request_id, 0));

        SELECT * INTO existing
        FROM privacybudget.export_audits
        WHERE request_id = p_request_id;

        IF FOUND THEN
          RETURN QUERY SELECT existing.decision, existing.reason,
            existing.epsilon_before, existing.epsilon_after, existing.status;
          RETURN;
        END IF;

        SELECT * INTO budget
        FROM privacybudget.budgets
        WHERE subject_id = p_subject_id
        FOR UPDATE;

        IF NOT FOUND THEN
          INSERT INTO privacybudget.export_audits (
            request_id, subject_id, requested_epsilon, decision, reason,
            status, s3_bucket, s3_key
          ) VALUES (
            p_request_id, p_subject_id, p_requested_epsilon, 'REJECTED',
            'NO_STORED_BUDGET', 'REJECTED', p_s3_bucket, p_s3_key
          );
          RETURN QUERY SELECT 'REJECTED'::text, 'NO_STORED_BUDGET'::text,
            NULL::numeric, NULL::numeric, 'REJECTED'::text;
          RETURN;
        END IF;

        IF budget.epsilon_spent + p_requested_epsilon > budget.epsilon_limit THEN
          INSERT INTO privacybudget.export_audits (
            request_id, subject_id, requested_epsilon, epsilon_before,
            epsilon_after, decision, reason, status, s3_bucket, s3_key
          ) VALUES (
            p_request_id, p_subject_id, p_requested_epsilon,
            budget.epsilon_limit - budget.epsilon_spent,
            budget.epsilon_limit - budget.epsilon_spent,
            'REJECTED', 'BUDGET_EXCEEDED', 'REJECTED', p_s3_bucket, p_s3_key
          );
          RETURN QUERY SELECT 'REJECTED'::text, 'BUDGET_EXCEEDED'::text,
            (budget.epsilon_limit - budget.epsilon_spent)::numeric,
            (budget.epsilon_limit - budget.epsilon_spent)::numeric,
            'REJECTED'::text;
          RETURN;
        END IF;

        UPDATE privacybudget.budgets
        SET epsilon_spent = epsilon_spent + p_requested_epsilon,
            updated_at = now()
        WHERE subject_id = p_subject_id;

        INSERT INTO privacybudget.export_audits (
          request_id, subject_id, requested_epsilon, epsilon_before,
          epsilon_after, decision, reason, status, s3_bucket, s3_key
        ) VALUES (
          p_request_id, p_subject_id, p_requested_epsilon,
          budget.epsilon_limit - budget.epsilon_spent,
          budget.epsilon_limit - budget.epsilon_spent - p_requested_epsilon,
          'APPROVED', 'WITHIN_BUDGET', 'RESERVED', p_s3_bucket, p_s3_key
        );

        RETURN QUERY SELECT 'APPROVED'::text, 'WITHIN_BUDGET'::text,
          (budget.epsilon_limit - budget.epsilon_spent)::numeric,
          (budget.epsilon_limit - budget.epsilon_spent - p_requested_epsilon)::numeric,
          'RESERVED'::text;
      END;
      $$;
    timeout: PT2M
    retry:
      type: exponential
      interval: PT2S
      maxInterval: PT10S
      maxAttempts: 5

  - id: seed_demo_fixtures
    type: io.kestra.plugin.core.flow.If
    description: Seed a fixed budget and bounded twelve-row analytics fixture only
      in disposable demo PostgreSQL.
    condition: "{{ inputs.demo_mode }}"
    then:
      - id: seed_postgres
        type: io.kestra.plugin.jdbc.postgresql.Queries
        url: "{{ render(vars.resolved_postgres_url) }}"
        username: "{{ render(vars.resolved_postgres_user) }}"
        password: "{{ render(vars.resolved_postgres_password) }}"
        fetchType: NONE
        sql: |
          INSERT INTO privacybudget.budgets (subject_id, epsilon_limit, epsilon_spent)
          VALUES ('{{ inputs.budget_subject }}', {{ inputs.demo_budget_limit }}, 0)
          ON CONFLICT (subject_id) DO UPDATE SET
            epsilon_limit = EXCLUDED.epsilon_limit,
            epsilon_spent = 0,
            updated_at = now();

          INSERT INTO privacybudget.daily_metrics (metric_date, order_value) VALUES
            ('2026-01-01', 84.25), ('2026-01-02', 91.10),
            ('2026-01-03', 77.90), ('2026-01-04', 102.40),
            ('2026-01-05', 88.70), ('2026-01-06', 95.20),
            ('2026-01-07', 110.00), ('2026-01-08', 72.35),
            ('2026-01-09', 98.80), ('2026-01-10', 86.45),
            ('2026-01-11', 105.60), ('2026-01-12', 93.15)
          ON CONFLICT (metric_date) DO UPDATE SET order_value = EXCLUDED.order_value;
        timeout: PT1M

      - id: create_demo_bucket
        type: io.kestra.plugin.aws.s3.CreateBucket
        description: Create the LocalStack bucket through the same native AWS plugin
          used for production S3.
        accessKeyId: "{{ render(vars.resolved_s3_access_key) }}"
        secretKeyId: "{{ render(vars.resolved_s3_secret_key) }}"
        region: "{{ inputs.s3_region }}"
        bucket: "{{ inputs.s3_bucket }}"
        endpointOverride: "{{ render(vars.resolved_s3_endpoint) }}"
        compatibilityMode: true
        forcePathStyle: true
        timeout: PT1M
        retry:
          type: exponential
          interval: PT2S
          maxInterval: PT10S
          maxAttempts: 5

  - id: reserve_budget
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: Lock the subject ledger and atomically approve and charge or reject
      the request; retries return the original decision without double-spending.
    url: "{{ render(vars.resolved_postgres_url) }}"
    username: "{{ render(vars.resolved_postgres_user) }}"
    password: "{{ render(vars.resolved_postgres_password) }}"
    fetchType: FETCH_ONE
    parameters:
      request_id: "{{ execution.id }}"
      subject_id: "{{ inputs.budget_subject }}"
      requested_epsilon: "{{ inputs.requested_epsilon }}"
      s3_bucket: "{{ inputs.s3_bucket }}"
      s3_key: "{{ render(vars.object_key) }}"
    sql: |
      SELECT *
      FROM privacybudget.reserve_export(
        :request_id,
        :subject_id,
        CAST(:requested_epsilon AS numeric),
        :s3_bucket,
        :s3_key
      );
    timeout: PT1M
    retry:
      type: exponential
      interval: PT2S
      maxInterval: PT10S
      maxAttempts: 3

  - id: privacy_export_gate
    type: io.kestra.plugin.core.flow.If
    description: Run mechanisms and publish only after PostgreSQL returns APPROVED;
      otherwise fail with the durable rejection reason.
    condition: "{{ outputs.reserve_budget.row.decision == 'APPROVED' }}"
    then:
      - id: calculate_private_aggregates
        type: io.kestra.plugin.core.flow.Sequential
        description: Run two bounded diffprivlib mechanisms in order so both reuse one
          pinned dependency cache, each consuming half of the reserved epsilon.
        tasks:
          - id: private_mean_order_value
            type: io.kestra.plugin.scripts.python.Script
            description: Query bounded values from PostgreSQL and calculate a differentially
              private mean without logging raw rows.
            taskRunner:
              type: io.kestra.plugin.scripts.runner.docker.Docker
              extraHosts:
                - host.docker.internal:host-gateway
            containerImage: python:3.12.7-slim
            dependencies:
              - kestra==0.23.0
              - diffprivlib==0.6.6
              - numpy==2.1.3
              - scipy==1.14.1
              - scikit-learn==1.6.1
              - psycopg[binary]==3.2.10
            env:
              POSTGRES_DSN: "{{ render(vars.resolved_postgres_url) | replace({'jdbc:': ''})
                }}"
              POSTGRES_USER: "{{ render(vars.resolved_postgres_user) }}"
              POSTGRES_PASSWORD: "{{ render(vars.resolved_postgres_password) }}"
              MECHANISM_EPSILON: "{{ inputs.requested_epsilon / 2 }}"
              DEMO_MODE: "{{ inputs.demo_mode }}"
            timeout: PT10M
            retry:
              type: exponential
              interval: PT2S
              maxInterval: PT15S
              maxAttempts: 3
            script: |
              import os
              import numpy as np
              import psycopg
              from diffprivlib.tools import mean
              from kestra import Kestra

              with psycopg.connect(
                  os.environ["POSTGRES_DSN"],
                  user=os.environ["POSTGRES_USER"],
                  password=os.environ["POSTGRES_PASSWORD"],
                  connect_timeout=10,
              ) as connection:
                  with connection.cursor() as cursor:
                      cursor.execute("SELECT order_value::double precision FROM privacybudget.daily_metrics ORDER BY metric_date")
                      values = np.asarray([row[0] for row in cursor.fetchall()], dtype=float)

              if values.size == 0:
                  raise RuntimeError("privacybudget.daily_metrics contains no rows")
              if np.any((values < 0) | (values > 500)):
                  raise RuntimeError("order_value violates the declared [0, 500] contribution bounds")

              epsilon = float(os.environ["MECHANISM_EPSILON"])
              seed = 7301 if os.environ["DEMO_MODE"].lower() == "true" else None
              private_value = float(mean(values, epsilon=epsilon, bounds=(0, 500), random_state=seed))
              Kestra.outputs({
                  "value": round(private_value, 6),
                  "algorithm": "diffprivlib.tools.mean",
                  "epsilon": epsilon,
                  "bounds": "[0,500]",
              })

          - id: private_record_count
            type: io.kestra.plugin.scripts.python.Script
            description: Calculate a differentially private row count with sensitivity one
              and no raw records in task outputs.
            taskRunner:
              type: io.kestra.plugin.scripts.runner.docker.Docker
              extraHosts:
                - host.docker.internal:host-gateway
            containerImage: python:3.12.7-slim
            dependencies:
              - kestra==0.23.0
              - diffprivlib==0.6.6
              - numpy==2.1.3
              - scipy==1.14.1
              - scikit-learn==1.6.1
              - psycopg[binary]==3.2.10
            env:
              POSTGRES_DSN: "{{ render(vars.resolved_postgres_url) | replace({'jdbc:': ''})
                }}"
              POSTGRES_USER: "{{ render(vars.resolved_postgres_user) }}"
              POSTGRES_PASSWORD: "{{ render(vars.resolved_postgres_password) }}"
              MECHANISM_EPSILON: "{{ inputs.requested_epsilon / 2 }}"
              DEMO_MODE: "{{ inputs.demo_mode }}"
            timeout: PT10M
            retry:
              type: exponential
              interval: PT2S
              maxInterval: PT15S
              maxAttempts: 3
            script: |
              import os
              import numpy as np
              import psycopg
              from diffprivlib.tools import count_nonzero
              from kestra import Kestra

              with psycopg.connect(
                  os.environ["POSTGRES_DSN"],
                  user=os.environ["POSTGRES_USER"],
                  password=os.environ["POSTGRES_PASSWORD"],
                  connect_timeout=10,
              ) as connection:
                  with connection.cursor() as cursor:
                      cursor.execute("SELECT 1 FROM privacybudget.daily_metrics")
                      indicators = np.ones(len(cursor.fetchall()), dtype=int)

              if indicators.size == 0:
                  raise RuntimeError("privacybudget.daily_metrics contains no rows")

              epsilon = float(os.environ["MECHANISM_EPSILON"])
              seed = 7302 if os.environ["DEMO_MODE"].lower() == "true" else None
              private_value = int(count_nonzero(indicators, epsilon=epsilon, random_state=seed))
              Kestra.outputs({
                  "value": max(0, private_value),
                  "algorithm": "diffprivlib.tools.count_nonzero",
                  "epsilon": epsilon,
                  "sensitivity": 1,
              })

      - id: build_private_export
        type: io.kestra.plugin.scripts.python.Script
        description: Build a privacy-safe JSON dataset containing only noisy aggregates
          and budget provenance.
        taskRunner:
          type: io.kestra.plugin.scripts.runner.docker.Docker
        containerImage: python:3.12.7-slim
        dependencies:
          - kestra==0.23.0
        env:
          EXECUTION_ID: "{{ execution.id }}"
          SUBJECT_ID: "{{ inputs.budget_subject }}"
          REQUESTED_EPSILON: "{{ inputs.requested_epsilon }}"
          REMAINING_EPSILON: "{{ outputs.reserve_budget.row.epsilon_after }}"
          PRIVATE_MEAN: "{{ outputs.private_mean_order_value.vars.value }}"
          PRIVATE_COUNT: "{{ outputs.private_record_count.vars.value }}"
        outputFiles:
          - aggregate.json
        timeout: PT2M
        script: |
          import json
          import os
          from kestra import Kestra

          export = {
              "request_id": os.environ["EXECUTION_ID"],
              "subject_id": os.environ["SUBJECT_ID"],
              "privacy": {
                  "epsilon_spent": float(os.environ["REQUESTED_EPSILON"]),
                  "epsilon_remaining": float(os.environ["REMAINING_EPSILON"]),
                  "composition": "sequential; epsilon split equally across two mechanisms",
              },
              "aggregates": {
                  "private_mean_order_value": float(os.environ["PRIVATE_MEAN"]),
                  "private_record_count": int(os.environ["PRIVATE_COUNT"]),
              },
              "algorithms": [
                  {"name": "diffprivlib.tools.mean", "epsilon_share": 0.5, "bounds": [0, 500]},
                  {"name": "diffprivlib.tools.count_nonzero", "epsilon_share": 0.5, "sensitivity": 1},
              ],
          }
          with open("aggregate.json", "w", encoding="utf-8") as output_file:
              json.dump(export, output_file, indent=2, sort_keys=True)
          Kestra.outputs({"aggregate_count": 2})

      - id: publish_private_export
        type: io.kestra.plugin.aws.s3.Upload
        description: Publish the permitted aggregate dataset through the native S3 API;
          no raw analytics rows leave PostgreSQL.
        accessKeyId: "{{ render(vars.resolved_s3_access_key) }}"
        secretKeyId: "{{ render(vars.resolved_s3_secret_key) }}"
        region: "{{ inputs.s3_region }}"
        bucket: "{{ inputs.s3_bucket }}"
        key: "{{ render(vars.object_key) }}"
        from: "{{ outputs.build_private_export.outputFiles['aggregate.json'] }}"
        endpointOverride: "{{ render(vars.resolved_s3_endpoint) }}"
        compatibilityMode: "{{ render(vars.resolved_s3_compatibility) }}"
        forcePathStyle: "{{ render(vars.resolved_s3_compatibility) }}"
        timeout: PT2M
        retry:
          type: exponential
          interval: PT2S
          maxInterval: PT15S
          maxAttempts: 5

      - id: download_published_export
        type: io.kestra.plugin.aws.s3.Download
        description: Read the object back through S3 before marking the audit row published.
        accessKeyId: "{{ render(vars.resolved_s3_access_key) }}"
        secretKeyId: "{{ render(vars.resolved_s3_secret_key) }}"
        region: "{{ inputs.s3_region }}"
        bucket: "{{ inputs.s3_bucket }}"
        key: "{{ render(vars.object_key) }}"
        endpointOverride: "{{ render(vars.resolved_s3_endpoint) }}"
        compatibilityMode: "{{ render(vars.resolved_s3_compatibility) }}"
        forcePathStyle: "{{ render(vars.resolved_s3_compatibility) }}"
        timeout: PT2M
        retry:
          type: exponential
          interval: PT2S
          maxInterval: PT15S
          maxAttempts: 5

      - id: verify_published_export
        type: io.kestra.plugin.scripts.python.Script
        description: Parse the S3 round trip, reject unexpected raw fields, and emit a
          SHA-256 evidence value.
        taskRunner:
          type: io.kestra.plugin.scripts.runner.docker.Docker
        containerImage: python:3.12.7-slim
        dependencies:
          - kestra==0.23.0
        inputFiles:
          downloaded.json: "{{ outputs.download_published_export.uri }}"
        timeout: PT2M
        script: |
          import hashlib
          import json
          from pathlib import Path
          from kestra import Kestra

          payload = Path("downloaded.json").read_bytes()
          document = json.loads(payload)
          expected = {"request_id", "subject_id", "privacy", "aggregates", "algorithms"}
          if set(document) != expected:
              raise RuntimeError(f"unexpected export fields: {sorted(set(document) - expected)}")
          if set(document["aggregates"]) != {"private_mean_order_value", "private_record_count"}:
              raise RuntimeError("export aggregate contract is invalid")
          Kestra.outputs({"sha256": hashlib.sha256(payload).hexdigest(), "verified": True})

      - id: mark_export_published
        type: io.kestra.plugin.jdbc.postgresql.Query
        description: Commit the algorithm manifest and verified object hash to the audit
          ledger only after the S3 round trip succeeds.
        url: "{{ render(vars.resolved_postgres_url) }}"
        username: "{{ render(vars.resolved_postgres_user) }}"
        password: "{{ render(vars.resolved_postgres_password) }}"
        fetchType: FETCH_ONE
        parameters:
          request_id: "{{ execution.id }}"
          export_sha256: "{{ outputs.verify_published_export.vars.sha256 }}"
          private_mean: "{{ outputs.private_mean_order_value.vars.value }}"
          private_count: "{{ outputs.private_record_count.vars.value }}"
        sql: |
          UPDATE privacybudget.export_audits
          SET status = 'PUBLISHED',
              algorithms = jsonb_build_object(
                'diffprivlib_version', '0.6.6',
                'private_mean_order_value', CAST(:private_mean AS numeric),
                'private_record_count', CAST(:private_count AS integer)
              ),
              export_sha256 = :export_sha256,
              published_at = now()
          WHERE request_id = :request_id AND decision = 'APPROVED'
          RETURNING request_id, decision, status, epsilon_after,
                    s3_bucket, s3_key, export_sha256;
        timeout: PT1M
        retry:
          type: exponential
          interval: PT2S
          maxInterval: PT10S
          maxAttempts: 3
    else:
      - id: reject_over_budget_export
        type: io.kestra.plugin.core.execution.Fail
        errorMessage: >-
          PrivacyBudget rejected request {{ execution.id }}: {{
          outputs.reserve_budget.row.reason }}; remaining epsilon is {{
          outputs.reserve_budget.row.epsilon_after }} and requested epsilon is
          {{ inputs.requested_epsilon }}. No algorithm ran and no S3 object was
          written.

  - id: teardown_demo_services
    type: io.kestra.plugin.core.flow.If
    description: Remove demo containers and volumes after a successful verified export.
    condition: "{{ inputs.demo_mode }}"
    then:
      - id: compose_down
        type: io.kestra.plugin.docker.cli.Compose
        taskRunner:
          type: io.kestra.plugin.scripts.runner.docker.Docker
          user: "0"
          volumes:
            - /var/run/docker.sock:/var/run/docker.sock
        composeFile: "{{ render(vars.demo_compose) }}"
        composeArgs:
          - --project-name
          - "privacybudget-{{ execution.id | lower }}"
          - down
          - -v
          - --remove-orphans
        timeout: PT2M

errors:
  - id: record_failed_export
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: Best-effort failure bookkeeping that preserves an intentional
      REJECTED decision and marks interrupted approved exports FAILED.
    allowFailure: true
    url: "{{ render(vars.resolved_postgres_url) }}"
    username: "{{ render(vars.resolved_postgres_user) }}"
    password: "{{ render(vars.resolved_postgres_password) }}"
    fetchType: NONE
    parameters:
      request_id: "{{ execution.id }}"
    sql: |
      UPDATE privacybudget.export_audits
      SET status = CASE WHEN decision = 'REJECTED' THEN 'REJECTED' ELSE 'FAILED' END,
          error_message = CASE
            WHEN decision = 'REJECTED' THEN reason
            ELSE 'Execution failed before a verified S3 publication'
          END
      WHERE request_id = :request_id;
    timeout: PT1M

  - id: teardown_demo_after_failure
    type: io.kestra.plugin.docker.cli.Compose
    description: Best-effort cleanup for rejected requests and infrastructure or
      algorithm failures.
    runIf: "{{ inputs.demo_mode }}"
    allowFailure: true
    taskRunner:
      type: io.kestra.plugin.scripts.runner.docker.Docker
      user: "0"
      volumes:
        - /var/run/docker.sock:/var/run/docker.sock
    composeFile: "{{ render(vars.demo_compose) }}"
    composeArgs:
      - --project-name
      - "privacybudget-{{ execution.id | lower }}"
      - down
      - -v
      - --remove-orphans
    timeout: PT2M

outputs:
  - id: decision
    type: STRING
    value: "{{ outputs.reserve_budget is defined ?
      outputs.reserve_budget.row.decision : 'ERROR' }}"

  - id: reason
    type: STRING
    value: "{{ outputs.reserve_budget is defined ? outputs.reserve_budget.row.reason
      : 'EXECUTION_FAILED_BEFORE_RESERVATION' }}"

  - id: remaining_epsilon
    type: FLOAT
    value: "{{ outputs.reserve_budget is defined and
      outputs.reserve_budget.row.epsilon_after is not null ?
      outputs.reserve_budget.row.epsilon_after : 0.0 }}"

  - id: s3_object
    type: STRING
    value: "s3://{{ inputs.s3_bucket }}/{{ render(vars.object_key) }}"