setup_blocks.py

"""
Create the Prefect Variable and Secret that prefect_flow.py expects.

Reads the same POSTGRES_* environment variables that sync.py uses and writes
them into your currently-configured Prefect server.

Usage:
    export POSTGRES_HOST=... POSTGRES_PORT=... POSTGRES_DB=...
    export POSTGRES_USER=... POSTGRES_PASSWORD=...
    python setup_blocks.py
"""

import json
import os
import sys

from prefect.blocks.system import Secret
from prefect.variables import Variable


def main() -> int:
    required = ["POSTGRES_HOST", "POSTGRES_PORT", "POSTGRES_DB",
                "POSTGRES_USER", "POSTGRES_PASSWORD"]
    missing = [v for v in required if not os.environ.get(v)]
    if missing:
        print(f"Missing environment variables: {', '.join(missing)}",
              file=sys.stderr)
        return 1

    config = {
        "host": os.environ["POSTGRES_HOST"],
        "port": int(os.environ["POSTGRES_PORT"]),
        "database": os.environ["POSTGRES_DB"],
        "user": os.environ["POSTGRES_USER"],
    }
    Variable.set("postgres_config", json.dumps(config), overwrite=True)
    Secret(value=os.environ["POSTGRES_PASSWORD"]).save(
        "postgres-password", overwrite=True)
    print("Created Variable 'postgres_config' and Secret 'postgres-password'.")
    return 0


if __name__ == "__main__":
    sys.exit(main())