agent_platform_oci/deploy/oke/load-test/scripts/prepare_env.py

#!/usr/bin/env python3
from __future__ import annotations
import argparse
from pathlib import Path
import secrets
import shlex


def parse_env(path: Path) -> dict[str, str]:
    out: dict[str, str] = {}
    if not path.exists():
        return out
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        # Support the common `export KEY=value` form too.
        if line.startswith("export "):
            line = line[7:].lstrip()
        k, v = line.split("=", 1)
        k = k.strip()
        if not k:
            continue
        out[k] = v.strip()
    return out


def write_env(path: Path, values: dict[str, str], *, shell_safe: bool = False):
    lines = ["# Generated by prepare_env.py. Do not commit."]
    for k in sorted(values):
        value = values[k]
        if shell_safe:
            value = shlex.quote(value)
        lines.append(f"{k}={value}")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def merge_missing(target: dict[str, str], defaults: dict[str, str]) -> None:
    """Add only keys that are missing from target; never replace backend values."""
    for k, v in defaults.items():
        target.setdefault(k, v)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--base", required=True,
                   help="Main application env. Existing values are preserved.")
    p.add_argument("--overlay", required=True,
                   help="Load-test defaults. Only fills keys absent from --base.")
    p.add_argument("--override", required=False,
                   help="Optional explicit load-test overrides. These replace base/default values.")
    p.add_argument("--output", required=True)
    p.add_argument("--langfuse-env", required=True)
    args = p.parse_args()

    base_path = Path(args.base)
    defaults_path = Path(args.overlay)
    override_path = Path(args.override) if args.override else None

    base = parse_env(base_path)
    defaults = parse_env(defaults_path)
    explicit = parse_env(override_path) if override_path and override_path.exists() else {}

    # Backend .env is authoritative. The example file provides only missing load-test keys.
    merged = dict(base)
    merge_missing(merged, defaults)

    # .env.loadtest, when present, is the deliberate operator override layer.
    merged.update(explicit)

    # Container/OKE invariants. These are location/routing details that cannot use
    # workstation-local values from templates/agent_template_backend/.env.
    merged["APP_ENV"] = "oke-load-test"
    merged["ADB_WALLET_LOCATION"] = "/app/wallet"
    merged["ENABLE_LANGFUSE"] = "true"
    merged["LANGFUSE_HOST"] = "http://langfuse-web.langfuse.svc.cluster.local:3000"

    # OKE authentication should normally use workload identity. Allow an explicit
    # .env.loadtest override if the operator intentionally selected another mode.
    if "OCI_AUTH_MODE" not in explicit:
        merged["OCI_AUTH_MODE"] = "oke_workload_identity"

    # Correct only the historical typo in Mongo-compatible ports; preserve host,
    # credentials, query params and database values from the backend .env.
    for key in ("MONGODB_URI", "PUBSUB_SEQUENCE_MONGODB_URI"):
        if key in merged:
            merged[key] = merged[key].replace(":37017", ":27017")

    # Langfuse bootstrap secrets are stable across reruns.
    lf_path = Path(args.langfuse_env)
    lf = parse_env(lf_path)
    public_key = lf.get("LANGFUSE_PUBLIC_KEY") or "pk-lf-" + secrets.token_hex(16)
    secret_key = lf.get("LANGFUSE_SECRET_KEY") or "sk-lf-" + secrets.token_hex(32)
    admin_password = lf.get("LANGFUSE_INIT_USER_PASSWORD") or secrets.token_urlsafe(24)

    lf.update({
        "LANGFUSE_INIT_ORG_ID": lf.get("LANGFUSE_INIT_ORG_ID", "agent-framework-org"),
        "LANGFUSE_INIT_ORG_NAME": lf.get("LANGFUSE_INIT_ORG_NAME", "Agent Framework OCI"),
        "LANGFUSE_INIT_PROJECT_ID": lf.get("LANGFUSE_INIT_PROJECT_ID", "agent-framework-loadtest"),
        "LANGFUSE_INIT_PROJECT_NAME": lf.get("LANGFUSE_INIT_PROJECT_NAME", "Agent Framework Load Test"),
        "LANGFUSE_PUBLIC_KEY": public_key,
        "LANGFUSE_SECRET_KEY": secret_key,
        "LANGFUSE_INIT_USER_EMAIL": lf.get("LANGFUSE_INIT_USER_EMAIL", "admin@loadtest.local"),
        "LANGFUSE_INIT_USER_NAME": lf.get("LANGFUSE_INIT_USER_NAME", "Load Test Admin"),
        "LANGFUSE_INIT_USER_PASSWORD": admin_password,
        "LANGFUSE_NEXTAUTH_SECRET": lf.get("LANGFUSE_NEXTAUTH_SECRET") or secrets.token_hex(32),
        "LANGFUSE_SALT": lf.get("LANGFUSE_SALT") or secrets.token_hex(16),
        "LANGFUSE_ENCRYPTION_KEY": lf.get("LANGFUSE_ENCRYPTION_KEY") or secrets.token_hex(32),
        "LANGFUSE_POSTGRES_PASSWORD": lf.get("LANGFUSE_POSTGRES_PASSWORD") or secrets.token_urlsafe(24),
        "LANGFUSE_REDIS_PASSWORD": lf.get("LANGFUSE_REDIS_PASSWORD") or secrets.token_urlsafe(24),
        "LANGFUSE_CLICKHOUSE_PASSWORD": lf.get("LANGFUSE_CLICKHOUSE_PASSWORD") or secrets.token_urlsafe(24),
        "LANGFUSE_MINIO_PASSWORD": lf.get("LANGFUSE_MINIO_PASSWORD") or secrets.token_urlsafe(24),
    })
    lf.update({
        "DATABASE_URL": f"postgresql://postgres:{lf['LANGFUSE_POSTGRES_PASSWORD']}@postgres.langfuse.svc.cluster.local:5432/postgres",
        "DIRECT_URL": f"postgresql://postgres:{lf['LANGFUSE_POSTGRES_PASSWORD']}@postgres.langfuse.svc.cluster.local:5432/postgres",
        "POSTGRES_PASSWORD": lf["LANGFUSE_POSTGRES_PASSWORD"],
        "SALT": lf["LANGFUSE_SALT"],
        "ENCRYPTION_KEY": lf["LANGFUSE_ENCRYPTION_KEY"],
        "NEXTAUTH_SECRET": lf["LANGFUSE_NEXTAUTH_SECRET"],
        "CLICKHOUSE_PASSWORD": lf["LANGFUSE_CLICKHOUSE_PASSWORD"],
        "REDIS_AUTH": lf["LANGFUSE_REDIS_PASSWORD"],
        "LANGFUSE_S3_EVENT_UPLOAD_SECRET_ACCESS_KEY": lf["LANGFUSE_MINIO_PASSWORD"],
        "LANGFUSE_S3_MEDIA_UPLOAD_SECRET_ACCESS_KEY": lf["LANGFUSE_MINIO_PASSWORD"],
        "LANGFUSE_S3_BATCH_EXPORT_SECRET_ACCESS_KEY": lf["LANGFUSE_MINIO_PASSWORD"],
        "LANGFUSE_INIT_PROJECT_PUBLIC_KEY": public_key,
        "LANGFUSE_INIT_PROJECT_SECRET_KEY": secret_key,
    })
    write_env(lf_path, lf, shell_safe=False)

    merged["LANGFUSE_PUBLIC_KEY"] = public_key
    merged["LANGFUSE_SECRET_KEY"] = secret_key
    write_env(Path(args.output), merged, shell_safe=True)

    preserved = sum(1 for k, v in base.items() if merged.get(k) == v)
    changed = sorted(k for k, v in base.items() if merged.get(k) != v)
    print(f"Base env: {base_path} ({len(base)} variables)")
    print(f"Defaults: {defaults_path} ({len(defaults)} variables)")
    if override_path:
        print(f"Explicit override: {override_path} ({len(explicit)} variables)")
    print(f"Runtime env written: {args.output} ({len(merged)} variables)")
    print(f"Backend values preserved unchanged: {preserved}/{len(base)}")
    if changed:
        print("Backend values intentionally changed for OKE/container compatibility:")
        for k in changed:
            print(f"  - {k}")
    print(f"Langfuse env written: {args.langfuse_env}")
    print("LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY synchronized automatically.")

if __name__ == "__main__":
    main()