Model/scripts/e2e_common.py
Khalim Conn-Kowlessar ea72ee97bf feat(scripts): add full AraFirstRunPipeline local runner
scripts/run_first_run_e2e.py runs the real Ingestion -> Baseline -> Modelling
pipeline against the DB by composing build_first_run_pipeline + dispatch_first_run
with the live source clients (the Lambda handler can't run locally — its
_source_clients_from_env still raises, #1136). Unlike run_modelling_e2e it runs
real ingestion (persists EPC/spatial/solar) and has no inspect-only mode, so it's
gated behind --confirm (preview otherwise); measure scoping comes only from the
Scenario's exclusions (the pipeline threads no --measures), and the modelling
batch is all-or-nothing, both documented.

Extract the shared env/engine/S3 plumbing into scripts/e2e_common.py (public
load_env/build_engine/s3_parquet_reader) so both runners share one source and
neither imports the other's privates.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-16 23:45:23 +00:00

68 lines
2.4 KiB
Python

"""Shared configuration + client plumbing for the local e2e runner scripts
(``run_modelling_e2e`` and ``run_first_run_e2e``).
Loads ``backend/.env`` and builds the DB engine from the FastAPI-layer ``DB_*``
vars (the ``infrastructure/postgres`` layer reads ``POSTGRES_*``, which the .env
does not carry), plus an S3-backed ``ParquetReader`` for the geospatial
repository. Secrets live in the .env and the ambient ``~/.aws`` profile; this
module never hard-codes them.
"""
from __future__ import annotations
import io
import os
from pathlib import Path
from typing import Any, cast
import boto3
import pandas as pd
from sqlalchemy import Engine, create_engine
from repositories.geospatial.geospatial_s3_repository import ParquetReader
_REPO_ROOT = Path(__file__).resolve().parents[1]
ENV_PATH = _REPO_ROOT / "backend" / ".env"
def load_env(path: Path = ENV_PATH) -> None:
"""Load `KEY=value` lines from `backend/.env` into the environment (without
overriding anything already set), so the DB creds + API tokens are present."""
if not path.exists():
return
for raw in path.read_text(encoding="utf-8").splitlines():
line = raw.strip()
if not line or line.startswith("#") or "=" not in line:
continue
key, value = line.split("=", 1)
os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))
def db_url() -> str:
"""The connection string from the FastAPI-layer `DB_*` env vars."""
env = os.environ
return (
f"postgresql+psycopg2://{env['DB_USERNAME']}:{env['DB_PASSWORD']}"
f"@{env['DB_HOST']}:{env['DB_PORT']}/{env['DB_NAME']}"
)
def build_engine() -> Engine:
"""A connection-pooled engine to the target DB (DB_* creds)."""
return create_engine(
db_url(), pool_pre_ping=True, connect_args={"connect_timeout": 10}
)
def s3_parquet_reader(bucket: str) -> ParquetReader:
"""A `ParquetReader` (key -> DataFrame) backed by `bucket` in S3, for the
`GeospatialS3Repository`. AWS creds come from the ambient `~/.aws` profile;
pyarrow reads the parquet bytes (s3fs is not installed here)."""
# boto3 ships only partial type stubs, so the client is an untyped boundary.
client = cast(Any, boto3.client("s3")) # pyright: ignore[reportUnknownMemberType]
def read(key: str) -> pd.DataFrame:
body = cast(bytes, client.get_object(Bucket=bucket, Key=key)["Body"].read())
return pd.read_parquet(io.BytesIO(body))
return read