From e92191e3d62b1bd50c12483b0992186a4f204f64 Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Tue, 25 Aug 2026 15:53:59 +0530 Subject: [PATCH 01/11] Initial changes: Add KPO compute for python models --- CHANGELOG.md | 4 + airflow/v2_10/etl_helper.py | 658 ++++++++++++------ airflow/v2_10/kpo_factory.py | 195 ++++++ airflow/v2_10/kpo_sizes.yml | 19 + airflow/v2_7/etl_helper.py | 658 ++++++++++++------ airflow/v2_7/kpo_factory.py | 195 ++++++ airflow/v2_7/kpo_sizes.yml | 19 + package-lock.json | 2 +- schemas/python-model.schema.json | 12 + .../__tests__/python-model-roundtrip.test.ts | 16 + src/services/framework/index.ts | 92 ++- .../utils/__tests__/kpo-sizes.test.ts | 23 + .../__tests__/python-model-config.test.ts | 12 + src/services/framework/utils/kpo-sizes.ts | 48 ++ .../framework/utils/python-model-utils.ts | 23 + src/shared/framework/types.ts | 25 + .../schema/types/python-model.schema.d.ts | 8 + .../skills/dj-create-python-model/_SKILL.md | 21 +- .../skills/dj-review-python-model/_SKILL.md | 9 +- .../references/review-checklist.md | 13 +- web/src/context/app.tsx | 20 + .../__tests__/formatKpoSizeInfo.test.ts | 34 + .../features/PythonModel/formatKpoSizeInfo.ts | 8 + web/src/pages/PythonModelCreate.tsx | 90 ++- 24 files changed, 1717 insertions(+), 487 deletions(-) create mode 100644 airflow/v2_10/kpo_factory.py create mode 100644 airflow/v2_10/kpo_sizes.yml create mode 100644 airflow/v2_7/kpo_factory.py create mode 100644 airflow/v2_7/kpo_sizes.yml create mode 100644 src/services/framework/utils/__tests__/kpo-sizes.test.ts create mode 100644 src/services/framework/utils/__tests__/python-model-config.test.ts create mode 100644 src/services/framework/utils/kpo-sizes.ts create mode 100644 web/src/features/PythonModel/__tests__/formatKpoSizeInfo.test.ts create mode 100644 web/src/features/PythonModel/formatKpoSizeInfo.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index db1899b..5b15125 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,10 @@ - **Fix: Model Lineage opens the correct view on first use.** Opening Model Lineage from a model `.sql`, `.model.json`, or `.yml` file (context menu, editor title, or `Cmd+Shift+L`) now shows the model lineage graph in Data Explorer instead of Column Lineage the first time the panel opens. +### Python models + +- **KPO compute is the default for new Python models.** Create Python Model lets you pick **Compute** (`kpo` or legacy `airflow` worker) and, for KPO, a **Pod size** (`small` through `xlarge`) with vCPU, memory, and workload guidance from `dags/_ext_/kpo_sizes.yml`. DJ ships `kpo_factory.py` and that YAML alongside `etl_helper.py`; set Airflow Variables `airflow_runner_cfg` (image/cluster) and `python_models_s3_bucket` (KPO S3 sync). Scaffolded `.python.py` files include `if __name__ == "__main__"`; Airflow and KPO both execute the script as `__main__` (not a fixed `run_etl` import). Until your KPO pipeline is ready, set `"compute": "airflow"` on existing models. + ### Agent skills - **New `dj-migrate-notebook-to-pymodel`, `dj-verify-pymodel-parity`, and `dj-document-pymodels` skills.** Migrate a legacy Jupyter notebook into a python model with migration plan, generate Trino SQL to verify a python model's output table matches a legacy table, and generate/refresh a topic-level README for a `python_models///` folder — rounding out the python model lifecycle alongside `dj-create-python-model` and `dj-review-python-model`. diff --git a/airflow/v2_10/etl_helper.py b/airflow/v2_10/etl_helper.py index 057a2fa..2bd749e 100644 --- a/airflow/v2_10/etl_helper.py +++ b/airflow/v2_10/etl_helper.py @@ -1,9 +1,10 @@ from __future__ import annotations -import importlib.util import json import logging +import os import re +import runpy import time from collections import defaultdict from datetime import datetime, timedelta @@ -11,13 +12,13 @@ from airflow.exceptions import AirflowSkipException from airflow.models import Variable +from airflow.models.dag import DagContext from airflow.operators.python import PythonOperator from airflow.utils.email import send_email from airflow.utils.task_group import TaskGroup from _ext_.services import trino_run -# from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator -# from kubernetes.client import models as k8s +from _ext_.kpo_factory import build_kpo_task log = logging.getLogger(__name__) @@ -152,7 +153,7 @@ def _build_task_model_registry( dag_id: str, python_source_config: dict | None = None, ) -> dict[str, dict]: - """Map Airflow task_id to topic and model_ids for mapped python model tasks.""" + """Map Airflow task_id to topic and model_ids for python model tasks.""" if python_source_config is None: python_source_config = get_python_source_config() @@ -166,11 +167,27 @@ def _build_task_model_registry( if python_source_config.get(topic, {}).get("skip", False): continue for model_tasks in _compute_model_tasks(topic_models): - task_name = "_".join(m["model_id"] for m in model_tasks) - registry[f"{topic}.{task_name}"] = { - "topic": topic, - "models": [m["model_id"] for m in model_tasks], - } + kpo_models = [m for m in model_tasks if _model_uses_kpo(m)] + airflow_models = [m for m in model_tasks if not _model_uses_kpo(m)] + + if kpo_models: + for m in kpo_models: + model_id = m["model_id"] + registry[f"{topic}.kpo__{model_id}"] = { + "topic": topic, + "models": [model_id], + } + registry[f"{topic}.gate__{model_id}"] = { + "topic": topic, + "models": [model_id], + } + + if airflow_models: + task_name = "_".join(m["model_id"] for m in airflow_models) + registry[f"{topic}.{task_name}"] = { + "topic": topic, + "models": [m["model_id"] for m in airflow_models], + } return registry @@ -516,7 +533,7 @@ def send_python_source_failure_notification( def discover_models(dag_id: str | None = None) -> list[dict]: - """Scan python/ for .python.json configs with run_etl() companion .py files. + """Scan python_models for .python.json configs with companion .python.py files. If dag_id is provided, only return models whose 'dags' field includes that DAG. Models with no dags (utility modules) are always skipped. @@ -545,18 +562,16 @@ def discover_models(dag_id: str | None = None) -> list[dict]: print(f"Skipping {json_file.name} (no companion .python.py)") continue - text = py_file.read_text() - if "def run_etl(" not in text: - print(f"Skipping {py_file.name} (no run_etl function)") - continue - model: dict = { "model_id": config.get("name", json_file.stem), "model_path": str(py_file), + "script_rel": str(py_file.relative_to(PYTHON_DIR)), "model_type": "python", "depends_on": config.get("depends_on", []), "task_group": config.get("task_group"), "topic": config.get("topic", ""), + "compute": config.get("compute", "kpo"), + "kpo_size": config.get("kpo_size", "small"), } if config.get("output"): @@ -572,7 +587,7 @@ def discover_models(dag_id: str | None = None) -> list[dict]: def execute_model(model: dict, context: dict) -> None: - """Dynamically import a model file and call its run_etl(context).""" + """Run a model .python.py file as __main__ with PYMODEL_* env from context.""" model_path = model["model_path"] model_id = model["model_id"] @@ -581,14 +596,23 @@ def execute_model(model: dict, context: dict) -> None: print(f"Executing model: {model_id}") - spec = importlib.util.spec_from_file_location(model_id, model_path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - - if not hasattr(module, "run_etl"): - raise AttributeError(f"Model missing run_etl() function: {model_path}") + pymodel_keys = ("PYMODEL_DS", "PYMODEL_DS_NODASH", "PYMODEL_DATES") + prev = {k: os.environ.get(k) for k in pymodel_keys} + try: + os.environ["PYMODEL_DS"] = context["ds"] + os.environ["PYMODEL_DS_NODASH"] = context["ds_nodash"] + if "dates" in context: + os.environ["PYMODEL_DATES"] = json.dumps(context["dates"]) + else: + os.environ.pop("PYMODEL_DATES", None) + runpy.run_path(model_path, run_name="__main__") + finally: + for key, value in prev.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value - module.run_etl(context) print(f"Completed model: {model_id}") @@ -605,32 +629,79 @@ def register_python_model_tasks(dag_id: str, dag=None): if not models: return None, None - tasks = {} + if dag is None: + dag = _current_dag() + + gate_tasks: dict[str, PythonOperator] = {} + kpo_tasks: dict[str, object] = {} + airflow_tasks: dict[str, PythonOperator] = {} + for model in models: - build_kw = dict( - task_id=f"python_model__{model['model_id']}", - python_callable=execute_model, - op_kwargs={ - "model": model, - "context": {"ds": "{{ ds }}", "ds_nodash": "{{ ds_nodash }}"}, - }, - ) - if dag is not None: - build_kw["dag"] = dag - task = PythonOperator(**build_kw) - tasks[model["model_id"]] = task + model_id = model["model_id"] + if _model_uses_kpo(model): + gate_task_id = f"gate__{model_id}" + gate_tasks[model_id] = PythonOperator( + task_id=gate_task_id, + python_callable=gate_model_run, + op_kwargs={"model": model}, + dag=dag, + on_failure_callback=_record_airflow_task_failure, + ) + kpo_tasks[model_id] = build_kpo_task( + task_id=f"kpo__{model_id}", + dag=dag, + python_model_name=model["script_rel"], + env_vars=build_kpo_env_vars(gate_task_id), + script_args=["--run-id", "{{ dag_run.run_id }}", "--logical-date", "{{ ds }}"], + size=model.get("kpo_size", "small"), + labels={"component": "python-model", "variant": model.get("topic", "")}, + on_success_callback=_record_kpo_success, + on_failure_callback=_record_kpo_failure, + ) + gate_tasks[model_id] >> kpo_tasks[model_id] + else: + airflow_tasks[model_id] = PythonOperator( + task_id=f"python_model__{model_id}", + python_callable=_execute_model_with_context, + op_kwargs={ + "model": model, + "context": {"ds": "{{ ds }}", "ds_nodash": "{{ ds_nodash }}"}, + }, + dag=dag, + ) - has_upstream = set() - is_depended_on = set() + all_ids = {m["model_id"] for m in models} for model in models: + model_id = model["model_id"] for dep_name in model.get("depends_on", []): - if dep_name in tasks: - tasks[dep_name] >> tasks[model["model_id"]] - has_upstream.add(model["model_id"]) + if dep_name not in all_ids: + continue + if dep_name in kpo_tasks and model_id in gate_tasks: + kpo_tasks[dep_name] >> gate_tasks[model_id] + elif dep_name in airflow_tasks and model_id in airflow_tasks: + airflow_tasks[dep_name] >> airflow_tasks[model_id] + + has_upstream: set[str] = set() + is_depended_on: set[str] = set() + for model in models: + model_id = model["model_id"] + for dep_name in model.get("depends_on", []): + if dep_name in all_ids: + has_upstream.add(model_id) is_depended_on.add(dep_name) - entry_tasks = [t for name, t in tasks.items() if name not in has_upstream] - exit_tasks = [t for name, t in tasks.items() if name not in is_depended_on] + entry_tasks = [] + exit_tasks = [] + for model in models: + model_id = model["model_id"] + if model_id in gate_tasks and model_id not in has_upstream: + entry_tasks.append(gate_tasks[model_id]) + if model_id in airflow_tasks and model_id not in has_upstream: + entry_tasks.append(airflow_tasks[model_id]) + if model_id in kpo_tasks and model_id not in is_depended_on: + exit_tasks.append(kpo_tasks[model_id]) + if model_id in airflow_tasks and model_id not in is_depended_on: + exit_tasks.append(airflow_tasks[model_id]) return entry_tasks, exit_tasks @@ -752,6 +823,283 @@ def resolve_dates_in(dates_in: str | list | dict, ds: str) -> list[str]: return dates +def _model_uses_kpo(model: dict) -> bool: + return model.get("compute", "kpo") != "airflow" + + +def _current_dag(): + dag = DagContext.get_current_dag() + if dag is None: + raise RuntimeError( + "Python model task registration must run inside an active DAG context" + ) + return dag + + +def _resolve_model_run( + model: dict, + topic: str, + context: dict, + *, + record_fn=None, +) -> tuple[dict, str, str] | None: + """Resolve partition context for a model run. + + Returns ``(ctx, etl_date, model_group)`` or ``None`` when ``skip_model`` + applies. Raises ``AirflowSkipException`` for topic-level skip / schedule gate. + """ + model_name = model["model_id"] + model_group = topic or model.get("topic", "") + dates_in: str | list | dict = "" + + if topic: + conf = (context.get("dag_run").conf or {}) if context.get("dag_run") else {} + runtime_config = conf.get("python_source_config") or get_python_source_config() + topic_cfg = runtime_config.get(topic, {}) + + if topic_cfg.get("skip", False): + if record_fn: + record_fn("skipped", 0.0, True) + raise AirflowSkipException(f"topic {topic} skip=True") + + date_to_run = topic_cfg.get("date_to_run", "") + if not passes_schedule_gate(date_to_run, context["ds"]): + if record_fn: + record_fn("skipped", 0.0, True) + raise AirflowSkipException( + f"topic {topic} schedule gate '{date_to_run}' not met" + ) + + skip_models = topic_cfg.get("skip_model", []) + if isinstance(skip_models, str): + skip_models = [skip_models] + skip_models = [m for m in skip_models if m and str(m).strip()] + if model_name in skip_models: + if record_fn: + record_fn("skipped", 0.0, True) + return None + + dates_in = topic_cfg.get("dates_in", "") + + dates = resolve_dates_in(dates_in, context["ds"]) + etl_date = dates[0] + ctx = {"ds": etl_date, "ds_nodash": etl_date.replace("-", "")} + if dates_in: + ctx["dates"] = dates + return ctx, etl_date, model_group + + +def gate_model_run(model: dict, topic: str = "", **context) -> None: + """Resolve gating and push run context to XCom for a downstream KPO task.""" + model_name = model["model_id"] + model_group = topic or model.get("topic", "") + etl_timestamp = _get_etl_timestamp_from_context(context) + etl_date = context["ds"] + + def _record( + run_status: str, + run_seconds: float, + is_skipped: bool, + run_message: str | None = None, + record_etl_date: str | None = None, + ) -> None: + try: + record_python_model_run( + model_name=model_name, + model_group=model_group, + etl_date=record_etl_date or etl_date, + etl_timestamp=etl_timestamp or "", + run_status=run_status, + run_seconds=run_seconds, + is_skipped=is_skipped, + run_message=run_message, + context=context, + ) + except Exception as exc: + log.warning("Failed to record python model run for %s: %s", model_name, exc) + + resolved = _resolve_model_run(model, topic, context, record_fn=_record) + if resolved is None: + return + + ctx, resolved_etl_date, resolved_group = resolved + ti = context["ti"] + ti.xcom_push(key="ds", value=ctx["ds"]) + ti.xcom_push(key="ds_nodash", value=ctx["ds_nodash"]) + if "dates" in ctx: + ti.xcom_push(key="dates_json", value=json.dumps(ctx["dates"])) + ti.xcom_push(key="etl_date", value=resolved_etl_date) + ti.xcom_push(key="model_group", value=resolved_group) + + +def build_kpo_env_vars(gate_task_id: str) -> dict[str, str]: + """Build container env vars for a KPO task, pulling partition dates from gate XCom.""" + env: dict[str, str] = { + "PYMODEL_DS": f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='ds') }}}}", + "PYMODEL_DS_NODASH": ( + f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='ds_nodash') }}}}" + ), + "DJ_PYTHON_MODEL_CATALOG_NAME": ( + "{{ var.value.get('dj_python_model_catalog_name', 'glue') }}" + ), + "python_models_s3_bucket": ( + "{{ var.value.get('python_models_s3_bucket') }}" + ), + "python_model_s3_region": ( + "{{ var.value.get('python_model_s3_region', 'us-west-2') }}" + ), + "OPUS-ETL-CLUSTER-HOST": "{{ var.value.get('opus-etl-cluster-host') }}", + "OPUS-ETL-USER-NAME": "{{ var.value.get('opus-etl-user-name') }}", + "OPUS-ETL-CLUSTER-PASSWORD": ( + "{{ var.value.get('opus-etl-cluster-password') }}" + ), + } + env["PYMODEL_DATES"] = ( + f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='dates_json') }}}}" + ) + return env + + +def _kpo_model_id_from_task_id(task_id: str) -> str | None: + marker = "kpo__" + if marker not in task_id: + return None + return task_id.rsplit(marker, 1)[-1] + + +def _gate_task_id_for_kpo(task_id: str) -> str: + if "kpo__" in task_id: + return task_id.replace("kpo__", "gate__", 1) + return task_id + + +def _record_kpo_success(context: dict) -> None: + _record_kpo_outcome(context, run_status="success") + + +def _record_kpo_failure(context: dict) -> None: + ti = context.get("ti") + run_message = None + if ti is not None: + run_message = getattr(ti, "note", None) or "KPO task failed" + _record_kpo_outcome(context, run_status="error", run_message=run_message) + + +def _record_kpo_outcome( + context: dict, + run_status: str, + run_message: str | None = None, +) -> None: + try: + ti = context["ti"] + model_id = _kpo_model_id_from_task_id(ti.task_id) + if not model_id: + return + + gate_task_id = _gate_task_id_for_kpo(ti.task_id) + etl_date = ti.xcom_pull(task_ids=gate_task_id, key="etl_date") or context["ds"] + model_group = ti.xcom_pull(task_ids=gate_task_id, key="model_group") or "" + etl_timestamp = _get_etl_timestamp_from_context(context) + + started_at = ti.start_date + run_seconds = 0.0 + if started_at is not None and ti.end_date is not None: + run_seconds = (ti.end_date - started_at).total_seconds() + + record_python_model_run( + model_name=model_id, + model_group=model_group, + etl_date=etl_date, + etl_timestamp=etl_timestamp or "", + run_status=run_status, + run_seconds=run_seconds, + is_skipped=False, + run_message=run_message, + context=context, + ) + except Exception as exc: + log.warning("Failed to record KPO model run: %s", exc) + + +def _wire_kpo_batch_depends_on( + model_tasks: list[dict], + gate_tasks: dict[str, PythonOperator], + kpo_tasks: dict[str, object], +) -> tuple[list[PythonOperator], list]: + """Wire depends_on across gate/KPO pairs; return entry gates and exit KPO tasks.""" + model_ids = {m["model_id"] for m in model_tasks} + for m in model_tasks: + model_id = m["model_id"] + for dep_id in m.get("depends_on", []): + if dep_id in model_ids and dep_id in kpo_tasks: + kpo_tasks[dep_id] >> gate_tasks[model_id] + + has_upstream: set[str] = set() + is_depended_on: set[str] = set() + for m in model_tasks: + model_id = m["model_id"] + for dep_id in m.get("depends_on", []): + if dep_id in model_ids: + has_upstream.add(model_id) + is_depended_on.add(dep_id) + + entry_gates = [ + gate_tasks[m["model_id"]] + for m in model_tasks + if m["model_id"] not in has_upstream + ] + exit_kpos = [ + kpo_tasks[m["model_id"]] + for m in model_tasks + if m["model_id"] not in is_depended_on + ] + return entry_gates, exit_kpos + + +def _register_kpo_model_batch( + dag, + topic: str, + model_tasks: list[dict], + prev_task, +) -> object: + """Create gate + KPO task pairs for a homogeneous KPO model batch.""" + gate_tasks: dict[str, PythonOperator] = {} + kpo_tasks: dict[str, object] = {} + + for m in model_tasks: + model_id = m["model_id"] + gate_task_id = f"gate__{model_id}" + gate_tasks[model_id] = PythonOperator( + task_id=gate_task_id, + python_callable=gate_model_run, + op_kwargs={"model": m, "topic": topic}, + dag=dag, + on_failure_callback=_record_airflow_task_failure, + ) + kpo_tasks[model_id] = build_kpo_task( + task_id=f"kpo__{model_id}", + dag=dag, + python_model_name=m["script_rel"], + env_vars=build_kpo_env_vars(gate_task_id), + script_args=["--run-id", "{{ dag_run.run_id }}", "--logical-date", "{{ ds }}"], + size=m.get("kpo_size", "small"), + labels={"component": "python-model", "variant": topic}, + on_success_callback=_record_kpo_success, + on_failure_callback=_record_kpo_failure, + ) + gate_tasks[model_id] >> kpo_tasks[model_id] + + entry_gates, exit_kpos = _wire_kpo_batch_depends_on( + model_tasks, gate_tasks, kpo_tasks + ) + if prev_task is not None: + prev_task >> entry_gates + + if len(exit_kpos) == 1: + return exit_kpos[0] + return exit_kpos + + def _execute_model_with_context( model: dict, topic: str = "", @@ -776,7 +1124,6 @@ def _execute_model_with_context( model_group = topic or model.get("topic", "") etl_timestamp = _get_etl_timestamp_from_context(context) etl_date = context["ds"] - dates_in: str | list | dict = "" started_at: float | None = None def _record( @@ -802,56 +1149,11 @@ def _record( log.warning("Failed to record python model run for %s: %s", model_name, exc) try: - if topic: - conf = (context.get("dag_run").conf or {}) if context.get("dag_run") else {} - runtime_config = conf.get("python_source_config") or get_python_source_config() - topic_cfg = runtime_config.get(topic, {}) - - if topic_cfg.get("skip", False): - log.info( - "Skipping model %s (topic %s skip=True at runtime)", - model_name, - topic, - ) - _record("skipped", 0.0, True) - raise AirflowSkipException(f"topic {topic} skip=True") - - date_to_run = topic_cfg.get("date_to_run", "") - if not passes_schedule_gate(date_to_run, context["ds"]): - log.info( - "Skipping model %s (topic %s schedule gate '%s' not met for ds %s)", - model_name, - topic, - date_to_run, - context["ds"], - ) - _record("skipped", 0.0, True) - raise AirflowSkipException( - f"topic {topic} schedule gate '{date_to_run}' not met" - ) - - skip_models = topic_cfg.get("skip_model", []) - if isinstance(skip_models, str): - skip_models = [skip_models] - skip_models = [m for m in skip_models if m and str(m).strip()] - if model_name in skip_models: - log.info( - "Skipping model %s (topic %s skip_model=%s)", - model_name, - topic, - skip_models, - ) - _record("skipped", 0.0, True) - return - - dates_in = topic_cfg.get("dates_in", "") - - dates = resolve_dates_in(dates_in, context["ds"]) - etl_date = dates[0] - ctx = {"ds": etl_date, "ds_nodash": etl_date.replace("-", "")} - if dates_in: - ctx["dates"] = dates + resolved = _resolve_model_run(model, topic, context, record_fn=_record) + if resolved is None: + return + ctx, etl_date, _ = resolved started_at = time.perf_counter() execute_model(model, ctx) _record("success", time.perf_counter() - started_at, False) @@ -863,6 +1165,43 @@ def _record( raise +def _register_model_task_batch( + dag, + topic: str, + model_tasks: list[dict], + prev_task, +) -> object: + """Register a topological batch using KPO or Airflow worker compute per model.""" + kpo_models = [m for m in model_tasks if _model_uses_kpo(m)] + airflow_models = [m for m in model_tasks if not _model_uses_kpo(m)] + batch_exit = prev_task + + if kpo_models: + batch_exit = _register_kpo_model_batch(dag, topic, kpo_models, batch_exit) + + if airflow_models: + task_name = "_".join(m["model_id"] for m in airflow_models) + op_kwargs_list = [ + { + "model": m, + "model_id": m["model_id"], + **({"topic": topic} if topic else {}), + } + for m in airflow_models + ] + mapped_task = PythonOperator.partial( + task_id=task_name, + python_callable=_execute_model_with_context, + dag=dag, + on_failure_callback=_record_airflow_task_failure, + ).expand(op_kwargs=op_kwargs_list) + if batch_exit is not None: + batch_exit >> mapped_task + batch_exit = mapped_task + + return batch_exit + + def register_python_model_mapped_tasks(dag_id: str, group_id: str | None = None): """Create a TaskGroup with model-task-batch mapped tasks for python models. @@ -883,20 +1222,12 @@ def register_python_model_mapped_tasks(dag_id: str, group_id: str | None = None) return None model_task_batches = _compute_model_tasks(models) + dag = _current_dag() with TaskGroup(group_id=group_id or dag_id) as tg: prev_task = None for model_tasks in model_task_batches: - task_name = "_".join(m["model_id"] for m in model_tasks) - mapped_task = PythonOperator.partial( - task_id=task_name, - python_callable=_execute_model_with_context, - on_failure_callback=_record_airflow_task_failure, - ).expand(op_kwargs=[{"model": m, "model_id": m["model_id"]} for m in model_tasks]) - - if prev_task is not None: - prev_task >> mapped_task - prev_task = mapped_task + prev_task = _register_model_task_batch(dag, "", model_tasks, prev_task) return tg @@ -940,128 +1271,15 @@ def register_python_model_mapped_tasks_by_topic( continue model_task_batches = _compute_model_tasks(topic_models) + dag = _current_dag() with TaskGroup(group_id=topic) as tg: prev_task = None for model_tasks in model_task_batches: - task_name = "_".join(m["model_id"] for m in model_tasks) - - op_kwargs_list = [ - { - "model": m, - "model_id": m["model_id"], - "topic": topic, - } - for m in model_tasks - ] - - mapped_task = PythonOperator.partial( - task_id=task_name, - python_callable=_execute_model_with_context, - on_failure_callback=_record_airflow_task_failure, - ).expand(op_kwargs=op_kwargs_list) - - if prev_task is not None: - prev_task >> mapped_task - prev_task = mapped_task + prev_task = _register_model_task_batch( + dag, topic, model_tasks, prev_task + ) topic_groups.append(tg) return topic_groups if topic_groups else None - - -# def register_k8s_model_tasks( -# dag_id: str, -# image: str, -# k8s_config_path: str, -# cluster_context: str, -# namespace: str = "mwaa", -# image_pull_secret: str = "wd-docker-artifactory-cred", -# dag=None, -# ) -> tuple: -# """Create KubernetesPodOperator tasks for each Python model and wire depends_on. -# -# Each pod runs dags/python_models/entrypoint.py inside the pymodels image. -# Airflow context (ds) and Trino credentials are injected as env vars via -# Jinja templates so values are resolved at execution time, not parse time. -# -# Returns (entry_tasks, exit_tasks) for chaining into the parent DAG. -# Returns (None, None) if no models are found. -# """ -# models = discover_models(dag_id) -# if not models: -# return None, None -# -# # Group models by task_group; models without one form solo groups keyed by model_id. -# groups: dict = {} -# for m in models: -# key = m.get("task_group") or m["model_id"] -# groups.setdefault(key, []).append(m) -# -# # model_id → group_key for cross-group dependency wiring -# model_to_group: dict = {m["model_id"]: (m.get("task_group") or m["model_id"]) for m in models} -# -# _shared_env = [ -# k8s.V1EnvVar(name="PYMODEL_DS", value="{{ ds }}"), -# k8s.V1EnvVar( -# name="OPUS-ETL-CLUSTER-HOST", -# value="{{ var.value.get('opus-etl-cluster-host') }}", -# ), -# k8s.V1EnvVar( -# name="OPUS-ETL-USER-NAME", -# value="{{ var.value.get('opus-etl-user-name') }}", -# ), -# k8s.V1EnvVar( -# name="OPUS-ETL-CLUSTER-PASSWORD", -# value="{{ var.value.get('opus-etl-cluster-password') }}", -# ), -# k8s.V1EnvVar( -# name="DJ_PYTHON_MODEL_CATALOG_NAME", -# value="{{ var.value.get('dj_python_model_catalog_name', 'glue') }}", -# ), -# ] -# -# tasks: dict = {} -# for group_key, group_models in groups.items(): -# if len(group_models) == 1: -# m = group_models[0] -# task_id = f"python_model__{m['model_id']}" -# group_env = [k8s.V1EnvVar(name="PYMODEL_MODEL_ID", value=m["model_id"])] -# else: -# task_id = f"python_model_group__{group_key}" -# ids = ",".join(m["model_id"] for m in group_models) -# group_env = [k8s.V1EnvVar(name="PYMODEL_MODEL_IDS", value=ids)] -# -# build_kw = dict( -# task_id=task_id, -# namespace=namespace, -# image=image, -# image_pull_secrets=[k8s.V1LocalObjectReference(name=image_pull_secret)], -# cmds=["python3.11", "/opt/py_models/entrypoint.py"], -# env_vars=_shared_env + group_env, -# is_delete_operator_pod=True, -# get_logs=True, -# config_file=k8s_config_path, -# in_cluster=False, -# cluster_context=cluster_context, -# startup_timeout_seconds=300, -# ) -# if dag is not None: -# build_kw["dag"] = dag -# tasks[group_key] = KubernetesPodOperator(**build_kw) -# -# # Wire cross-group Airflow dependencies. -# has_upstream: set = set() -# is_depended_on: set = set() -# for m in models: -# group_key = model_to_group[m["model_id"]] -# for dep_model_id in m.get("depends_on", []): -# dep_group = model_to_group.get(dep_model_id) -# if dep_group and dep_group != group_key and dep_group in tasks: -# tasks[dep_group] >> tasks[group_key] -# has_upstream.add(group_key) -# is_depended_on.add(dep_group) -# -# entry_tasks = [t for k, t in tasks.items() if k not in has_upstream] -# exit_tasks = [t for k, t in tasks.items() if k not in is_depended_on] -# return entry_tasks, exit_tasks diff --git a/airflow/v2_10/kpo_factory.py b/airflow/v2_10/kpo_factory.py new file mode 100644 index 0000000..ca561a7 --- /dev/null +++ b/airflow/v2_10/kpo_factory.py @@ -0,0 +1,195 @@ +"""Reusable KubernetesPodOperator factory for DJ python model execution. + +Target: Airflow 2.10 with ``apache-airflow-providers-cncf-kubernetes`` 8.x +(uses ``on_finish_action``, not ``is_delete_operator_pod``). + +Cluster/image config is resolved at task-execution time from the shared +Airflow Variable ``airflow_runner_cfg``. Pod sizing presets are loaded from +``kpo_sizes.yml`` beside this module (optional ``size_specs`` override in +the Variable JSON). +""" + +from __future__ import annotations + +from datetime import timedelta +from functools import lru_cache +from pathlib import Path +from typing import Literal, Mapping, Sequence + +from airflow import DAG +from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator +from kubernetes.client import models as k8s + +_DEFAULT_IMAGE_PULL_SECRETS: tuple[str, ...] = ("wd-docker-artifactory-cred",) + +_RUNNER_CFG_VAR_NAME: str = "airflow_runner_cfg" +_IMAGE_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.image }}}}" +_TAG_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.tag }}}}" +_NAMESPACE_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.namespace }}}}" +_CONFIG_FILE_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.config_file }}}}" +_CLUSTER_CONTEXT_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.cluster_context }}}}" + +_IMAGE_REF: str = f"{_IMAGE_TPL}:{_TAG_TPL}" + +PodSize = Literal["small", "medium", "large", "xlarge"] + +_DEFAULT_SIZE_SPECS: dict[str, tuple[str, str]] = { + "small": ("1", "2Gi"), + "medium": ("2", "4Gi"), + "large": ("2", "8Gi"), + "xlarge": ("2", "16Gi"), +} + +_SIZES_YAML_PATH = Path(__file__).resolve().parent / "kpo_sizes.yml" + + +def _parse_size_specs_from_yaml(path: Path) -> dict[str, tuple[str, str]]: + if not path.is_file(): + return dict(_DEFAULT_SIZE_SPECS) + + try: + import yaml + except ImportError: + return dict(_DEFAULT_SIZE_SPECS) + + try: + raw = yaml.safe_load(path.read_text()) or {} + except Exception: + return dict(_DEFAULT_SIZE_SPECS) + + sizes = raw.get("sizes") or {} + parsed: dict[str, tuple[str, str]] = {} + for name, spec in sizes.items(): + if not isinstance(spec, dict): + continue + cpu = spec.get("cpu") + memory = spec.get("memory") + if cpu and memory: + parsed[str(name)] = (str(cpu), str(memory)) + return parsed if parsed else dict(_DEFAULT_SIZE_SPECS) + + +def _merge_variable_size_specs(specs: dict[str, tuple[str, str]]) -> dict[str, tuple[str, str]]: + try: + from airflow.models import Variable + + cfg = Variable.get(_RUNNER_CFG_VAR_NAME, deserialize_json=True, default_var={}) + if not isinstance(cfg, dict): + return specs + overrides = cfg.get("size_specs") or {} + if not isinstance(overrides, dict): + return specs + merged = dict(specs) + for name, spec in overrides.items(): + if isinstance(spec, dict) and spec.get("cpu") and spec.get("memory"): + merged[str(name)] = (str(spec["cpu"]), str(spec["memory"])) + return merged + except Exception: + return specs + + +@lru_cache(maxsize=1) +def load_size_specs() -> dict[str, tuple[str, str]]: + """Load pod CPU/memory presets from kpo_sizes.yml with optional Variable override.""" + return _merge_variable_size_specs(_parse_size_specs_from_yaml(_SIZES_YAML_PATH)) + + +def _sanitize_pod_name(dag_id: str, task_id: str) -> str: + raw = f"{dag_id}-{task_id}".replace("_", "-").lower() + return raw[:63].rstrip("-") + + +def _container_resources(size: PodSize) -> k8s.V1ResourceRequirements: + specs = load_size_specs() + if size not in specs: + raise ValueError( + f"Invalid pod size {size!r}; expected one of {sorted(specs)}." + ) + cpu, memory = specs[size] + return k8s.V1ResourceRequirements( + requests={"cpu": cpu, "memory": memory}, + limits={"cpu": cpu, "memory": memory}, + ) + + +def _env_var_list(env: Mapping[str, str]) -> list[k8s.V1EnvVar]: + return [k8s.V1EnvVar(name=name, value=value) for name, value in env.items()] + + +def build_kpo_task( + *, + task_id: str, + dag: DAG, + python_model_name: str, + env_vars: Mapping[str, str], + script_args: Sequence[str] | None = None, + size: PodSize = "small", + image_pull_policy: str = "IfNotPresent", + image_pull_secrets: Sequence[str] = _DEFAULT_IMAGE_PULL_SECRETS, + service_account_name: str = "default", + retries: int = 1, + retry_delay: timedelta = timedelta(minutes=5), + retry_exponential_backoff: bool = True, + startup_timeout_seconds: int = 600, + on_finish_action: str = "delete_pod", + deferrable: bool = False, + logging_interval: int | None = 10, + do_xcom_push: bool = False, + labels: Mapping[str, str] | None = None, + annotations: Mapping[str, str] | None = None, + on_success_callback=None, + on_failure_callback=None, +) -> KubernetesPodOperator: + """Build a KubernetesPodOperator that invokes the image ENTRYPOINT.""" + env_map: dict[str, str] = {"PYTHONUNBUFFERED": "1"} + env_map.update(env_vars) + + dag_labels: dict[str, str] = { + "dag_id": dag.dag_id, + "task_id": task_id, + "run_id": "{{ dag_run.run_id | truncate(63, true, '') }}", + "component": "airflow-kpo", + "variant": dag.dag_id, + } + if labels: + dag_labels.update(labels) + + arguments = [python_model_name, *(list(script_args) if script_args else [])] + + deferred_log_kwargs: dict[str, int] = {} + if deferrable and logging_interval is not None: + deferred_log_kwargs["logging_interval"] = logging_interval + + return KubernetesPodOperator( + task_id=task_id, + dag=dag, + image=_IMAGE_REF, + image_pull_policy=image_pull_policy, + image_pull_secrets=[ + k8s.V1LocalObjectReference(name=n) for n in image_pull_secrets + ], + in_cluster=False, + config_file=_CONFIG_FILE_TPL, + cluster_context=_CLUSTER_CONTEXT_TPL, + namespace=_NAMESPACE_TPL, + name=_sanitize_pod_name(dag.dag_id, task_id), + arguments=arguments, + env_vars=_env_var_list(env_map), + container_resources=_container_resources(size), + service_account_name=service_account_name, + get_logs=True, + log_events_on_failure=True, + on_finish_action=on_finish_action, + do_xcom_push=do_xcom_push, + deferrable=deferrable, + **deferred_log_kwargs, + retries=retries, + retry_delay=retry_delay, + retry_exponential_backoff=retry_exponential_backoff, + startup_timeout_seconds=startup_timeout_seconds, + reattach_on_restart=True, + labels=dag_labels, + annotations=dict(annotations or {}), + on_success_callback=on_success_callback, + on_failure_callback=on_failure_callback, + ) diff --git a/airflow/v2_10/kpo_sizes.yml b/airflow/v2_10/kpo_sizes.yml new file mode 100644 index 0000000..868dd7f --- /dev/null +++ b/airflow/v2_10/kpo_sizes.yml @@ -0,0 +1,19 @@ +# Pod resource presets for KPO python models. +# CPU requests == limits (Guaranteed QoS). CPU capped at 2 on shared EKS nodes. +sizes: + small: + cpu: "1" + memory: "2Gi" + guidance: "Light jobs — small API responses, narrow CSV extracts, < ~500 MB processed in-memory." + medium: + cpu: "2" + memory: "4Gi" + guidance: "Typical ETL — moderate API pagination, joins on staging tables, ~500 MB–2 GB." + large: + cpu: "2" + memory: "8Gi" + guidance: "Heavier transforms — wider schemas, multi-step pandas work, ~2–8 GB." + xlarge: + cpu: "2" + memory: "16Gi" + guidance: "Large workloads — full-table scans, wide fin-style exports, ~8–16 GB." diff --git a/airflow/v2_7/etl_helper.py b/airflow/v2_7/etl_helper.py index 057a2fa..2bd749e 100644 --- a/airflow/v2_7/etl_helper.py +++ b/airflow/v2_7/etl_helper.py @@ -1,9 +1,10 @@ from __future__ import annotations -import importlib.util import json import logging +import os import re +import runpy import time from collections import defaultdict from datetime import datetime, timedelta @@ -11,13 +12,13 @@ from airflow.exceptions import AirflowSkipException from airflow.models import Variable +from airflow.models.dag import DagContext from airflow.operators.python import PythonOperator from airflow.utils.email import send_email from airflow.utils.task_group import TaskGroup from _ext_.services import trino_run -# from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator -# from kubernetes.client import models as k8s +from _ext_.kpo_factory import build_kpo_task log = logging.getLogger(__name__) @@ -152,7 +153,7 @@ def _build_task_model_registry( dag_id: str, python_source_config: dict | None = None, ) -> dict[str, dict]: - """Map Airflow task_id to topic and model_ids for mapped python model tasks.""" + """Map Airflow task_id to topic and model_ids for python model tasks.""" if python_source_config is None: python_source_config = get_python_source_config() @@ -166,11 +167,27 @@ def _build_task_model_registry( if python_source_config.get(topic, {}).get("skip", False): continue for model_tasks in _compute_model_tasks(topic_models): - task_name = "_".join(m["model_id"] for m in model_tasks) - registry[f"{topic}.{task_name}"] = { - "topic": topic, - "models": [m["model_id"] for m in model_tasks], - } + kpo_models = [m for m in model_tasks if _model_uses_kpo(m)] + airflow_models = [m for m in model_tasks if not _model_uses_kpo(m)] + + if kpo_models: + for m in kpo_models: + model_id = m["model_id"] + registry[f"{topic}.kpo__{model_id}"] = { + "topic": topic, + "models": [model_id], + } + registry[f"{topic}.gate__{model_id}"] = { + "topic": topic, + "models": [model_id], + } + + if airflow_models: + task_name = "_".join(m["model_id"] for m in airflow_models) + registry[f"{topic}.{task_name}"] = { + "topic": topic, + "models": [m["model_id"] for m in airflow_models], + } return registry @@ -516,7 +533,7 @@ def send_python_source_failure_notification( def discover_models(dag_id: str | None = None) -> list[dict]: - """Scan python/ for .python.json configs with run_etl() companion .py files. + """Scan python_models for .python.json configs with companion .python.py files. If dag_id is provided, only return models whose 'dags' field includes that DAG. Models with no dags (utility modules) are always skipped. @@ -545,18 +562,16 @@ def discover_models(dag_id: str | None = None) -> list[dict]: print(f"Skipping {json_file.name} (no companion .python.py)") continue - text = py_file.read_text() - if "def run_etl(" not in text: - print(f"Skipping {py_file.name} (no run_etl function)") - continue - model: dict = { "model_id": config.get("name", json_file.stem), "model_path": str(py_file), + "script_rel": str(py_file.relative_to(PYTHON_DIR)), "model_type": "python", "depends_on": config.get("depends_on", []), "task_group": config.get("task_group"), "topic": config.get("topic", ""), + "compute": config.get("compute", "kpo"), + "kpo_size": config.get("kpo_size", "small"), } if config.get("output"): @@ -572,7 +587,7 @@ def discover_models(dag_id: str | None = None) -> list[dict]: def execute_model(model: dict, context: dict) -> None: - """Dynamically import a model file and call its run_etl(context).""" + """Run a model .python.py file as __main__ with PYMODEL_* env from context.""" model_path = model["model_path"] model_id = model["model_id"] @@ -581,14 +596,23 @@ def execute_model(model: dict, context: dict) -> None: print(f"Executing model: {model_id}") - spec = importlib.util.spec_from_file_location(model_id, model_path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - - if not hasattr(module, "run_etl"): - raise AttributeError(f"Model missing run_etl() function: {model_path}") + pymodel_keys = ("PYMODEL_DS", "PYMODEL_DS_NODASH", "PYMODEL_DATES") + prev = {k: os.environ.get(k) for k in pymodel_keys} + try: + os.environ["PYMODEL_DS"] = context["ds"] + os.environ["PYMODEL_DS_NODASH"] = context["ds_nodash"] + if "dates" in context: + os.environ["PYMODEL_DATES"] = json.dumps(context["dates"]) + else: + os.environ.pop("PYMODEL_DATES", None) + runpy.run_path(model_path, run_name="__main__") + finally: + for key, value in prev.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value - module.run_etl(context) print(f"Completed model: {model_id}") @@ -605,32 +629,79 @@ def register_python_model_tasks(dag_id: str, dag=None): if not models: return None, None - tasks = {} + if dag is None: + dag = _current_dag() + + gate_tasks: dict[str, PythonOperator] = {} + kpo_tasks: dict[str, object] = {} + airflow_tasks: dict[str, PythonOperator] = {} + for model in models: - build_kw = dict( - task_id=f"python_model__{model['model_id']}", - python_callable=execute_model, - op_kwargs={ - "model": model, - "context": {"ds": "{{ ds }}", "ds_nodash": "{{ ds_nodash }}"}, - }, - ) - if dag is not None: - build_kw["dag"] = dag - task = PythonOperator(**build_kw) - tasks[model["model_id"]] = task + model_id = model["model_id"] + if _model_uses_kpo(model): + gate_task_id = f"gate__{model_id}" + gate_tasks[model_id] = PythonOperator( + task_id=gate_task_id, + python_callable=gate_model_run, + op_kwargs={"model": model}, + dag=dag, + on_failure_callback=_record_airflow_task_failure, + ) + kpo_tasks[model_id] = build_kpo_task( + task_id=f"kpo__{model_id}", + dag=dag, + python_model_name=model["script_rel"], + env_vars=build_kpo_env_vars(gate_task_id), + script_args=["--run-id", "{{ dag_run.run_id }}", "--logical-date", "{{ ds }}"], + size=model.get("kpo_size", "small"), + labels={"component": "python-model", "variant": model.get("topic", "")}, + on_success_callback=_record_kpo_success, + on_failure_callback=_record_kpo_failure, + ) + gate_tasks[model_id] >> kpo_tasks[model_id] + else: + airflow_tasks[model_id] = PythonOperator( + task_id=f"python_model__{model_id}", + python_callable=_execute_model_with_context, + op_kwargs={ + "model": model, + "context": {"ds": "{{ ds }}", "ds_nodash": "{{ ds_nodash }}"}, + }, + dag=dag, + ) - has_upstream = set() - is_depended_on = set() + all_ids = {m["model_id"] for m in models} for model in models: + model_id = model["model_id"] for dep_name in model.get("depends_on", []): - if dep_name in tasks: - tasks[dep_name] >> tasks[model["model_id"]] - has_upstream.add(model["model_id"]) + if dep_name not in all_ids: + continue + if dep_name in kpo_tasks and model_id in gate_tasks: + kpo_tasks[dep_name] >> gate_tasks[model_id] + elif dep_name in airflow_tasks and model_id in airflow_tasks: + airflow_tasks[dep_name] >> airflow_tasks[model_id] + + has_upstream: set[str] = set() + is_depended_on: set[str] = set() + for model in models: + model_id = model["model_id"] + for dep_name in model.get("depends_on", []): + if dep_name in all_ids: + has_upstream.add(model_id) is_depended_on.add(dep_name) - entry_tasks = [t for name, t in tasks.items() if name not in has_upstream] - exit_tasks = [t for name, t in tasks.items() if name not in is_depended_on] + entry_tasks = [] + exit_tasks = [] + for model in models: + model_id = model["model_id"] + if model_id in gate_tasks and model_id not in has_upstream: + entry_tasks.append(gate_tasks[model_id]) + if model_id in airflow_tasks and model_id not in has_upstream: + entry_tasks.append(airflow_tasks[model_id]) + if model_id in kpo_tasks and model_id not in is_depended_on: + exit_tasks.append(kpo_tasks[model_id]) + if model_id in airflow_tasks and model_id not in is_depended_on: + exit_tasks.append(airflow_tasks[model_id]) return entry_tasks, exit_tasks @@ -752,6 +823,283 @@ def resolve_dates_in(dates_in: str | list | dict, ds: str) -> list[str]: return dates +def _model_uses_kpo(model: dict) -> bool: + return model.get("compute", "kpo") != "airflow" + + +def _current_dag(): + dag = DagContext.get_current_dag() + if dag is None: + raise RuntimeError( + "Python model task registration must run inside an active DAG context" + ) + return dag + + +def _resolve_model_run( + model: dict, + topic: str, + context: dict, + *, + record_fn=None, +) -> tuple[dict, str, str] | None: + """Resolve partition context for a model run. + + Returns ``(ctx, etl_date, model_group)`` or ``None`` when ``skip_model`` + applies. Raises ``AirflowSkipException`` for topic-level skip / schedule gate. + """ + model_name = model["model_id"] + model_group = topic or model.get("topic", "") + dates_in: str | list | dict = "" + + if topic: + conf = (context.get("dag_run").conf or {}) if context.get("dag_run") else {} + runtime_config = conf.get("python_source_config") or get_python_source_config() + topic_cfg = runtime_config.get(topic, {}) + + if topic_cfg.get("skip", False): + if record_fn: + record_fn("skipped", 0.0, True) + raise AirflowSkipException(f"topic {topic} skip=True") + + date_to_run = topic_cfg.get("date_to_run", "") + if not passes_schedule_gate(date_to_run, context["ds"]): + if record_fn: + record_fn("skipped", 0.0, True) + raise AirflowSkipException( + f"topic {topic} schedule gate '{date_to_run}' not met" + ) + + skip_models = topic_cfg.get("skip_model", []) + if isinstance(skip_models, str): + skip_models = [skip_models] + skip_models = [m for m in skip_models if m and str(m).strip()] + if model_name in skip_models: + if record_fn: + record_fn("skipped", 0.0, True) + return None + + dates_in = topic_cfg.get("dates_in", "") + + dates = resolve_dates_in(dates_in, context["ds"]) + etl_date = dates[0] + ctx = {"ds": etl_date, "ds_nodash": etl_date.replace("-", "")} + if dates_in: + ctx["dates"] = dates + return ctx, etl_date, model_group + + +def gate_model_run(model: dict, topic: str = "", **context) -> None: + """Resolve gating and push run context to XCom for a downstream KPO task.""" + model_name = model["model_id"] + model_group = topic or model.get("topic", "") + etl_timestamp = _get_etl_timestamp_from_context(context) + etl_date = context["ds"] + + def _record( + run_status: str, + run_seconds: float, + is_skipped: bool, + run_message: str | None = None, + record_etl_date: str | None = None, + ) -> None: + try: + record_python_model_run( + model_name=model_name, + model_group=model_group, + etl_date=record_etl_date or etl_date, + etl_timestamp=etl_timestamp or "", + run_status=run_status, + run_seconds=run_seconds, + is_skipped=is_skipped, + run_message=run_message, + context=context, + ) + except Exception as exc: + log.warning("Failed to record python model run for %s: %s", model_name, exc) + + resolved = _resolve_model_run(model, topic, context, record_fn=_record) + if resolved is None: + return + + ctx, resolved_etl_date, resolved_group = resolved + ti = context["ti"] + ti.xcom_push(key="ds", value=ctx["ds"]) + ti.xcom_push(key="ds_nodash", value=ctx["ds_nodash"]) + if "dates" in ctx: + ti.xcom_push(key="dates_json", value=json.dumps(ctx["dates"])) + ti.xcom_push(key="etl_date", value=resolved_etl_date) + ti.xcom_push(key="model_group", value=resolved_group) + + +def build_kpo_env_vars(gate_task_id: str) -> dict[str, str]: + """Build container env vars for a KPO task, pulling partition dates from gate XCom.""" + env: dict[str, str] = { + "PYMODEL_DS": f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='ds') }}}}", + "PYMODEL_DS_NODASH": ( + f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='ds_nodash') }}}}" + ), + "DJ_PYTHON_MODEL_CATALOG_NAME": ( + "{{ var.value.get('dj_python_model_catalog_name', 'glue') }}" + ), + "python_models_s3_bucket": ( + "{{ var.value.get('python_models_s3_bucket') }}" + ), + "python_model_s3_region": ( + "{{ var.value.get('python_model_s3_region', 'us-west-2') }}" + ), + "OPUS-ETL-CLUSTER-HOST": "{{ var.value.get('opus-etl-cluster-host') }}", + "OPUS-ETL-USER-NAME": "{{ var.value.get('opus-etl-user-name') }}", + "OPUS-ETL-CLUSTER-PASSWORD": ( + "{{ var.value.get('opus-etl-cluster-password') }}" + ), + } + env["PYMODEL_DATES"] = ( + f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='dates_json') }}}}" + ) + return env + + +def _kpo_model_id_from_task_id(task_id: str) -> str | None: + marker = "kpo__" + if marker not in task_id: + return None + return task_id.rsplit(marker, 1)[-1] + + +def _gate_task_id_for_kpo(task_id: str) -> str: + if "kpo__" in task_id: + return task_id.replace("kpo__", "gate__", 1) + return task_id + + +def _record_kpo_success(context: dict) -> None: + _record_kpo_outcome(context, run_status="success") + + +def _record_kpo_failure(context: dict) -> None: + ti = context.get("ti") + run_message = None + if ti is not None: + run_message = getattr(ti, "note", None) or "KPO task failed" + _record_kpo_outcome(context, run_status="error", run_message=run_message) + + +def _record_kpo_outcome( + context: dict, + run_status: str, + run_message: str | None = None, +) -> None: + try: + ti = context["ti"] + model_id = _kpo_model_id_from_task_id(ti.task_id) + if not model_id: + return + + gate_task_id = _gate_task_id_for_kpo(ti.task_id) + etl_date = ti.xcom_pull(task_ids=gate_task_id, key="etl_date") or context["ds"] + model_group = ti.xcom_pull(task_ids=gate_task_id, key="model_group") or "" + etl_timestamp = _get_etl_timestamp_from_context(context) + + started_at = ti.start_date + run_seconds = 0.0 + if started_at is not None and ti.end_date is not None: + run_seconds = (ti.end_date - started_at).total_seconds() + + record_python_model_run( + model_name=model_id, + model_group=model_group, + etl_date=etl_date, + etl_timestamp=etl_timestamp or "", + run_status=run_status, + run_seconds=run_seconds, + is_skipped=False, + run_message=run_message, + context=context, + ) + except Exception as exc: + log.warning("Failed to record KPO model run: %s", exc) + + +def _wire_kpo_batch_depends_on( + model_tasks: list[dict], + gate_tasks: dict[str, PythonOperator], + kpo_tasks: dict[str, object], +) -> tuple[list[PythonOperator], list]: + """Wire depends_on across gate/KPO pairs; return entry gates and exit KPO tasks.""" + model_ids = {m["model_id"] for m in model_tasks} + for m in model_tasks: + model_id = m["model_id"] + for dep_id in m.get("depends_on", []): + if dep_id in model_ids and dep_id in kpo_tasks: + kpo_tasks[dep_id] >> gate_tasks[model_id] + + has_upstream: set[str] = set() + is_depended_on: set[str] = set() + for m in model_tasks: + model_id = m["model_id"] + for dep_id in m.get("depends_on", []): + if dep_id in model_ids: + has_upstream.add(model_id) + is_depended_on.add(dep_id) + + entry_gates = [ + gate_tasks[m["model_id"]] + for m in model_tasks + if m["model_id"] not in has_upstream + ] + exit_kpos = [ + kpo_tasks[m["model_id"]] + for m in model_tasks + if m["model_id"] not in is_depended_on + ] + return entry_gates, exit_kpos + + +def _register_kpo_model_batch( + dag, + topic: str, + model_tasks: list[dict], + prev_task, +) -> object: + """Create gate + KPO task pairs for a homogeneous KPO model batch.""" + gate_tasks: dict[str, PythonOperator] = {} + kpo_tasks: dict[str, object] = {} + + for m in model_tasks: + model_id = m["model_id"] + gate_task_id = f"gate__{model_id}" + gate_tasks[model_id] = PythonOperator( + task_id=gate_task_id, + python_callable=gate_model_run, + op_kwargs={"model": m, "topic": topic}, + dag=dag, + on_failure_callback=_record_airflow_task_failure, + ) + kpo_tasks[model_id] = build_kpo_task( + task_id=f"kpo__{model_id}", + dag=dag, + python_model_name=m["script_rel"], + env_vars=build_kpo_env_vars(gate_task_id), + script_args=["--run-id", "{{ dag_run.run_id }}", "--logical-date", "{{ ds }}"], + size=m.get("kpo_size", "small"), + labels={"component": "python-model", "variant": topic}, + on_success_callback=_record_kpo_success, + on_failure_callback=_record_kpo_failure, + ) + gate_tasks[model_id] >> kpo_tasks[model_id] + + entry_gates, exit_kpos = _wire_kpo_batch_depends_on( + model_tasks, gate_tasks, kpo_tasks + ) + if prev_task is not None: + prev_task >> entry_gates + + if len(exit_kpos) == 1: + return exit_kpos[0] + return exit_kpos + + def _execute_model_with_context( model: dict, topic: str = "", @@ -776,7 +1124,6 @@ def _execute_model_with_context( model_group = topic or model.get("topic", "") etl_timestamp = _get_etl_timestamp_from_context(context) etl_date = context["ds"] - dates_in: str | list | dict = "" started_at: float | None = None def _record( @@ -802,56 +1149,11 @@ def _record( log.warning("Failed to record python model run for %s: %s", model_name, exc) try: - if topic: - conf = (context.get("dag_run").conf or {}) if context.get("dag_run") else {} - runtime_config = conf.get("python_source_config") or get_python_source_config() - topic_cfg = runtime_config.get(topic, {}) - - if topic_cfg.get("skip", False): - log.info( - "Skipping model %s (topic %s skip=True at runtime)", - model_name, - topic, - ) - _record("skipped", 0.0, True) - raise AirflowSkipException(f"topic {topic} skip=True") - - date_to_run = topic_cfg.get("date_to_run", "") - if not passes_schedule_gate(date_to_run, context["ds"]): - log.info( - "Skipping model %s (topic %s schedule gate '%s' not met for ds %s)", - model_name, - topic, - date_to_run, - context["ds"], - ) - _record("skipped", 0.0, True) - raise AirflowSkipException( - f"topic {topic} schedule gate '{date_to_run}' not met" - ) - - skip_models = topic_cfg.get("skip_model", []) - if isinstance(skip_models, str): - skip_models = [skip_models] - skip_models = [m for m in skip_models if m and str(m).strip()] - if model_name in skip_models: - log.info( - "Skipping model %s (topic %s skip_model=%s)", - model_name, - topic, - skip_models, - ) - _record("skipped", 0.0, True) - return - - dates_in = topic_cfg.get("dates_in", "") - - dates = resolve_dates_in(dates_in, context["ds"]) - etl_date = dates[0] - ctx = {"ds": etl_date, "ds_nodash": etl_date.replace("-", "")} - if dates_in: - ctx["dates"] = dates + resolved = _resolve_model_run(model, topic, context, record_fn=_record) + if resolved is None: + return + ctx, etl_date, _ = resolved started_at = time.perf_counter() execute_model(model, ctx) _record("success", time.perf_counter() - started_at, False) @@ -863,6 +1165,43 @@ def _record( raise +def _register_model_task_batch( + dag, + topic: str, + model_tasks: list[dict], + prev_task, +) -> object: + """Register a topological batch using KPO or Airflow worker compute per model.""" + kpo_models = [m for m in model_tasks if _model_uses_kpo(m)] + airflow_models = [m for m in model_tasks if not _model_uses_kpo(m)] + batch_exit = prev_task + + if kpo_models: + batch_exit = _register_kpo_model_batch(dag, topic, kpo_models, batch_exit) + + if airflow_models: + task_name = "_".join(m["model_id"] for m in airflow_models) + op_kwargs_list = [ + { + "model": m, + "model_id": m["model_id"], + **({"topic": topic} if topic else {}), + } + for m in airflow_models + ] + mapped_task = PythonOperator.partial( + task_id=task_name, + python_callable=_execute_model_with_context, + dag=dag, + on_failure_callback=_record_airflow_task_failure, + ).expand(op_kwargs=op_kwargs_list) + if batch_exit is not None: + batch_exit >> mapped_task + batch_exit = mapped_task + + return batch_exit + + def register_python_model_mapped_tasks(dag_id: str, group_id: str | None = None): """Create a TaskGroup with model-task-batch mapped tasks for python models. @@ -883,20 +1222,12 @@ def register_python_model_mapped_tasks(dag_id: str, group_id: str | None = None) return None model_task_batches = _compute_model_tasks(models) + dag = _current_dag() with TaskGroup(group_id=group_id or dag_id) as tg: prev_task = None for model_tasks in model_task_batches: - task_name = "_".join(m["model_id"] for m in model_tasks) - mapped_task = PythonOperator.partial( - task_id=task_name, - python_callable=_execute_model_with_context, - on_failure_callback=_record_airflow_task_failure, - ).expand(op_kwargs=[{"model": m, "model_id": m["model_id"]} for m in model_tasks]) - - if prev_task is not None: - prev_task >> mapped_task - prev_task = mapped_task + prev_task = _register_model_task_batch(dag, "", model_tasks, prev_task) return tg @@ -940,128 +1271,15 @@ def register_python_model_mapped_tasks_by_topic( continue model_task_batches = _compute_model_tasks(topic_models) + dag = _current_dag() with TaskGroup(group_id=topic) as tg: prev_task = None for model_tasks in model_task_batches: - task_name = "_".join(m["model_id"] for m in model_tasks) - - op_kwargs_list = [ - { - "model": m, - "model_id": m["model_id"], - "topic": topic, - } - for m in model_tasks - ] - - mapped_task = PythonOperator.partial( - task_id=task_name, - python_callable=_execute_model_with_context, - on_failure_callback=_record_airflow_task_failure, - ).expand(op_kwargs=op_kwargs_list) - - if prev_task is not None: - prev_task >> mapped_task - prev_task = mapped_task + prev_task = _register_model_task_batch( + dag, topic, model_tasks, prev_task + ) topic_groups.append(tg) return topic_groups if topic_groups else None - - -# def register_k8s_model_tasks( -# dag_id: str, -# image: str, -# k8s_config_path: str, -# cluster_context: str, -# namespace: str = "mwaa", -# image_pull_secret: str = "wd-docker-artifactory-cred", -# dag=None, -# ) -> tuple: -# """Create KubernetesPodOperator tasks for each Python model and wire depends_on. -# -# Each pod runs dags/python_models/entrypoint.py inside the pymodels image. -# Airflow context (ds) and Trino credentials are injected as env vars via -# Jinja templates so values are resolved at execution time, not parse time. -# -# Returns (entry_tasks, exit_tasks) for chaining into the parent DAG. -# Returns (None, None) if no models are found. -# """ -# models = discover_models(dag_id) -# if not models: -# return None, None -# -# # Group models by task_group; models without one form solo groups keyed by model_id. -# groups: dict = {} -# for m in models: -# key = m.get("task_group") or m["model_id"] -# groups.setdefault(key, []).append(m) -# -# # model_id → group_key for cross-group dependency wiring -# model_to_group: dict = {m["model_id"]: (m.get("task_group") or m["model_id"]) for m in models} -# -# _shared_env = [ -# k8s.V1EnvVar(name="PYMODEL_DS", value="{{ ds }}"), -# k8s.V1EnvVar( -# name="OPUS-ETL-CLUSTER-HOST", -# value="{{ var.value.get('opus-etl-cluster-host') }}", -# ), -# k8s.V1EnvVar( -# name="OPUS-ETL-USER-NAME", -# value="{{ var.value.get('opus-etl-user-name') }}", -# ), -# k8s.V1EnvVar( -# name="OPUS-ETL-CLUSTER-PASSWORD", -# value="{{ var.value.get('opus-etl-cluster-password') }}", -# ), -# k8s.V1EnvVar( -# name="DJ_PYTHON_MODEL_CATALOG_NAME", -# value="{{ var.value.get('dj_python_model_catalog_name', 'glue') }}", -# ), -# ] -# -# tasks: dict = {} -# for group_key, group_models in groups.items(): -# if len(group_models) == 1: -# m = group_models[0] -# task_id = f"python_model__{m['model_id']}" -# group_env = [k8s.V1EnvVar(name="PYMODEL_MODEL_ID", value=m["model_id"])] -# else: -# task_id = f"python_model_group__{group_key}" -# ids = ",".join(m["model_id"] for m in group_models) -# group_env = [k8s.V1EnvVar(name="PYMODEL_MODEL_IDS", value=ids)] -# -# build_kw = dict( -# task_id=task_id, -# namespace=namespace, -# image=image, -# image_pull_secrets=[k8s.V1LocalObjectReference(name=image_pull_secret)], -# cmds=["python3.11", "/opt/py_models/entrypoint.py"], -# env_vars=_shared_env + group_env, -# is_delete_operator_pod=True, -# get_logs=True, -# config_file=k8s_config_path, -# in_cluster=False, -# cluster_context=cluster_context, -# startup_timeout_seconds=300, -# ) -# if dag is not None: -# build_kw["dag"] = dag -# tasks[group_key] = KubernetesPodOperator(**build_kw) -# -# # Wire cross-group Airflow dependencies. -# has_upstream: set = set() -# is_depended_on: set = set() -# for m in models: -# group_key = model_to_group[m["model_id"]] -# for dep_model_id in m.get("depends_on", []): -# dep_group = model_to_group.get(dep_model_id) -# if dep_group and dep_group != group_key and dep_group in tasks: -# tasks[dep_group] >> tasks[group_key] -# has_upstream.add(group_key) -# is_depended_on.add(dep_group) -# -# entry_tasks = [t for k, t in tasks.items() if k not in has_upstream] -# exit_tasks = [t for k, t in tasks.items() if k not in is_depended_on] -# return entry_tasks, exit_tasks diff --git a/airflow/v2_7/kpo_factory.py b/airflow/v2_7/kpo_factory.py new file mode 100644 index 0000000..ca561a7 --- /dev/null +++ b/airflow/v2_7/kpo_factory.py @@ -0,0 +1,195 @@ +"""Reusable KubernetesPodOperator factory for DJ python model execution. + +Target: Airflow 2.10 with ``apache-airflow-providers-cncf-kubernetes`` 8.x +(uses ``on_finish_action``, not ``is_delete_operator_pod``). + +Cluster/image config is resolved at task-execution time from the shared +Airflow Variable ``airflow_runner_cfg``. Pod sizing presets are loaded from +``kpo_sizes.yml`` beside this module (optional ``size_specs`` override in +the Variable JSON). +""" + +from __future__ import annotations + +from datetime import timedelta +from functools import lru_cache +from pathlib import Path +from typing import Literal, Mapping, Sequence + +from airflow import DAG +from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator +from kubernetes.client import models as k8s + +_DEFAULT_IMAGE_PULL_SECRETS: tuple[str, ...] = ("wd-docker-artifactory-cred",) + +_RUNNER_CFG_VAR_NAME: str = "airflow_runner_cfg" +_IMAGE_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.image }}}}" +_TAG_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.tag }}}}" +_NAMESPACE_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.namespace }}}}" +_CONFIG_FILE_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.config_file }}}}" +_CLUSTER_CONTEXT_TPL: str = f"{{{{ var.json.{_RUNNER_CFG_VAR_NAME}.cluster_context }}}}" + +_IMAGE_REF: str = f"{_IMAGE_TPL}:{_TAG_TPL}" + +PodSize = Literal["small", "medium", "large", "xlarge"] + +_DEFAULT_SIZE_SPECS: dict[str, tuple[str, str]] = { + "small": ("1", "2Gi"), + "medium": ("2", "4Gi"), + "large": ("2", "8Gi"), + "xlarge": ("2", "16Gi"), +} + +_SIZES_YAML_PATH = Path(__file__).resolve().parent / "kpo_sizes.yml" + + +def _parse_size_specs_from_yaml(path: Path) -> dict[str, tuple[str, str]]: + if not path.is_file(): + return dict(_DEFAULT_SIZE_SPECS) + + try: + import yaml + except ImportError: + return dict(_DEFAULT_SIZE_SPECS) + + try: + raw = yaml.safe_load(path.read_text()) or {} + except Exception: + return dict(_DEFAULT_SIZE_SPECS) + + sizes = raw.get("sizes") or {} + parsed: dict[str, tuple[str, str]] = {} + for name, spec in sizes.items(): + if not isinstance(spec, dict): + continue + cpu = spec.get("cpu") + memory = spec.get("memory") + if cpu and memory: + parsed[str(name)] = (str(cpu), str(memory)) + return parsed if parsed else dict(_DEFAULT_SIZE_SPECS) + + +def _merge_variable_size_specs(specs: dict[str, tuple[str, str]]) -> dict[str, tuple[str, str]]: + try: + from airflow.models import Variable + + cfg = Variable.get(_RUNNER_CFG_VAR_NAME, deserialize_json=True, default_var={}) + if not isinstance(cfg, dict): + return specs + overrides = cfg.get("size_specs") or {} + if not isinstance(overrides, dict): + return specs + merged = dict(specs) + for name, spec in overrides.items(): + if isinstance(spec, dict) and spec.get("cpu") and spec.get("memory"): + merged[str(name)] = (str(spec["cpu"]), str(spec["memory"])) + return merged + except Exception: + return specs + + +@lru_cache(maxsize=1) +def load_size_specs() -> dict[str, tuple[str, str]]: + """Load pod CPU/memory presets from kpo_sizes.yml with optional Variable override.""" + return _merge_variable_size_specs(_parse_size_specs_from_yaml(_SIZES_YAML_PATH)) + + +def _sanitize_pod_name(dag_id: str, task_id: str) -> str: + raw = f"{dag_id}-{task_id}".replace("_", "-").lower() + return raw[:63].rstrip("-") + + +def _container_resources(size: PodSize) -> k8s.V1ResourceRequirements: + specs = load_size_specs() + if size not in specs: + raise ValueError( + f"Invalid pod size {size!r}; expected one of {sorted(specs)}." + ) + cpu, memory = specs[size] + return k8s.V1ResourceRequirements( + requests={"cpu": cpu, "memory": memory}, + limits={"cpu": cpu, "memory": memory}, + ) + + +def _env_var_list(env: Mapping[str, str]) -> list[k8s.V1EnvVar]: + return [k8s.V1EnvVar(name=name, value=value) for name, value in env.items()] + + +def build_kpo_task( + *, + task_id: str, + dag: DAG, + python_model_name: str, + env_vars: Mapping[str, str], + script_args: Sequence[str] | None = None, + size: PodSize = "small", + image_pull_policy: str = "IfNotPresent", + image_pull_secrets: Sequence[str] = _DEFAULT_IMAGE_PULL_SECRETS, + service_account_name: str = "default", + retries: int = 1, + retry_delay: timedelta = timedelta(minutes=5), + retry_exponential_backoff: bool = True, + startup_timeout_seconds: int = 600, + on_finish_action: str = "delete_pod", + deferrable: bool = False, + logging_interval: int | None = 10, + do_xcom_push: bool = False, + labels: Mapping[str, str] | None = None, + annotations: Mapping[str, str] | None = None, + on_success_callback=None, + on_failure_callback=None, +) -> KubernetesPodOperator: + """Build a KubernetesPodOperator that invokes the image ENTRYPOINT.""" + env_map: dict[str, str] = {"PYTHONUNBUFFERED": "1"} + env_map.update(env_vars) + + dag_labels: dict[str, str] = { + "dag_id": dag.dag_id, + "task_id": task_id, + "run_id": "{{ dag_run.run_id | truncate(63, true, '') }}", + "component": "airflow-kpo", + "variant": dag.dag_id, + } + if labels: + dag_labels.update(labels) + + arguments = [python_model_name, *(list(script_args) if script_args else [])] + + deferred_log_kwargs: dict[str, int] = {} + if deferrable and logging_interval is not None: + deferred_log_kwargs["logging_interval"] = logging_interval + + return KubernetesPodOperator( + task_id=task_id, + dag=dag, + image=_IMAGE_REF, + image_pull_policy=image_pull_policy, + image_pull_secrets=[ + k8s.V1LocalObjectReference(name=n) for n in image_pull_secrets + ], + in_cluster=False, + config_file=_CONFIG_FILE_TPL, + cluster_context=_CLUSTER_CONTEXT_TPL, + namespace=_NAMESPACE_TPL, + name=_sanitize_pod_name(dag.dag_id, task_id), + arguments=arguments, + env_vars=_env_var_list(env_map), + container_resources=_container_resources(size), + service_account_name=service_account_name, + get_logs=True, + log_events_on_failure=True, + on_finish_action=on_finish_action, + do_xcom_push=do_xcom_push, + deferrable=deferrable, + **deferred_log_kwargs, + retries=retries, + retry_delay=retry_delay, + retry_exponential_backoff=retry_exponential_backoff, + startup_timeout_seconds=startup_timeout_seconds, + reattach_on_restart=True, + labels=dag_labels, + annotations=dict(annotations or {}), + on_success_callback=on_success_callback, + on_failure_callback=on_failure_callback, + ) diff --git a/airflow/v2_7/kpo_sizes.yml b/airflow/v2_7/kpo_sizes.yml new file mode 100644 index 0000000..868dd7f --- /dev/null +++ b/airflow/v2_7/kpo_sizes.yml @@ -0,0 +1,19 @@ +# Pod resource presets for KPO python models. +# CPU requests == limits (Guaranteed QoS). CPU capped at 2 on shared EKS nodes. +sizes: + small: + cpu: "1" + memory: "2Gi" + guidance: "Light jobs — small API responses, narrow CSV extracts, < ~500 MB processed in-memory." + medium: + cpu: "2" + memory: "4Gi" + guidance: "Typical ETL — moderate API pagination, joins on staging tables, ~500 MB–2 GB." + large: + cpu: "2" + memory: "8Gi" + guidance: "Heavier transforms — wider schemas, multi-step pandas work, ~2–8 GB." + xlarge: + cpu: "2" + memory: "16Gi" + guidance: "Large workloads — full-table scans, wide fin-style exports, ~8–16 GB." diff --git a/package-lock.json b/package-lock.json index 5501236..86ff31b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -15460,4 +15460,4 @@ } } } -} +} \ No newline at end of file diff --git a/schemas/python-model.schema.json b/schemas/python-model.schema.json index 8a7c06b..b89e9f3 100644 --- a/schemas/python-model.schema.json +++ b/schemas/python-model.schema.json @@ -41,6 +41,18 @@ "default": [], "description": "Names of other Python models that must complete before this model runs" }, + "compute": { + "type": "string", + "enum": ["kpo", "airflow"], + "default": "kpo", + "description": "Execution backend: kpo (Kubernetes pod, default) or airflow (MWAA worker)" + }, + "kpo_size": { + "type": "string", + "enum": ["small", "medium", "large", "xlarge"], + "default": "small", + "description": "KPO pod resource preset (only when compute is kpo)" + }, "output_type": { "type": "string", "enum": ["iceberg", "s3"], diff --git a/src/services/framework/__tests__/python-model-roundtrip.test.ts b/src/services/framework/__tests__/python-model-roundtrip.test.ts index 04cf89c..a19b5ee 100644 --- a/src/services/framework/__tests__/python-model-roundtrip.test.ts +++ b/src/services/framework/__tests__/python-model-roundtrip.test.ts @@ -25,6 +25,8 @@ const CONFIG: PythonModelConfig = { dags: ['full_source_etl'], namespace: 'opus_python_source', table_name: 'my_model', + compute: 'kpo', + kpo_size: 'medium', }; describe('parsePythonToCells / serializeCellsToPython', () => { @@ -201,4 +203,18 @@ describe('.py -> notebook -> .py stability', () => { const cells = buildNotebookCells(py, CONFIG); expect(serializeCellsToPython(cells)).toBe(py); }); + + test('compute and kpo_size survive scaffold generation', () => { + const py = generatePythonModelScaffoldPy(CONFIG); + expect(py).toContain('my_model'); + expect(CONFIG.compute).toBe('kpo'); + expect(CONFIG.kpo_size).toBe('medium'); + }); + + test('scaffold includes __main__ entry with build_context_from_env', () => { + const py = generatePythonModelScaffoldPy(CONFIG); + expect(py).toContain('if __name__ == "__main__":'); + expect(py).toContain('build_context_from_env'); + expect(py).toContain('run_etl(build_context_from_env())'); + }); }); diff --git a/src/services/framework/index.ts b/src/services/framework/index.ts index 285cdea..8803e2d 100644 --- a/src/services/framework/index.ts +++ b/src/services/framework/index.ts @@ -75,6 +75,10 @@ import { generateTrinoIoPy, workspaceHasPythonModels, } from './utils'; +import { + ETL_HELPER_BUNDLE_FILES, + readKpoSizesTemplate, +} from './utils/kpo-sizes'; import { resolveModelJsonPath } from './utils/resolve-model-json-path'; /** @@ -411,6 +415,8 @@ export class Framework implements ApiEnabledService<'framework'> { tags, namespace: reqNamespace, table_name: reqTableName, + compute: reqCompute, + kpo_size: reqKpoSize, create_dag: createDag, dag_config: dagConfig, } = payload.request; @@ -458,6 +464,10 @@ export class Framework implements ApiEnabledService<'framework'> { ...(description && { description }), model_type, dags, + compute: reqCompute ?? 'kpo', + ...(reqCompute === 'kpo' || reqCompute === undefined + ? { kpo_size: reqKpoSize ?? 'small' } + : {}), ...(enable_notebook && { enable_notebook }), tags: tags && tags.length > 0 ? tags : ['python-model', group], ...(reqNamespace && { namespace: reqNamespace }), @@ -601,6 +611,8 @@ export class Framework implements ApiEnabledService<'framework'> { case 'framework-get-python-model-groups': return { groups: getDjConfig().pythonModelGroups }; + case 'framework-get-kpo-sizes': + return { sizes: await readKpoSizesTemplate() }; case 'framework-model-cte-analysis': return await this.cteAnalysisHandlers.handleCteAnalysis(payload); @@ -2176,6 +2188,8 @@ export class Framework implements ApiEnabledService<'framework'> { const { airflowDagsPath } = getDjConfig(); const scaffoldingPaths = [ path.join(WORKSPACE_ROOT, 'dags', '_ext_', 'etl_helper.py'), + path.join(WORKSPACE_ROOT, 'dags', '_ext_', 'kpo_factory.py'), + path.join(WORKSPACE_ROOT, 'dags', '_ext_', 'kpo_sizes.yml'), path.join(WORKSPACE_ROOT, 'dags', 'python_models', '_config.py'), path.join(WORKSPACE_ROOT, 'dags', 'python_models', '_trino_io.py'), path.join(WORKSPACE_ROOT, 'dags', 'python_models', '_model_tasks.py'), @@ -2183,6 +2197,8 @@ export class Framework implements ApiEnabledService<'framework'> { if (airflowDagsPath) { scaffoldingPaths.push( path.join(WORKSPACE_ROOT, airflowDagsPath, 'etl_helper.py'), + path.join(WORKSPACE_ROOT, airflowDagsPath, 'kpo_factory.py'), + path.join(WORKSPACE_ROOT, airflowDagsPath, 'kpo_sizes.yml'), ); } @@ -2258,16 +2274,15 @@ export class Framework implements ApiEnabledService<'framework'> { } /** - * Read the etl_helper.py template for the configured Airflow version. - * Used both to seed the file and to revert unauthorized edits. + * Read an Airflow template file for the configured target version. */ - private async readEtlHelperTemplate(): Promise { + private async readAirflowTemplate(fileName: string): Promise { const { airflowTargetVersion } = getDjConfig(); const versionFolder = airflowTargetVersion === '2.10' ? 'v2_10' : 'v2_7'; const templatePath = path.join( BASE_AIRFLOW_PATH, versionFolder, - 'etl_helper.py', + fileName, ); const content = await vscode.workspace.fs.readFile( vscode.Uri.file(templatePath), @@ -2275,6 +2290,13 @@ export class Framework implements ApiEnabledService<'framework'> { return Buffer.from(content).toString(); } + /** + * Read the etl_helper.py template for the configured Airflow version. + */ + private async readEtlHelperTemplate(): Promise { + return this.readAirflowTemplate('etl_helper.py'); + } + /** * Watch all DJ-generated python model shared files and revert unauthorized * edits. These files are auto-generated and must not be hand-edited: @@ -2306,6 +2328,16 @@ export class Framework implements ApiEnabledService<'framework'> { generate: () => this.readEtlHelperTemplate(), label: 'etl_helper.py', }, + { + glob: '**/kpo_factory.py', + generate: () => this.readAirflowTemplate('kpo_factory.py'), + label: 'kpo_factory.py', + }, + { + glob: '**/kpo_sizes.yml', + generate: () => this.readAirflowTemplate('kpo_sizes.yml'), + label: 'kpo_sizes.yml', + }, ]; for (const guard of guards) { @@ -2341,6 +2373,7 @@ export class Framework implements ApiEnabledService<'framework'> { '_model_tasks.py', '__init__.py', 'etl_helper.py', + 'kpo_factory.py', ]); /** Whether a path is a python model file that the SQL guard should scan. */ @@ -2478,50 +2511,49 @@ export class Framework implements ApiEnabledService<'framework'> { private async ensureEtlHelperFile(): Promise { const { airflowDagsPath, airflowTargetVersion } = getDjConfig(); - // Check possible locations for etl_helper.py const possibleTargets: string[] = []; if (airflowDagsPath) { - possibleTargets.push( - path.join(WORKSPACE_ROOT, airflowDagsPath, 'etl_helper.py'), - ); + possibleTargets.push(path.join(WORKSPACE_ROOT, airflowDagsPath)); } - possibleTargets.push( - path.join(WORKSPACE_ROOT, 'dags', '_ext_', 'etl_helper.py'), - ); + possibleTargets.push(path.join(WORKSPACE_ROOT, 'dags', '_ext_')); - // Determine best target: overwrite an existing file wherever it lives, - // otherwise prefer airflowDagsPath, then workspace dags/_ext_/. - let targetPath = possibleTargets[0]; + let targetDir = possibleTargets[0]; for (const target of possibleTargets) { try { - await vscode.workspace.fs.stat(vscode.Uri.file(target)); - targetPath = target; + await vscode.workspace.fs.stat( + vscode.Uri.file(path.join(target, 'etl_helper.py')), + ); + targetDir = target; break; } catch { // Not found, continue } } - // Determine source template const versionFolder = airflowTargetVersion === '2.10' ? 'v2_10' : 'v2_7'; - const templatePath = path.join( - BASE_AIRFLOW_PATH, - versionFolder, - 'etl_helper.py', - ); try { - const content = await vscode.workspace.fs.readFile( - vscode.Uri.file(templatePath), - ); - // Ensure target directory exists - const targetDir = path.dirname(targetPath); await vscode.workspace.fs.createDirectory(vscode.Uri.file(targetDir)); - await vscode.workspace.fs.writeFile(vscode.Uri.file(targetPath), content); - this.log.info(`Generated etl_helper.py at: ${targetPath}`); + + for (const fileName of ETL_HELPER_BUNDLE_FILES) { + const templatePath = path.join( + BASE_AIRFLOW_PATH, + versionFolder, + fileName, + ); + const content = await vscode.workspace.fs.readFile( + vscode.Uri.file(templatePath), + ); + const targetPath = path.join(targetDir, fileName); + await vscode.workspace.fs.writeFile( + vscode.Uri.file(targetPath), + content, + ); + this.log.info(`Generated ${fileName} at: ${targetPath}`); + } } catch (err) { - this.log.warn(`Could not generate etl_helper.py: ${String(err)}`); + this.log.warn(`Could not generate etl helper bundle: ${String(err)}`); } } diff --git a/src/services/framework/utils/__tests__/kpo-sizes.test.ts b/src/services/framework/utils/__tests__/kpo-sizes.test.ts new file mode 100644 index 0000000..1419d50 --- /dev/null +++ b/src/services/framework/utils/__tests__/kpo-sizes.test.ts @@ -0,0 +1,23 @@ +import { describe, expect, test } from '@jest/globals'; + +import { readKpoSizesTemplate } from '@services/framework/utils/kpo-sizes'; + +describe('readKpoSizesTemplate', () => { + test('loads cpu, memory, and guidance for all preset sizes', async () => { + const sizes = await readKpoSizesTemplate(); + + expect(sizes.small).toEqual({ + cpu: '1', + memory: '2Gi', + guidance: expect.stringContaining('Light jobs'), + }); + expect(sizes.medium).toEqual({ + cpu: '2', + memory: '4Gi', + guidance: expect.stringContaining('Typical ETL'), + }); + expect(sizes.large?.cpu).toBe('2'); + expect(sizes.large?.memory).toBe('8Gi'); + expect(sizes.xlarge?.memory).toBe('16Gi'); + }); +}); diff --git a/src/services/framework/utils/__tests__/python-model-config.test.ts b/src/services/framework/utils/__tests__/python-model-config.test.ts new file mode 100644 index 0000000..65e8be7 --- /dev/null +++ b/src/services/framework/utils/__tests__/python-model-config.test.ts @@ -0,0 +1,12 @@ +import { describe, expect, test } from '@jest/globals'; + +import { generatePythonModelConfigPy } from '@services/framework/utils/python-model-utils'; + +describe('generatePythonModelConfigPy', () => { + test('includes build_context_from_env for __main__ script execution', () => { + const py = generatePythonModelConfigPy(); + expect(py).toContain('def build_context_from_env()'); + expect(py).toContain('PYMODEL_DS'); + expect(py).toContain('PYMODEL_DATES'); + }); +}); diff --git a/src/services/framework/utils/kpo-sizes.ts b/src/services/framework/utils/kpo-sizes.ts new file mode 100644 index 0000000..2b8b656 --- /dev/null +++ b/src/services/framework/utils/kpo-sizes.ts @@ -0,0 +1,48 @@ +import fs from 'fs/promises'; +import path from 'path'; +import { parse as parseYaml } from 'yaml'; + +import { BASE_AIRFLOW_PATH } from '@services/constants'; +import { getDjConfig } from '@services/config'; +import type { KpoSizeSpec } from '@shared/framework/types'; + +/** + * Read KPO pod size presets from the extension's airflow template YAML. + */ +export async function readKpoSizesTemplate(): Promise< + Record +> { + const { airflowTargetVersion } = getDjConfig(); + const versionFolder = airflowTargetVersion === '2.10' ? 'v2_10' : 'v2_7'; + const templatePath = path.join( + BASE_AIRFLOW_PATH, + versionFolder, + 'kpo_sizes.yml', + ); + + const content = await fs.readFile(templatePath, 'utf8'); + const raw = parseYaml(content) as { + sizes?: Record>; + }; + const sizes = raw?.sizes ?? {}; + const result: Record = {}; + + for (const [name, spec] of Object.entries(sizes)) { + if (!spec?.cpu || !spec?.memory) { + continue; + } + result[name] = { + cpu: String(spec.cpu), + memory: String(spec.memory), + guidance: spec.guidance ? String(spec.guidance) : '', + }; + } + + return result; +} + +export const ETL_HELPER_BUNDLE_FILES = [ + 'etl_helper.py', + 'kpo_factory.py', + 'kpo_sizes.yml', +] as const; diff --git a/src/services/framework/utils/python-model-utils.ts b/src/services/framework/utils/python-model-utils.ts index d31784a..155fedf 100644 --- a/src/services/framework/utils/python-model-utils.ts +++ b/src/services/framework/utils/python-model-utils.ts @@ -47,6 +47,22 @@ def get_namespace_name() -> str: return _require_variable(NAMESPACE_VARIABLE) +def build_context_from_env() -> dict: + """Build Airflow run context from PYMODEL_* env vars (KPO pod or MWAA worker).""" + import json + import os + + ds = os.environ.get("PYMODEL_DS") or "" + ctx: dict = { + "ds": ds, + "ds_nodash": os.environ.get("PYMODEL_DS_NODASH") or ds.replace("-", ""), + } + dates_raw = os.environ.get("PYMODEL_DATES") + if dates_raw: + ctx["dates"] = json.loads(dates_raw) + return ctx + + @dataclass(frozen=True) class PythonModelConfig: model_name: str @@ -709,6 +725,12 @@ function buildScaffoldCodeCells(config: PythonModelConfig): PythonModelCell[] { ' cleanup(context)', ].join('\n'); + const mainCell = [ + 'if __name__ == "__main__":', + ' from python_models._config import build_context_from_env', + ' run_etl(build_context_from_env())', + ].join('\n'); + return [ makeCodeCell(headerCell), makeCodeCell(inputVarsCell), @@ -716,6 +738,7 @@ function buildScaffoldCodeCells(config: PythonModelConfig): PythonModelCell[] { makeCodeCell(transformCell), makeCodeCell(cleanupCell), makeCodeCell(runEtlCell), + makeCodeCell(mainCell), ]; } diff --git a/src/shared/framework/types.ts b/src/shared/framework/types.ts index 760db89..c57342e 100644 --- a/src/shared/framework/types.ts +++ b/src/shared/framework/types.ts @@ -58,6 +58,15 @@ export interface PythonModelCell { execution_count?: number | null; } +export type PythonModelCompute = 'kpo' | 'airflow'; +export type PythonModelKpoSize = 'small' | 'medium' | 'large' | 'xlarge'; + +export interface KpoSizeSpec { + cpu: string; + memory: string; + guidance: string; +} + export interface PythonModelConfig { name: string; group: PythonModelGroup; @@ -66,6 +75,8 @@ export interface PythonModelConfig { model_type: 'python'; dags?: string[]; depends_on?: string[]; + compute?: PythonModelCompute; + kpo_size?: PythonModelKpoSize; output_type?: PythonModelOutputType; output?: PythonModelOutputConfig; enable_notebook?: boolean; @@ -524,6 +535,8 @@ export type FrameworkApi = model_type: 'python'; dags?: string[]; depends_on?: string[]; + compute?: PythonModelCompute; + kpo_size?: PythonModelKpoSize; output_type?: PythonModelOutputType; enable_notebook?: boolean; tags?: string[]; @@ -564,6 +577,14 @@ export type FrameworkApi = request: null; response: { groups: string[] }; } + | { + type: 'framework-get-kpo-sizes'; + service: 'framework'; + request: null; + response: { + sizes: Record; + }; + } | { /** * Read-only CTE column inference + diagnostics for the model wizard. @@ -688,6 +709,10 @@ async function apiHandler(p: { type: 'framework-get-python-model-groups'; request: ApiRequest<'framework-get-python-model-groups'>; }): Promise>; +async function apiHandler(p: { + type: 'framework-get-kpo-sizes'; + request: ApiRequest<'framework-get-kpo-sizes'>; +}): Promise>; async function apiHandler(p: { type: 'framework-model-cte-analysis'; request: ApiRequest<'framework-model-cte-analysis'>; diff --git a/src/shared/schema/types/python-model.schema.d.ts b/src/shared/schema/types/python-model.schema.d.ts index 88c821d..b44fedd 100644 --- a/src/shared/schema/types/python-model.schema.d.ts +++ b/src/shared/schema/types/python-model.schema.d.ts @@ -46,6 +46,14 @@ export interface SchemaPythonModel { * Names of other Python models that must complete before this model runs */ depends_on?: string[]; + /** + * Execution backend: kpo (Kubernetes pod, default) or airflow (MWAA worker) + */ + compute?: 'kpo' | 'airflow'; + /** + * KPO pod resource preset (only when compute is kpo) + */ + kpo_size?: 'small' | 'medium' | 'large' | 'xlarge'; /** * Output backend type for data writes */ diff --git a/templates/skills/dj-create-python-model/_SKILL.md b/templates/skills/dj-create-python-model/_SKILL.md index 4da1146..1a05b05 100644 --- a/templates/skills/dj-create-python-model/_SKILL.md +++ b/templates/skills/dj-create-python-model/_SKILL.md @@ -169,6 +169,21 @@ dags/python_models///.python.json Model ID convention: `python______` +## Compute and KPO image naming + +New models default to **`compute: "kpo"`** (Kubernetes pod) instead of running on the MWAA worker. Set **`compute: "airflow"`** explicitly when the KPO image pipeline is not ready. + +| Field | Purpose | +| ----- | ------- | +| `compute` | `kpo` (default) or `airflow` | +| `kpo_size` | Pod preset when `compute` is `kpo`: `small`, `medium`, `large`, `xlarge` | + +**KPO script path:** The `mwaa-opus-kpo-runner` image syncs `s3:///dags/python_models/` into the pod and runs `python //.python.py`. Airflow passes that relative path as the container argument. Set Airflow Variable **`python_models_s3_bucket`** (DAG bucket name) on KPO tasks; optional **`python_model_s3_region`** (default `us-west-2`). + +**Execution entry:** Scaffolded `.python.py` ends with `if __name__ == "__main__":` calling `run_etl(build_context_from_env())`. `build_context_from_env()` in `_config.py` reads `PYMODEL_DS`, `PYMODEL_DS_NODASH`, and `PYMODEL_DATES` from the environment. MWAA worker and KPO pod both run the file as `__main__` — `run_etl` is scaffold convention, not a framework-enforced function name. + +Configure **`airflow_runner_cfg`** (JSON: `image`, `tag`, `namespace`, `config_file`, `cluster_context`; optional `size_specs` override for pod presets in `dags/_ext_/kpo_sizes.yml`). + ## Schema reference Before writing, read `.dj/schemas/python-model.schema.json` to validate field shapes. Required fields: `name`, `group`, `topic`. @@ -203,7 +218,8 @@ Use this when data comes from an external source (API, CSV, S3) but all transfor 6. **Transform + Load** — `def transform_and_load(context)` runs SQL `INSERT INTO ... SELECT` with all transformations, using `overwrite_partition()` or other DML helpers from `_trino_io` 7. **Cleanup** — `def cleanup(context)` drops the temporary staging table via `execute_trino()` 8. **Main** — `def run_etl(context)` calls extract -> transform_and_load -> cleanup -9. **Runner** — notebook-only `run_etl(context)` call +9. **Script entry** — `if __name__ == "__main__":` calls `run_etl(build_context_from_env())` (included in `.python.py`) +10. **Runner** — notebook-only `run_etl(context)` call for interactive use ### `_trino_io.py` DML operations reference @@ -315,7 +331,8 @@ The output table must be **registered as a `.source.json`** before a downstream - **SQL-first** — prefer Trino SQL for all transformations that SQL can express. Use DataFrames only for external data ingestion and Python-only logic - **`cells` mode:** only edit `.python.json` — the extension regenerates `.python.py` (and `.python.ipynb`) from the `cells` array. **`hand-written` mode:** edit `.python.py` directly; keep JSON metadata in sync -- **`run_etl()` is required** — Airflow's `etl_helper.py` scans for `def run_etl(` to discover models +- **`if __name__ == "__main__"`** — scaffold includes this block in `.python.py`; Airflow and KPO execute the script as `__main__` (use `build_context_from_env()` from `_config.py` for partition dates) +- **`run_etl()`** — scaffold convention for orchestrating extract/transform/cleanup; not required by discovery or execution - **`OUTPUT_CONFIG`** must use `PythonModelConfig` from `python_models/_config.py` — the extension auto-generates this file - **`portal_partition_daily`** must be set as a column (either via SQL `'{ds}' AS portal_partition_daily` or pandas assignment) for partition-based writes - **Always suggest optimizations** — partitioning strategy, column pruning, predicate pushdown, write mode selection diff --git a/templates/skills/dj-review-python-model/_SKILL.md b/templates/skills/dj-review-python-model/_SKILL.md index dcaae68..4a0ab0a 100644 --- a/templates/skills/dj-review-python-model/_SKILL.md +++ b/templates/skills/dj-review-python-model/_SKILL.md @@ -57,7 +57,7 @@ When `.dj/bin/dj system.ping` succeeds, use `dbt.sources` and `model.lineage` to | F1 | `.python.json` has required fields: `name`, `group`, `topic` | | F2 | Name/group/topic match pattern `^[a-z][a-z0-9_]*$` | | F3 | **`cells` mode:** `cells` array is present and non-empty in JSON. **`hand-written` mode:** empty or absent `cells` is OK | -| F4 | `.python.py` contains `def run_etl(context)` function | +| F4 | `.python.py` contains `if __name__ == "__main__":` script entry (scaffold calls `run_etl(build_context_from_env())`) | | F5 | Uses `_trino_io` helpers (`from python_models._trino_io import ...`) — no inline Trino connection code (`trino.dbapi.connect`, `create_engine`, raw `requests.post` to Trino) | | F6 | `OUTPUT_CONFIG` uses `PythonModelConfig` from `python_models._config` | | F7 | **`cells` mode:** `.python.py` derivable from JSON `cells` (no hand-edits lost on sync). **`hand-written` mode:** `.python.py` is source of truth — hand edits expected | @@ -84,7 +84,7 @@ The review validates the **full chain**: JSON metadata → `PythonModelConfig` | L6 | `dags` field is populated in JSON (model is scheduled, discoverable by Airflow DAG lineage). Empty `dags` = utility module, not a lineage participant | | L7 | `depends_on` correctly lists upstream Python models whose output tables this model reads (task dependency mirrors data dependency) | | L8 | Model ID derivable from file path matches `python_model_name` in properties: `python______` | -| L9 | Companion `.python.py` exists alongside `.python.json` (required for Airflow `etl_helper.py` discovery via `def run_etl(`) | +| L9 | Companion `.python.py` exists alongside `.python.json` (required for Airflow `etl_helper.py` discovery) | | L10 | `OUTPUT_CONFIG.table_name` (or `model_name` fallback) matches the table name referenced in downstream `.source.json` files, if any exist in the project | ### 3. Downstream Integration (D) @@ -174,7 +174,7 @@ Render the full report in this structure. Adapt headings to the actual model. Om | Priority | Criteria | | -------- | ----------------------------------------------------------------------- | -| CRITICAL | Lineage broken (L2-L5 failures), missing `run_etl`, no partition column | +| CRITICAL | Lineage broken (L2-L5 failures), missing `__main__` entry, no partition column | | HIGH | Inline Trino code, missing cleanup, stale upstream_sources | | MEDIUM | Non-SQL-first transforms, missing batch staging, no error handling | | LOW | Naming style, column order, documentation gaps | @@ -205,7 +205,8 @@ To validate upstream source completeness: - **Model ID:** `python______` - **Source of truth:** `cells` mode → `.python.json`; `hand-written` mode → `.python.py` (metadata in JSON) - **`_trino_io.py` DML helpers:** `execute_trino`, `overwrite_partition` / `overwrite` (keyword `insert_sql=` or `source_query=` + optional `columns=`), `append`, `merge`, `delete`, `update` -- **Required function:** `def run_etl(context)` — Airflow discovers models by scanning for this +- **Script entry:** `if __name__ == "__main__":` with `build_context_from_env()` from `_config.py` — Airflow and KPO run the file as `__main__` +- **Scaffold convention:** `def run_etl(context)` orchestrates extract/transform/cleanup; recommended but not enforced at discovery - **Partition contract:** downstream dbt models expect `portal_partition_daily` column in output - **Table properties for lineage:** `python_model_name`, `python_model_type`, `python_model_namespace`, `python_model_table`, `python_model_description`, `python_model_upstream_sources` diff --git a/templates/skills/dj-review-python-model/references/review-checklist.md b/templates/skills/dj-review-python-model/references/review-checklist.md index a75020f..8a9fe76 100644 --- a/templates/skills/dj-review-python-model/references/review-checklist.md +++ b/templates/skills/dj-review-python-model/references/review-checklist.md @@ -33,18 +33,13 @@ Extended explanations, code examples, and edge cases for each check in the DJ Py **Why it matters:** The framework regenerates `.python.py` from `cells`. If `cells` is empty/missing and `.python.py` exists, the sync service skips overwriting (safety valve) — but this means the JSON is no longer the source of truth. -### F4: `run_etl` function +### F4: `__main__` script entry -**Pass condition:** `.python.py` contains `def run_etl(context)` (exact signature). +**Pass condition:** `.python.py` contains `if __name__ == "__main__":` and uses `build_context_from_env()` from `_config.py` (scaffold default: `run_etl(build_context_from_env())`). -**Why it matters:** Airflow's `etl_helper.py` scans for this function to discover executable models: -```python -# etl_helper.py discovery logic -if hasattr(module, 'run_etl') and callable(module.run_etl): - # register as PythonOperator task -``` +**Why it matters:** Airflow worker and KPO pods execute the script as `__main__` via `runpy.run_path` or `python ` — not import + callable lookup. -**Common failure:** Function named `main()`, `run()`, or `etl()` instead of `run_etl`. +**Common failure:** Missing `__main__` block after migrating from the old import-based execution model. ### F5: `_trino_io` usage diff --git a/web/src/context/app.tsx b/web/src/context/app.tsx index bc9afe3..18d797e 100644 --- a/web/src/context/app.tsx +++ b/web/src/context/app.tsx @@ -1872,6 +1872,26 @@ SELECT * FROM final`, }), ); } + case 'framework-get-kpo-sizes': { + return resolve( + apiResponse({ + sizes: { + small: { + cpu: '1', + memory: '2Gi', + guidance: + 'Light jobs — small API responses, narrow CSV extracts, < ~500 MB processed in-memory.', + }, + medium: { + cpu: '2', + memory: '4Gi', + guidance: + 'Typical ETL — moderate API pagination, joins on staging tables, ~500 MB–2 GB.', + }, + }, + }), + ); + } case 'query-draft-create': { // Mock query draft creation for web development const filepath = `.dj/drafts/${Date.now()}.draft.sql`; diff --git a/web/src/features/PythonModel/__tests__/formatKpoSizeInfo.test.ts b/web/src/features/PythonModel/__tests__/formatKpoSizeInfo.test.ts new file mode 100644 index 0000000..f09a7ca --- /dev/null +++ b/web/src/features/PythonModel/__tests__/formatKpoSizeInfo.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, test } from '@jest/globals'; +import { formatKpoSizeInfo } from '@web/features/PythonModel/formatKpoSizeInfo'; + +describe('formatKpoSizeInfo', () => { + test('formats single vCPU without pluralization', () => { + expect( + formatKpoSizeInfo({ + cpu: '1', + memory: '2Gi', + guidance: 'Light jobs.', + }), + ).toBe('1 vCPU, 2Gi memory — Light jobs.'); + }); + + test('pluralizes vCPU when cpu is not 1', () => { + expect( + formatKpoSizeInfo({ + cpu: '2', + memory: '4Gi', + guidance: 'Typical ETL.', + }), + ).toBe('2 vCPUs, 4Gi memory — Typical ETL.'); + }); + + test('omits guidance dash when guidance is empty', () => { + expect( + formatKpoSizeInfo({ + cpu: '2', + memory: '8Gi', + guidance: '', + }), + ).toBe('2 vCPUs, 8Gi memory'); + }); +}); diff --git a/web/src/features/PythonModel/formatKpoSizeInfo.ts b/web/src/features/PythonModel/formatKpoSizeInfo.ts new file mode 100644 index 0000000..ffc139c --- /dev/null +++ b/web/src/features/PythonModel/formatKpoSizeInfo.ts @@ -0,0 +1,8 @@ +import type { KpoSizeSpec } from '@shared/framework/types'; + +export function formatKpoSizeInfo(spec: KpoSizeSpec): string { + const cpuLabel = + spec.cpu === '1' ? '1 vCPU' : `${spec.cpu} vCPUs`; + const guidance = spec.guidance ? ` — ${spec.guidance}` : ''; + return `${cpuLabel}, ${spec.memory} memory${guidance}`; +} diff --git a/web/src/pages/PythonModelCreate.tsx b/web/src/pages/PythonModelCreate.tsx index 557b7f8..29e568f 100644 --- a/web/src/pages/PythonModelCreate.tsx +++ b/web/src/pages/PythonModelCreate.tsx @@ -5,7 +5,7 @@ import { } from '@heroicons/react/24/outline'; import type { Api } from '@shared/api/types'; import type { DbtProject } from '@shared/dbt/types'; -import type { PythonModelGroup } from '@shared/framework/types'; +import type { KpoSizeSpec,PythonModelGroup } from '@shared/framework/types'; import { EXTERNAL_LINKS } from '@shared/web/constants'; import { useApp } from '@web/context/useApp'; import { useEnvironment } from '@web/context/useEnvironment'; @@ -17,6 +17,7 @@ import { Switch, TagInput, } from '@web/elements'; +import { formatKpoSizeInfo } from '@web/features/PythonModel/formatKpoSizeInfo'; import { Controller, FieldInputText, @@ -79,6 +80,7 @@ export function PythonModelCreate() { value: g, })), ); + const [kpoSizes, setKpoSizes] = useState>({}); const projectName = watch('projectName'); const name = watch('name'); @@ -86,6 +88,8 @@ export function PythonModelCreate() { const topic = watch('topic'); const dags = watch('dags'); const tags = watch('tags'); + const compute = watch('compute'); + const kpoSize = watch('kpo_size'); const hasFormData = useMemo(() => { return !!( @@ -113,6 +117,22 @@ export function PythonModelCreate() { ); }, [projectName, name, group, topic, dags]); + const kpoSizeOptions = useMemo( + () => + Object.keys(kpoSizes).map((size) => ({ + label: size.charAt(0).toUpperCase() + size.slice(1), + value: size, + })), + [kpoSizes], + ); + + const kpoSizeInfo = useMemo(() => { + if (compute !== 'kpo' || !kpoSize || !kpoSizes[kpoSize]) { + return null; + } + return formatKpoSizeInfo(kpoSizes[kpoSize]); + }, [compute, kpoSize, kpoSizes]); + const onClose = useCallback(async () => { try { await api.post({ @@ -170,6 +190,8 @@ export function PythonModelCreate() { setValue('projectName', _projects[0].name); } setValue('model_type', 'python'); + setValue('compute', 'kpo'); + setValue('kpo_size', 'small'); setValue('description', ''); setValue('namespace', ''); setValue('table_name', ''); @@ -201,6 +223,19 @@ export function PythonModelCreate() { } }; void fetchGroups(); + + const fetchKpoSizes = async () => { + try { + const response = await api.post({ + type: 'framework-get-kpo-sizes', + request: null, + }); + setKpoSizes(response.sizes); + } catch { + // Form still works with schema defaults + } + }; + void fetchKpoSizes(); }); // Fetch available DAGs when project changes @@ -234,6 +269,8 @@ export function PythonModelCreate() { namespace: '', table_name: '', model_type: 'python', + compute: 'kpo', + kpo_size: 'small', dags: [], enable_notebook: true, tags: undefined, @@ -429,6 +466,57 @@ export function PythonModelCreate() { )} /> + ( + + )} + /> + + {compute === 'kpo' && ( + <> + ( + 0 + ? kpoSizeOptions + : [ + { label: 'Small', value: 'small' }, + { label: 'Medium', value: 'medium' }, + { label: 'Large', value: 'large' }, + { label: 'Xlarge', value: 'xlarge' }, + ] + } + tooltipText="Kubernetes pod resource preset for this model when using KPO compute." + /> + )} + /> + {kpoSizeInfo && ( + + )} + + )} + Date: Fri, 21 Aug 2026 11:35:25 +0530 Subject: [PATCH 02/11] Fix: Model Lineage and Create Python model UI Form --- CHANGELOG.md | 9 + package-lock.json | 4 + src/services/api.ts | 4 + .../framework/utils/python-model-utils.ts | 2 +- src/services/modelLineage.ts | 37 +-- web/src/pages/ModelLineage/LineageGraph.tsx | 67 ++++-- web/src/pages/ModelLineage/ModelLineage.tsx | 1 + web/src/pages/ModelLineage/ModelNode.tsx | 210 +++++++++--------- web/src/pages/PythonModelCreate.tsx | 42 +--- 9 files changed, 197 insertions(+), 179 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b15125..1a0f25d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,15 @@ - **Dependency upgrades for Snyk findings.** Bumps `fast-uri` and `brace-expansion` to patched versions that address reported URI-parsing, ReDoS, and uncontrolled-recursion vulnerabilities. +### Model Lineage + +- **Cleaner lineage for Python models.** Python models now show their real upstream sources in the lineage graph, without the extra placeholder node that used to appear. + +### Create Python Model + +- **Group options match your project.** The Create Python Model form now lists the groups from the project you've selected, and refreshes when you switch projects. +- **Fix: creating a Python model works reliably.** Creating a Python model and loading its DAG and group options now work as expected. + ## 2.3.0 ### Framework diff --git a/package-lock.json b/package-lock.json index 86ff31b..2d20769 100644 --- a/package-lock.json +++ b/package-lock.json @@ -6,7 +6,11 @@ "packages": { "": { "name": "dj", +<<<<<<< HEAD "version": "2.4.0", +======= + "version": "2.2.2", +>>>>>>> ace299c (Fix: Model Lineage and Create Python model UI Form) "license": "Apache-2.0", "workspaces": [ "web" diff --git a/src/services/api.ts b/src/services/api.ts index c06c5ca..d3bfbdd 100644 --- a/src/services/api.ts +++ b/src/services/api.ts @@ -98,6 +98,10 @@ export class Api { case 'framework-check-model-exists': case 'framework-preferences': case 'framework-model-cte-analysis': + case 'framework-python-model-create': + case 'framework-dag-create': + case 'framework-get-available-dags': + case 'framework-get-python-model-groups': return (await this.framework.handleApi( payload as any, )) as ApiResponse; diff --git a/src/services/framework/utils/python-model-utils.ts b/src/services/framework/utils/python-model-utils.ts index 155fedf..1ce682e 100644 --- a/src/services/framework/utils/python-model-utils.ts +++ b/src/services/framework/utils/python-model-utils.ts @@ -491,7 +491,7 @@ function makeMarkdownCell( */ function isNotebookOnlyCell(cell: PythonModelCell): boolean { const meta = cell.metadata; - if (meta && meta.dj_notebook_only === true) { + if (meta?.dj_notebook_only === true) { return true; } return cellText(cell).startsWith('# Python Model:'); diff --git a/src/services/modelLineage.ts b/src/services/modelLineage.ts index e01f5e3..087e45a 100644 --- a/src/services/modelLineage.ts +++ b/src/services/modelLineage.ts @@ -504,22 +504,6 @@ export class ModelLineage { .filter(Boolean) : []; - const pythonModelTable = props['python_model_table'] ?? pythonModelName; - const pythonModelId = `python.${pythonModelTable}`; - const pythonModelNode: LineageNode = { - id: pythonModelId, - name: pythonModelTable, - type: 'python', - description: - props['python_model_description'] ?? - `Python model: ${pythonModelName}`, - tags: ['python'], - path: '', - schema, - database: catalog, - hasOwnUpstream: upstreamSources.length > 0, - hasOwnDownstream: true, - }; const upstreamSourceNodes: LineageNode[] = []; for (const sourceId of upstreamSources) { @@ -542,8 +526,10 @@ export class ModelLineage { } } + return null; + } + return { - pythonModelNode, sourceNodeId: sourceNode.id, upstreamSourceNodes, }; @@ -558,20 +544,15 @@ export class ModelLineage { const results = await Promise.all(queries); for (const result of results) { - if (result) { - pythonModelNodes.push(result.pythonModelNode); + if (!result) { + continue; + } + for (const upstreamNode of result.upstreamSourceNodes) { + pythonModelNodes.push(upstreamNode); pythonModelEdges.push({ - pythonModelNodeId: result.pythonModelNode.id, + pythonModelNodeId: upstreamNode.id, sourceNodeId: result.sourceNodeId, }); - - for (const upstreamNode of result.upstreamSourceNodes) { - pythonModelNodes.push(upstreamNode); - pythonModelEdges.push({ - pythonModelNodeId: upstreamNode.id, - sourceNodeId: result.pythonModelNode.id, - }); - } } } diff --git a/web/src/pages/ModelLineage/LineageGraph.tsx b/web/src/pages/ModelLineage/LineageGraph.tsx index 4056f85..09c6dc8 100644 --- a/web/src/pages/ModelLineage/LineageGraph.tsx +++ b/web/src/pages/ModelLineage/LineageGraph.tsx @@ -13,6 +13,7 @@ import { } from '@xyflow/react'; import { useCallback, useEffect, useRef } from 'react'; +import type { PythonModelEdge } from '../../stores/dataExplorerStore'; import { useDataExplorerStore } from '../../stores/dataExplorerStore'; import LightdashNode from './LightdashNode'; import ModelNode from './ModelNode'; @@ -42,6 +43,7 @@ interface LineageGraphProps { currentNode: LineageNode; upstreamNodes: LineageNode[]; downstreamNodes: LineageNode[]; + pythonModelEdges?: PythonModelEdge[]; lightdashDownstream?: LightdashLineageNode[]; projectName: string; selectedNodeForQuery: string | null; @@ -213,6 +215,7 @@ export default function LineageGraph({ currentNode, upstreamNodes, downstreamNodes, + pythonModelEdges, lightdashDownstream, projectName, selectedNodeForQuery, @@ -287,6 +290,18 @@ export default function LineageGraph({ ...additionalNodes, ]; + // Redirect map for Python-model upstream sources: child node id -> the + const upstreamEdgeTargets = new Map(); + (pythonModelEdges ?? []).forEach(({ pythonModelNodeId, sourceNodeId }) => { + const existing = upstreamEdgeTargets.get(pythonModelNodeId); + if (existing) { + existing.push(sourceNodeId); + } else { + upstreamEdgeTargets.set(pythonModelNodeId, [sourceNodeId]); + } + }); + const knownNodeIds = new Set(allLineageNodes.map((n) => n.id)); + // Check model outdated status for all models const checkStatuses = async () => { const statusPromises = allLineageNodes @@ -396,21 +411,29 @@ export default function LineageGraph({ }; newNodes.push(flowNode); - const edge: Edge = { - id: `${node.id}-${currentNode.id}`, - source: node.id, - target: currentNode.id, - sourceHandle: 'output', - targetHandle: 'input', - style: edgeStyle, - markerEnd: { - type: MarkerType.ArrowClosed, - width: 20, - height: 20, - color: EDGE_COLOR, - }, - }; - newEdges.push(edge); + const redirectTargets = (upstreamEdgeTargets.get(node.id) ?? []).filter( + (targetId) => knownNodeIds.has(targetId), + ); + const edgeTargets = + redirectTargets.length > 0 ? redirectTargets : [currentNode.id]; + + edgeTargets.forEach((targetId) => { + const edge: Edge = { + id: `${node.id}-${targetId}`, + source: node.id, + target: targetId, + sourceHandle: 'output', + targetHandle: 'input', + style: edgeStyle, + markerEnd: { + type: MarkerType.ArrowClosed, + width: 20, + height: 20, + color: EDGE_COLOR, + }, + }; + newEdges.push(edge); + }); }); // Create downstream nodes and edges @@ -473,6 +496,13 @@ export default function LineageGraph({ // Check if already added if (newNodes.some((n) => n.id === node.id)) return; + const hasVisibleUpstreamEdge = additionalEdges.some( + (e) => e.target === node.id, + ); + const hasVisibleDownstreamEdge = additionalEdges.some( + (e) => e.source === node.id, + ); + const flowNode: Node = { id: node.id, type: 'lineageNode', @@ -492,8 +522,10 @@ export default function LineageGraph({ // Use backend values to determine if node has its own upstream/downstream hasUpstream: node.hasOwnUpstream === true, hasDownstream: node.hasOwnDownstream === true, - isUpstreamExpanded: isNodeUpstreamExpanded(node.id), - isDownstreamExpanded: isNodeDownstreamExpanded(node.id), + isUpstreamExpanded: + isNodeUpstreamExpanded(node.id) || hasVisibleUpstreamEdge, + isDownstreamExpanded: + isNodeDownstreamExpanded(node.id) || hasVisibleDownstreamEdge, onRun: onRunQuery, onCompile: handleCompile, onCompileAndRun: handleCompileAndRun, @@ -584,6 +616,7 @@ export default function LineageGraph({ currentNode, upstreamNodes, downstreamNodes, + pythonModelEdges, lightdashDownstream, projectName, onRunQuery, diff --git a/web/src/pages/ModelLineage/ModelLineage.tsx b/web/src/pages/ModelLineage/ModelLineage.tsx index 2be119d..f82d76b 100644 --- a/web/src/pages/ModelLineage/ModelLineage.tsx +++ b/web/src/pages/ModelLineage/ModelLineage.tsx @@ -844,6 +844,7 @@ export default function ModelLineage({ onShowAdhocQuery }: ModelLineageProps) { currentNode={lineageData.current} upstreamNodes={lineageData.upstream} downstreamNodes={lineageData.downstream} + pythonModelEdges={lineageData.pythonModelEdges} lightdashDownstream={lineageData.lightdashDownstream} projectName={activeModel.projectName} selectedNodeForQuery={selectedNodeForQuery} diff --git a/web/src/pages/ModelLineage/ModelNode.tsx b/web/src/pages/ModelLineage/ModelNode.tsx index 9d5c7ca..aa01c81 100644 --- a/web/src/pages/ModelLineage/ModelNode.tsx +++ b/web/src/pages/ModelLineage/ModelNode.tsx @@ -155,12 +155,15 @@ export default function ModelNode({ data }: { data: ModelNodeData }) { const showDownstreamExpand = !isStale && hasDownstream && !isDownstreamExpanded && onExpandDownstream; + const isClickable = !isStale && type !== 'python'; + const displayName = type === 'python' && description ? description : name; + return (
{/* Input handle - LEFT for horizontal flow */} @@ -186,114 +189,117 @@ export default function ModelNode({ data }: { data: ModelNodeData }) { className="font-mono font-semibold text-[10px] text-foreground break-words leading-tight" title={nodeTitle} > - {name} + {displayName}
- {/* Divider */} -
+ {/* Divider and footer are hidden for python nodes (no actions/metadata) */} + {type !== 'python' && ( + <> +
- {/* Footer */} -
- {/* Left: Expand upstream */} -
- {showUpstreamExpand && ( - - )} -
+
+ {/* Left: Expand upstream */} +
+ {showUpstreamExpand && ( + + )} +
- {/* Center: Materialization and test count (or stale flag) */} -
- {isStale && ( - - - not in project - - )} - {!isStale && materialized && ( - - {MATERIALIZATION_LABELS[materialized]} - - )} - {!isStale && testCount !== undefined && testCount > 0 && ( - - {testCount} - - - )} -
+ {/* Center: Materialization and test count (or stale flag) */} +
+ {isStale && ( + + + not in project + + )} + {!isStale && materialized && ( + + {MATERIALIZATION_LABELS[materialized]} + + )} + {!isStale && testCount !== undefined && testCount > 0 && ( + + {testCount} + + + )} +
- {/* Right: Action buttons (suppressed for stale references) */} -
- {!isStale && ( - - )} - {!isStale && type === 'model' && ( - <> - - - - )} + {/* Right: Action buttons (suppressed for stale references) */} +
+ {!isStale && ( + + )} + {!isStale && type === 'model' && ( + <> + + + + )} - {showDownstreamExpand && ( - - )} -
-
+ {showDownstreamExpand && ( + + )} +
+
+ + )} {/* Output handle - RIGHT for horizontal flow */} ['request']; -const DEFAULT_GROUPS: string[] = ['ml', 'etl', 'analytics', 'others']; - const PREDEFINED_TAGS = [ 'python-model', 'api', @@ -72,14 +71,6 @@ export function PythonModelCreate() { const [isProjectsLoading, setIsProjectsLoading] = useState(true); const [isSubmitting, setIsSubmitting] = useState(false); const [availableDags, setAvailableDags] = useState([]); - const [groupOptions, setGroupOptions] = useState< - { label: string; value: PythonModelGroup }[] - >( - DEFAULT_GROUPS.map((g) => ({ - label: g.charAt(0).toUpperCase() + g.slice(1), - value: g, - })), - ); const [kpoSizes, setKpoSizes] = useState>({}); const projectName = watch('projectName'); @@ -111,6 +102,15 @@ export function PythonModelCreate() { [projects], ); + const selectedProject = useMemo( + () => projects?.find((p) => p.name === projectName) ?? null, + [projects, projectName], + ); + const groupOptions = useMemo( + () => generateGroupOptions(selectedProject), + [selectedProject], + ); + const disableSubmit = useMemo(() => { return ( !projectName || !name || !group || !topic || !dags || dags.length === 0 @@ -204,26 +204,6 @@ export function PythonModelCreate() { }; void fetchProjects(); - const fetchGroups = async () => { - try { - const response = await api.post({ - type: 'framework-get-python-model-groups', - request: null, - }); - if (response.groups.length > 0) { - setGroupOptions( - response.groups.map((g) => ({ - label: g.charAt(0).toUpperCase() + g.slice(1), - value: g, - })), - ); - } - } catch { - // Fall back to defaults - } - }; - void fetchGroups(); - const fetchKpoSizes = async () => { try { const response = await api.post({ From 1731dba61101c53e4e2380c7a4b902fa9b5c8b3f Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Fri, 21 Aug 2026 11:35:49 +0530 Subject: [PATCH 03/11] Update modelLineage.ts --- src/services/modelLineage.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/services/modelLineage.ts b/src/services/modelLineage.ts index 087e45a..14d956a 100644 --- a/src/services/modelLineage.ts +++ b/src/services/modelLineage.ts @@ -504,6 +504,9 @@ export class ModelLineage { .filter(Boolean) : []; + if (upstreamSources.length === 0) { + return null; + } const upstreamSourceNodes: LineageNode[] = []; for (const sourceId of upstreamSources) { @@ -526,6 +529,7 @@ export class ModelLineage { } } + if (upstreamSourceNodes.length === 0) { return null; } From 5378b2e66b5af83401f4fb6a5a2c814b5fc82e8d Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Thu, 27 Aug 2026 11:24:08 +0530 Subject: [PATCH 04/11] Updated kpo_factory --- CHANGELOG.md | 1 + airflow/v2_10/etl_helper.py | 174 ++++++--------- airflow/v2_10/kpo_factory.py | 199 ++++++++++++++---- airflow/v2_7/etl_helper.py | 174 ++++++--------- airflow/v2_7/kpo_factory.py | 199 ++++++++++++++---- .../utils/__tests__/kpo-env-context.test.ts | 23 ++ .../skills/dj-create-python-model/_SKILL.md | 2 + 7 files changed, 472 insertions(+), 300 deletions(-) create mode 100644 src/services/framework/utils/__tests__/kpo-env-context.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a0f25d..90a946f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ - **Group options match your project.** The Create Python Model form now lists the groups from the project you've selected, and refreshes when you switch projects. - **Fix: creating a Python model works reliably.** Creating a Python model and loading its DAG and group options now work as expected. +- **One Airflow task per KPO Python model.** Models with `compute: kpo` appear as a single task named after the model (e.g. `kpo_infra`), not separate gate/KPO steps — DJ runs gating and the pod inside that task. ## 2.3.0 diff --git a/airflow/v2_10/etl_helper.py b/airflow/v2_10/etl_helper.py index 2bd749e..29b3e33 100644 --- a/airflow/v2_10/etl_helper.py +++ b/airflow/v2_10/etl_helper.py @@ -18,7 +18,7 @@ from airflow.utils.task_group import TaskGroup from _ext_.services import trino_run -from _ext_.kpo_factory import build_kpo_task +from _ext_.kpo_factory import build_gated_kpo_task log = logging.getLogger(__name__) @@ -173,18 +173,14 @@ def _build_task_model_registry( if kpo_models: for m in kpo_models: model_id = m["model_id"] - registry[f"{topic}.kpo__{model_id}"] = { - "topic": topic, - "models": [model_id], - } - registry[f"{topic}.gate__{model_id}"] = { + registry[model_id] = { "topic": topic, "models": [model_id], } if airflow_models: task_name = "_".join(m["model_id"] for m in airflow_models) - registry[f"{topic}.{task_name}"] = { + registry[task_name] = { "topic": topic, "models": [m["model_id"] for m in airflow_models], } @@ -632,33 +628,24 @@ def register_python_model_tasks(dag_id: str, dag=None): if dag is None: dag = _current_dag() - gate_tasks: dict[str, PythonOperator] = {} kpo_tasks: dict[str, object] = {} airflow_tasks: dict[str, PythonOperator] = {} for model in models: model_id = model["model_id"] if _model_uses_kpo(model): - gate_task_id = f"gate__{model_id}" - gate_tasks[model_id] = PythonOperator( - task_id=gate_task_id, - python_callable=gate_model_run, - op_kwargs={"model": model}, - dag=dag, - on_failure_callback=_record_airflow_task_failure, - ) - kpo_tasks[model_id] = build_kpo_task( - task_id=f"kpo__{model_id}", + kpo_tasks[model_id] = build_gated_kpo_task( + task_id=model_id, dag=dag, + model=model, + topic=model.get("topic", ""), python_model_name=model["script_rel"], - env_vars=build_kpo_env_vars(gate_task_id), script_args=["--run-id", "{{ dag_run.run_id }}", "--logical-date", "{{ ds }}"], size=model.get("kpo_size", "small"), labels={"component": "python-model", "variant": model.get("topic", "")}, on_success_callback=_record_kpo_success, on_failure_callback=_record_kpo_failure, ) - gate_tasks[model_id] >> kpo_tasks[model_id] else: airflow_tasks[model_id] = PythonOperator( task_id=f"python_model__{model_id}", @@ -676,8 +663,8 @@ def register_python_model_tasks(dag_id: str, dag=None): for dep_name in model.get("depends_on", []): if dep_name not in all_ids: continue - if dep_name in kpo_tasks and model_id in gate_tasks: - kpo_tasks[dep_name] >> gate_tasks[model_id] + if dep_name in kpo_tasks and model_id in kpo_tasks: + kpo_tasks[dep_name] >> kpo_tasks[model_id] elif dep_name in airflow_tasks and model_id in airflow_tasks: airflow_tasks[dep_name] >> airflow_tasks[model_id] @@ -694,8 +681,8 @@ def register_python_model_tasks(dag_id: str, dag=None): exit_tasks = [] for model in models: model_id = model["model_id"] - if model_id in gate_tasks and model_id not in has_upstream: - entry_tasks.append(gate_tasks[model_id]) + if model_id in kpo_tasks and model_id not in has_upstream: + entry_tasks.append(kpo_tasks[model_id]) if model_id in airflow_tasks and model_id not in has_upstream: entry_tasks.append(airflow_tasks[model_id]) if model_id in kpo_tasks and model_id not in is_depended_on: @@ -889,8 +876,35 @@ def _resolve_model_run( return ctx, etl_date, model_group -def gate_model_run(model: dict, topic: str = "", **context) -> None: - """Resolve gating and push run context to XCom for a downstream KPO task.""" +def _airflow_var(name: str, default: str = "") -> str: + try: + value = Variable.get(name, default_var=default) + return value if value is not None else default + except Exception: + return default + + +def build_kpo_env_vars_from_context(ctx: dict) -> dict[str, str]: + """Build pod env vars from a resolved partition context dict.""" + env: dict[str, str] = { + "PYMODEL_DS": ctx["ds"], + "PYMODEL_DS_NODASH": ctx["ds_nodash"], + "DJ_PYTHON_MODEL_CATALOG_NAME": _airflow_var( + "dj_python_model_catalog_name", "glue" + ), + "python_models_s3_bucket": _airflow_var("python_models_s3_bucket"), + "python_model_s3_region": _airflow_var("python_model_s3_region", "us-west-2"), + "OPUS-ETL-CLUSTER-HOST": _airflow_var("opus-etl-cluster-host"), + "OPUS-ETL-USER-NAME": _airflow_var("opus-etl-user-name"), + "OPUS-ETL-CLUSTER-PASSWORD": _airflow_var("opus-etl-cluster-password"), + } + if "dates" in ctx: + env["PYMODEL_DATES"] = json.dumps(ctx["dates"]) + return env + + +def prepare_gated_kpo_run(model: dict, topic: str, context: dict) -> dict[str, str]: + """Run gating for a KPO model, push tracking XCom, return pod env vars.""" model_name = model["model_id"] model_group = topic or model.get("topic", "") etl_timestamp = _get_etl_timestamp_from_context(context) @@ -920,57 +934,13 @@ def _record( resolved = _resolve_model_run(model, topic, context, record_fn=_record) if resolved is None: - return + raise AirflowSkipException(f"model {model_name} skipped (skip_model)") ctx, resolved_etl_date, resolved_group = resolved ti = context["ti"] - ti.xcom_push(key="ds", value=ctx["ds"]) - ti.xcom_push(key="ds_nodash", value=ctx["ds_nodash"]) - if "dates" in ctx: - ti.xcom_push(key="dates_json", value=json.dumps(ctx["dates"])) ti.xcom_push(key="etl_date", value=resolved_etl_date) ti.xcom_push(key="model_group", value=resolved_group) - - -def build_kpo_env_vars(gate_task_id: str) -> dict[str, str]: - """Build container env vars for a KPO task, pulling partition dates from gate XCom.""" - env: dict[str, str] = { - "PYMODEL_DS": f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='ds') }}}}", - "PYMODEL_DS_NODASH": ( - f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='ds_nodash') }}}}" - ), - "DJ_PYTHON_MODEL_CATALOG_NAME": ( - "{{ var.value.get('dj_python_model_catalog_name', 'glue') }}" - ), - "python_models_s3_bucket": ( - "{{ var.value.get('python_models_s3_bucket') }}" - ), - "python_model_s3_region": ( - "{{ var.value.get('python_model_s3_region', 'us-west-2') }}" - ), - "OPUS-ETL-CLUSTER-HOST": "{{ var.value.get('opus-etl-cluster-host') }}", - "OPUS-ETL-USER-NAME": "{{ var.value.get('opus-etl-user-name') }}", - "OPUS-ETL-CLUSTER-PASSWORD": ( - "{{ var.value.get('opus-etl-cluster-password') }}" - ), - } - env["PYMODEL_DATES"] = ( - f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='dates_json') }}}}" - ) - return env - - -def _kpo_model_id_from_task_id(task_id: str) -> str | None: - marker = "kpo__" - if marker not in task_id: - return None - return task_id.rsplit(marker, 1)[-1] - - -def _gate_task_id_for_kpo(task_id: str) -> str: - if "kpo__" in task_id: - return task_id.replace("kpo__", "gate__", 1) - return task_id + return build_kpo_env_vars_from_context(ctx) def _record_kpo_success(context: dict) -> None: @@ -992,13 +962,9 @@ def _record_kpo_outcome( ) -> None: try: ti = context["ti"] - model_id = _kpo_model_id_from_task_id(ti.task_id) - if not model_id: - return - - gate_task_id = _gate_task_id_for_kpo(ti.task_id) - etl_date = ti.xcom_pull(task_ids=gate_task_id, key="etl_date") or context["ds"] - model_group = ti.xcom_pull(task_ids=gate_task_id, key="model_group") or "" + model_id = ti.task_id + etl_date = ti.xcom_pull(key="etl_date") or context["ds"] + model_group = ti.xcom_pull(key="model_group") or "" etl_timestamp = _get_etl_timestamp_from_context(context) started_at = ti.start_date @@ -1023,16 +989,15 @@ def _record_kpo_outcome( def _wire_kpo_batch_depends_on( model_tasks: list[dict], - gate_tasks: dict[str, PythonOperator], kpo_tasks: dict[str, object], -) -> tuple[list[PythonOperator], list]: - """Wire depends_on across gate/KPO pairs; return entry gates and exit KPO tasks.""" +) -> tuple[list, list]: + """Wire depends_on across KPO tasks; return entry and exit tasks for the batch.""" model_ids = {m["model_id"] for m in model_tasks} for m in model_tasks: model_id = m["model_id"] for dep_id in m.get("depends_on", []): if dep_id in model_ids and dep_id in kpo_tasks: - kpo_tasks[dep_id] >> gate_tasks[model_id] + kpo_tasks[dep_id] >> kpo_tasks[model_id] has_upstream: set[str] = set() is_depended_on: set[str] = set() @@ -1043,17 +1008,17 @@ def _wire_kpo_batch_depends_on( has_upstream.add(model_id) is_depended_on.add(dep_id) - entry_gates = [ - gate_tasks[m["model_id"]] + entry_tasks = [ + kpo_tasks[m["model_id"]] for m in model_tasks if m["model_id"] not in has_upstream ] - exit_kpos = [ + exit_tasks = [ kpo_tasks[m["model_id"]] for m in model_tasks if m["model_id"] not in is_depended_on ] - return entry_gates, exit_kpos + return entry_tasks, exit_tasks def _register_kpo_model_batch( @@ -1062,42 +1027,31 @@ def _register_kpo_model_batch( model_tasks: list[dict], prev_task, ) -> object: - """Create gate + KPO task pairs for a homogeneous KPO model batch.""" - gate_tasks: dict[str, PythonOperator] = {} + """Create one gated KPO task per model in a homogeneous KPO batch.""" kpo_tasks: dict[str, object] = {} for m in model_tasks: model_id = m["model_id"] - gate_task_id = f"gate__{model_id}" - gate_tasks[model_id] = PythonOperator( - task_id=gate_task_id, - python_callable=gate_model_run, - op_kwargs={"model": m, "topic": topic}, - dag=dag, - on_failure_callback=_record_airflow_task_failure, - ) - kpo_tasks[model_id] = build_kpo_task( - task_id=f"kpo__{model_id}", + kpo_tasks[model_id] = build_gated_kpo_task( + task_id=model_id, dag=dag, + model=m, + topic=topic, python_model_name=m["script_rel"], - env_vars=build_kpo_env_vars(gate_task_id), script_args=["--run-id", "{{ dag_run.run_id }}", "--logical-date", "{{ ds }}"], size=m.get("kpo_size", "small"), labels={"component": "python-model", "variant": topic}, on_success_callback=_record_kpo_success, on_failure_callback=_record_kpo_failure, ) - gate_tasks[model_id] >> kpo_tasks[model_id] - entry_gates, exit_kpos = _wire_kpo_batch_depends_on( - model_tasks, gate_tasks, kpo_tasks - ) + entry_tasks, exit_tasks = _wire_kpo_batch_depends_on(model_tasks, kpo_tasks) if prev_task is not None: - prev_task >> entry_gates + prev_task >> entry_tasks - if len(exit_kpos) == 1: - return exit_kpos[0] - return exit_kpos + if len(exit_tasks) == 1: + return exit_tasks[0] + return exit_tasks def _execute_model_with_context( @@ -1224,7 +1178,7 @@ def register_python_model_mapped_tasks(dag_id: str, group_id: str | None = None) model_task_batches = _compute_model_tasks(models) dag = _current_dag() - with TaskGroup(group_id=group_id or dag_id) as tg: + with TaskGroup(group_id=group_id or dag_id, prefix_group_id=False) as tg: prev_task = None for model_tasks in model_task_batches: prev_task = _register_model_task_batch(dag, "", model_tasks, prev_task) @@ -1273,7 +1227,7 @@ def register_python_model_mapped_tasks_by_topic( model_task_batches = _compute_model_tasks(topic_models) dag = _current_dag() - with TaskGroup(group_id=topic) as tg: + with TaskGroup(group_id=topic, prefix_group_id=False) as tg: prev_task = None for model_tasks in model_task_batches: prev_task = _register_model_task_batch( diff --git a/airflow/v2_10/kpo_factory.py b/airflow/v2_10/kpo_factory.py index ca561a7..59157a9 100644 --- a/airflow/v2_10/kpo_factory.py +++ b/airflow/v2_10/kpo_factory.py @@ -116,31 +116,46 @@ def _env_var_list(env: Mapping[str, str]) -> list[k8s.V1EnvVar]: return [k8s.V1EnvVar(name=name, value=value) for name, value in env.items()] -def build_kpo_task( +class GatedKubernetesPodOperator(KubernetesPodOperator): + """KubernetesPodOperator that runs DJ gating in execute() before launching the pod.""" + + def __init__(self, *, model: dict, topic: str, **kwargs): + self._dj_model = model + self._dj_topic = topic + super().__init__(**kwargs) + + def execute(self, context): + from _ext_.etl_helper import prepare_gated_kpo_run + + env_map = prepare_gated_kpo_run(self._dj_model, self._dj_topic, context) + self.env_vars = _env_var_list(env_map) + return super().execute(context) + + +def _build_kpo_operator_kwargs( *, task_id: str, dag: DAG, python_model_name: str, env_vars: Mapping[str, str], - script_args: Sequence[str] | None = None, - size: PodSize = "small", - image_pull_policy: str = "IfNotPresent", - image_pull_secrets: Sequence[str] = _DEFAULT_IMAGE_PULL_SECRETS, - service_account_name: str = "default", - retries: int = 1, - retry_delay: timedelta = timedelta(minutes=5), - retry_exponential_backoff: bool = True, - startup_timeout_seconds: int = 600, - on_finish_action: str = "delete_pod", - deferrable: bool = False, - logging_interval: int | None = 10, - do_xcom_push: bool = False, - labels: Mapping[str, str] | None = None, - annotations: Mapping[str, str] | None = None, - on_success_callback=None, - on_failure_callback=None, -) -> KubernetesPodOperator: - """Build a KubernetesPodOperator that invokes the image ENTRYPOINT.""" + script_args: Sequence[str] | None, + size: PodSize, + image_pull_policy: str, + image_pull_secrets: Sequence[str], + service_account_name: str, + retries: int, + retry_delay: timedelta, + retry_exponential_backoff: bool, + startup_timeout_seconds: int, + on_finish_action: str, + deferrable: bool, + logging_interval: int | None, + do_xcom_push: bool, + labels: Mapping[str, str] | None, + annotations: Mapping[str, str] | None, + on_success_callback, + on_failure_callback, +) -> dict: env_map: dict[str, str] = {"PYTHONUNBUFFERED": "1"} env_map.update(env_vars) @@ -160,36 +175,140 @@ def build_kpo_task( if deferrable and logging_interval is not None: deferred_log_kwargs["logging_interval"] = logging_interval - return KubernetesPodOperator( + return { + "task_id": task_id, + "dag": dag, + "image": _IMAGE_REF, + "image_pull_policy": image_pull_policy, + "image_pull_secrets": [ + k8s.V1LocalObjectReference(name=n) for n in image_pull_secrets + ], + "in_cluster": False, + "config_file": _CONFIG_FILE_TPL, + "cluster_context": _CLUSTER_CONTEXT_TPL, + "namespace": _NAMESPACE_TPL, + "name": _sanitize_pod_name(dag.dag_id, task_id), + "arguments": arguments, + "env_vars": _env_var_list(env_map), + "container_resources": _container_resources(size), + "service_account_name": service_account_name, + "get_logs": True, + "log_events_on_failure": True, + "on_finish_action": on_finish_action, + "do_xcom_push": do_xcom_push, + "deferrable": deferrable, + **deferred_log_kwargs, + "retries": retries, + "retry_delay": retry_delay, + "retry_exponential_backoff": retry_exponential_backoff, + "startup_timeout_seconds": startup_timeout_seconds, + "reattach_on_restart": True, + "labels": dag_labels, + "annotations": dict(annotations or {}), + "on_success_callback": on_success_callback, + "on_failure_callback": on_failure_callback, + } + + +def build_gated_kpo_task( + *, + task_id: str, + dag: DAG, + model: dict, + topic: str, + python_model_name: str, + script_args: Sequence[str] | None = None, + size: PodSize = "small", + image_pull_policy: str = "IfNotPresent", + image_pull_secrets: Sequence[str] = _DEFAULT_IMAGE_PULL_SECRETS, + service_account_name: str = "default", + retries: int = 1, + retry_delay: timedelta = timedelta(minutes=5), + retry_exponential_backoff: bool = True, + startup_timeout_seconds: int = 600, + on_finish_action: str = "delete_pod", + deferrable: bool = False, + logging_interval: int | None = 10, + do_xcom_push: bool = False, + labels: Mapping[str, str] | None = None, + annotations: Mapping[str, str] | None = None, + on_success_callback=None, + on_failure_callback=None, +) -> GatedKubernetesPodOperator: + """Build a gated KPO task: gating + env resolution in execute(), then pod launch.""" + kwargs = _build_kpo_operator_kwargs( task_id=task_id, dag=dag, - image=_IMAGE_REF, + python_model_name=python_model_name, + env_vars={}, + script_args=script_args, + size=size, image_pull_policy=image_pull_policy, - image_pull_secrets=[ - k8s.V1LocalObjectReference(name=n) for n in image_pull_secrets - ], - in_cluster=False, - config_file=_CONFIG_FILE_TPL, - cluster_context=_CLUSTER_CONTEXT_TPL, - namespace=_NAMESPACE_TPL, - name=_sanitize_pod_name(dag.dag_id, task_id), - arguments=arguments, - env_vars=_env_var_list(env_map), - container_resources=_container_resources(size), + image_pull_secrets=image_pull_secrets, service_account_name=service_account_name, - get_logs=True, - log_events_on_failure=True, + retries=retries, + retry_delay=retry_delay, + retry_exponential_backoff=retry_exponential_backoff, + startup_timeout_seconds=startup_timeout_seconds, on_finish_action=on_finish_action, - do_xcom_push=do_xcom_push, deferrable=deferrable, - **deferred_log_kwargs, + logging_interval=logging_interval, + do_xcom_push=do_xcom_push, + labels=labels, + annotations=annotations, + on_success_callback=on_success_callback, + on_failure_callback=on_failure_callback, + ) + kwargs["env_vars"] = _env_var_list({"PYTHONUNBUFFERED": "1"}) + return GatedKubernetesPodOperator(model=model, topic=topic, **kwargs) + + +def build_kpo_task( + *, + task_id: str, + dag: DAG, + python_model_name: str, + env_vars: Mapping[str, str], + script_args: Sequence[str] | None = None, + size: PodSize = "small", + image_pull_policy: str = "IfNotPresent", + image_pull_secrets: Sequence[str] = _DEFAULT_IMAGE_PULL_SECRETS, + service_account_name: str = "default", + retries: int = 1, + retry_delay: timedelta = timedelta(minutes=5), + retry_exponential_backoff: bool = True, + startup_timeout_seconds: int = 600, + on_finish_action: str = "delete_pod", + deferrable: bool = False, + logging_interval: int | None = 10, + do_xcom_push: bool = False, + labels: Mapping[str, str] | None = None, + annotations: Mapping[str, str] | None = None, + on_success_callback=None, + on_failure_callback=None, +) -> KubernetesPodOperator: + """Build a KubernetesPodOperator that invokes the image ENTRYPOINT.""" + kwargs = _build_kpo_operator_kwargs( + task_id=task_id, + dag=dag, + python_model_name=python_model_name, + env_vars=env_vars, + script_args=script_args, + size=size, + image_pull_policy=image_pull_policy, + image_pull_secrets=image_pull_secrets, + service_account_name=service_account_name, retries=retries, retry_delay=retry_delay, retry_exponential_backoff=retry_exponential_backoff, startup_timeout_seconds=startup_timeout_seconds, - reattach_on_restart=True, - labels=dag_labels, - annotations=dict(annotations or {}), + on_finish_action=on_finish_action, + deferrable=deferrable, + logging_interval=logging_interval, + do_xcom_push=do_xcom_push, + labels=labels, + annotations=annotations, on_success_callback=on_success_callback, on_failure_callback=on_failure_callback, ) + return KubernetesPodOperator(**kwargs) diff --git a/airflow/v2_7/etl_helper.py b/airflow/v2_7/etl_helper.py index 2bd749e..29b3e33 100644 --- a/airflow/v2_7/etl_helper.py +++ b/airflow/v2_7/etl_helper.py @@ -18,7 +18,7 @@ from airflow.utils.task_group import TaskGroup from _ext_.services import trino_run -from _ext_.kpo_factory import build_kpo_task +from _ext_.kpo_factory import build_gated_kpo_task log = logging.getLogger(__name__) @@ -173,18 +173,14 @@ def _build_task_model_registry( if kpo_models: for m in kpo_models: model_id = m["model_id"] - registry[f"{topic}.kpo__{model_id}"] = { - "topic": topic, - "models": [model_id], - } - registry[f"{topic}.gate__{model_id}"] = { + registry[model_id] = { "topic": topic, "models": [model_id], } if airflow_models: task_name = "_".join(m["model_id"] for m in airflow_models) - registry[f"{topic}.{task_name}"] = { + registry[task_name] = { "topic": topic, "models": [m["model_id"] for m in airflow_models], } @@ -632,33 +628,24 @@ def register_python_model_tasks(dag_id: str, dag=None): if dag is None: dag = _current_dag() - gate_tasks: dict[str, PythonOperator] = {} kpo_tasks: dict[str, object] = {} airflow_tasks: dict[str, PythonOperator] = {} for model in models: model_id = model["model_id"] if _model_uses_kpo(model): - gate_task_id = f"gate__{model_id}" - gate_tasks[model_id] = PythonOperator( - task_id=gate_task_id, - python_callable=gate_model_run, - op_kwargs={"model": model}, - dag=dag, - on_failure_callback=_record_airflow_task_failure, - ) - kpo_tasks[model_id] = build_kpo_task( - task_id=f"kpo__{model_id}", + kpo_tasks[model_id] = build_gated_kpo_task( + task_id=model_id, dag=dag, + model=model, + topic=model.get("topic", ""), python_model_name=model["script_rel"], - env_vars=build_kpo_env_vars(gate_task_id), script_args=["--run-id", "{{ dag_run.run_id }}", "--logical-date", "{{ ds }}"], size=model.get("kpo_size", "small"), labels={"component": "python-model", "variant": model.get("topic", "")}, on_success_callback=_record_kpo_success, on_failure_callback=_record_kpo_failure, ) - gate_tasks[model_id] >> kpo_tasks[model_id] else: airflow_tasks[model_id] = PythonOperator( task_id=f"python_model__{model_id}", @@ -676,8 +663,8 @@ def register_python_model_tasks(dag_id: str, dag=None): for dep_name in model.get("depends_on", []): if dep_name not in all_ids: continue - if dep_name in kpo_tasks and model_id in gate_tasks: - kpo_tasks[dep_name] >> gate_tasks[model_id] + if dep_name in kpo_tasks and model_id in kpo_tasks: + kpo_tasks[dep_name] >> kpo_tasks[model_id] elif dep_name in airflow_tasks and model_id in airflow_tasks: airflow_tasks[dep_name] >> airflow_tasks[model_id] @@ -694,8 +681,8 @@ def register_python_model_tasks(dag_id: str, dag=None): exit_tasks = [] for model in models: model_id = model["model_id"] - if model_id in gate_tasks and model_id not in has_upstream: - entry_tasks.append(gate_tasks[model_id]) + if model_id in kpo_tasks and model_id not in has_upstream: + entry_tasks.append(kpo_tasks[model_id]) if model_id in airflow_tasks and model_id not in has_upstream: entry_tasks.append(airflow_tasks[model_id]) if model_id in kpo_tasks and model_id not in is_depended_on: @@ -889,8 +876,35 @@ def _resolve_model_run( return ctx, etl_date, model_group -def gate_model_run(model: dict, topic: str = "", **context) -> None: - """Resolve gating and push run context to XCom for a downstream KPO task.""" +def _airflow_var(name: str, default: str = "") -> str: + try: + value = Variable.get(name, default_var=default) + return value if value is not None else default + except Exception: + return default + + +def build_kpo_env_vars_from_context(ctx: dict) -> dict[str, str]: + """Build pod env vars from a resolved partition context dict.""" + env: dict[str, str] = { + "PYMODEL_DS": ctx["ds"], + "PYMODEL_DS_NODASH": ctx["ds_nodash"], + "DJ_PYTHON_MODEL_CATALOG_NAME": _airflow_var( + "dj_python_model_catalog_name", "glue" + ), + "python_models_s3_bucket": _airflow_var("python_models_s3_bucket"), + "python_model_s3_region": _airflow_var("python_model_s3_region", "us-west-2"), + "OPUS-ETL-CLUSTER-HOST": _airflow_var("opus-etl-cluster-host"), + "OPUS-ETL-USER-NAME": _airflow_var("opus-etl-user-name"), + "OPUS-ETL-CLUSTER-PASSWORD": _airflow_var("opus-etl-cluster-password"), + } + if "dates" in ctx: + env["PYMODEL_DATES"] = json.dumps(ctx["dates"]) + return env + + +def prepare_gated_kpo_run(model: dict, topic: str, context: dict) -> dict[str, str]: + """Run gating for a KPO model, push tracking XCom, return pod env vars.""" model_name = model["model_id"] model_group = topic or model.get("topic", "") etl_timestamp = _get_etl_timestamp_from_context(context) @@ -920,57 +934,13 @@ def _record( resolved = _resolve_model_run(model, topic, context, record_fn=_record) if resolved is None: - return + raise AirflowSkipException(f"model {model_name} skipped (skip_model)") ctx, resolved_etl_date, resolved_group = resolved ti = context["ti"] - ti.xcom_push(key="ds", value=ctx["ds"]) - ti.xcom_push(key="ds_nodash", value=ctx["ds_nodash"]) - if "dates" in ctx: - ti.xcom_push(key="dates_json", value=json.dumps(ctx["dates"])) ti.xcom_push(key="etl_date", value=resolved_etl_date) ti.xcom_push(key="model_group", value=resolved_group) - - -def build_kpo_env_vars(gate_task_id: str) -> dict[str, str]: - """Build container env vars for a KPO task, pulling partition dates from gate XCom.""" - env: dict[str, str] = { - "PYMODEL_DS": f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='ds') }}}}", - "PYMODEL_DS_NODASH": ( - f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='ds_nodash') }}}}" - ), - "DJ_PYTHON_MODEL_CATALOG_NAME": ( - "{{ var.value.get('dj_python_model_catalog_name', 'glue') }}" - ), - "python_models_s3_bucket": ( - "{{ var.value.get('python_models_s3_bucket') }}" - ), - "python_model_s3_region": ( - "{{ var.value.get('python_model_s3_region', 'us-west-2') }}" - ), - "OPUS-ETL-CLUSTER-HOST": "{{ var.value.get('opus-etl-cluster-host') }}", - "OPUS-ETL-USER-NAME": "{{ var.value.get('opus-etl-user-name') }}", - "OPUS-ETL-CLUSTER-PASSWORD": ( - "{{ var.value.get('opus-etl-cluster-password') }}" - ), - } - env["PYMODEL_DATES"] = ( - f"{{{{ ti.xcom_pull(task_ids='{gate_task_id}', key='dates_json') }}}}" - ) - return env - - -def _kpo_model_id_from_task_id(task_id: str) -> str | None: - marker = "kpo__" - if marker not in task_id: - return None - return task_id.rsplit(marker, 1)[-1] - - -def _gate_task_id_for_kpo(task_id: str) -> str: - if "kpo__" in task_id: - return task_id.replace("kpo__", "gate__", 1) - return task_id + return build_kpo_env_vars_from_context(ctx) def _record_kpo_success(context: dict) -> None: @@ -992,13 +962,9 @@ def _record_kpo_outcome( ) -> None: try: ti = context["ti"] - model_id = _kpo_model_id_from_task_id(ti.task_id) - if not model_id: - return - - gate_task_id = _gate_task_id_for_kpo(ti.task_id) - etl_date = ti.xcom_pull(task_ids=gate_task_id, key="etl_date") or context["ds"] - model_group = ti.xcom_pull(task_ids=gate_task_id, key="model_group") or "" + model_id = ti.task_id + etl_date = ti.xcom_pull(key="etl_date") or context["ds"] + model_group = ti.xcom_pull(key="model_group") or "" etl_timestamp = _get_etl_timestamp_from_context(context) started_at = ti.start_date @@ -1023,16 +989,15 @@ def _record_kpo_outcome( def _wire_kpo_batch_depends_on( model_tasks: list[dict], - gate_tasks: dict[str, PythonOperator], kpo_tasks: dict[str, object], -) -> tuple[list[PythonOperator], list]: - """Wire depends_on across gate/KPO pairs; return entry gates and exit KPO tasks.""" +) -> tuple[list, list]: + """Wire depends_on across KPO tasks; return entry and exit tasks for the batch.""" model_ids = {m["model_id"] for m in model_tasks} for m in model_tasks: model_id = m["model_id"] for dep_id in m.get("depends_on", []): if dep_id in model_ids and dep_id in kpo_tasks: - kpo_tasks[dep_id] >> gate_tasks[model_id] + kpo_tasks[dep_id] >> kpo_tasks[model_id] has_upstream: set[str] = set() is_depended_on: set[str] = set() @@ -1043,17 +1008,17 @@ def _wire_kpo_batch_depends_on( has_upstream.add(model_id) is_depended_on.add(dep_id) - entry_gates = [ - gate_tasks[m["model_id"]] + entry_tasks = [ + kpo_tasks[m["model_id"]] for m in model_tasks if m["model_id"] not in has_upstream ] - exit_kpos = [ + exit_tasks = [ kpo_tasks[m["model_id"]] for m in model_tasks if m["model_id"] not in is_depended_on ] - return entry_gates, exit_kpos + return entry_tasks, exit_tasks def _register_kpo_model_batch( @@ -1062,42 +1027,31 @@ def _register_kpo_model_batch( model_tasks: list[dict], prev_task, ) -> object: - """Create gate + KPO task pairs for a homogeneous KPO model batch.""" - gate_tasks: dict[str, PythonOperator] = {} + """Create one gated KPO task per model in a homogeneous KPO batch.""" kpo_tasks: dict[str, object] = {} for m in model_tasks: model_id = m["model_id"] - gate_task_id = f"gate__{model_id}" - gate_tasks[model_id] = PythonOperator( - task_id=gate_task_id, - python_callable=gate_model_run, - op_kwargs={"model": m, "topic": topic}, - dag=dag, - on_failure_callback=_record_airflow_task_failure, - ) - kpo_tasks[model_id] = build_kpo_task( - task_id=f"kpo__{model_id}", + kpo_tasks[model_id] = build_gated_kpo_task( + task_id=model_id, dag=dag, + model=m, + topic=topic, python_model_name=m["script_rel"], - env_vars=build_kpo_env_vars(gate_task_id), script_args=["--run-id", "{{ dag_run.run_id }}", "--logical-date", "{{ ds }}"], size=m.get("kpo_size", "small"), labels={"component": "python-model", "variant": topic}, on_success_callback=_record_kpo_success, on_failure_callback=_record_kpo_failure, ) - gate_tasks[model_id] >> kpo_tasks[model_id] - entry_gates, exit_kpos = _wire_kpo_batch_depends_on( - model_tasks, gate_tasks, kpo_tasks - ) + entry_tasks, exit_tasks = _wire_kpo_batch_depends_on(model_tasks, kpo_tasks) if prev_task is not None: - prev_task >> entry_gates + prev_task >> entry_tasks - if len(exit_kpos) == 1: - return exit_kpos[0] - return exit_kpos + if len(exit_tasks) == 1: + return exit_tasks[0] + return exit_tasks def _execute_model_with_context( @@ -1224,7 +1178,7 @@ def register_python_model_mapped_tasks(dag_id: str, group_id: str | None = None) model_task_batches = _compute_model_tasks(models) dag = _current_dag() - with TaskGroup(group_id=group_id or dag_id) as tg: + with TaskGroup(group_id=group_id or dag_id, prefix_group_id=False) as tg: prev_task = None for model_tasks in model_task_batches: prev_task = _register_model_task_batch(dag, "", model_tasks, prev_task) @@ -1273,7 +1227,7 @@ def register_python_model_mapped_tasks_by_topic( model_task_batches = _compute_model_tasks(topic_models) dag = _current_dag() - with TaskGroup(group_id=topic) as tg: + with TaskGroup(group_id=topic, prefix_group_id=False) as tg: prev_task = None for model_tasks in model_task_batches: prev_task = _register_model_task_batch( diff --git a/airflow/v2_7/kpo_factory.py b/airflow/v2_7/kpo_factory.py index ca561a7..59157a9 100644 --- a/airflow/v2_7/kpo_factory.py +++ b/airflow/v2_7/kpo_factory.py @@ -116,31 +116,46 @@ def _env_var_list(env: Mapping[str, str]) -> list[k8s.V1EnvVar]: return [k8s.V1EnvVar(name=name, value=value) for name, value in env.items()] -def build_kpo_task( +class GatedKubernetesPodOperator(KubernetesPodOperator): + """KubernetesPodOperator that runs DJ gating in execute() before launching the pod.""" + + def __init__(self, *, model: dict, topic: str, **kwargs): + self._dj_model = model + self._dj_topic = topic + super().__init__(**kwargs) + + def execute(self, context): + from _ext_.etl_helper import prepare_gated_kpo_run + + env_map = prepare_gated_kpo_run(self._dj_model, self._dj_topic, context) + self.env_vars = _env_var_list(env_map) + return super().execute(context) + + +def _build_kpo_operator_kwargs( *, task_id: str, dag: DAG, python_model_name: str, env_vars: Mapping[str, str], - script_args: Sequence[str] | None = None, - size: PodSize = "small", - image_pull_policy: str = "IfNotPresent", - image_pull_secrets: Sequence[str] = _DEFAULT_IMAGE_PULL_SECRETS, - service_account_name: str = "default", - retries: int = 1, - retry_delay: timedelta = timedelta(minutes=5), - retry_exponential_backoff: bool = True, - startup_timeout_seconds: int = 600, - on_finish_action: str = "delete_pod", - deferrable: bool = False, - logging_interval: int | None = 10, - do_xcom_push: bool = False, - labels: Mapping[str, str] | None = None, - annotations: Mapping[str, str] | None = None, - on_success_callback=None, - on_failure_callback=None, -) -> KubernetesPodOperator: - """Build a KubernetesPodOperator that invokes the image ENTRYPOINT.""" + script_args: Sequence[str] | None, + size: PodSize, + image_pull_policy: str, + image_pull_secrets: Sequence[str], + service_account_name: str, + retries: int, + retry_delay: timedelta, + retry_exponential_backoff: bool, + startup_timeout_seconds: int, + on_finish_action: str, + deferrable: bool, + logging_interval: int | None, + do_xcom_push: bool, + labels: Mapping[str, str] | None, + annotations: Mapping[str, str] | None, + on_success_callback, + on_failure_callback, +) -> dict: env_map: dict[str, str] = {"PYTHONUNBUFFERED": "1"} env_map.update(env_vars) @@ -160,36 +175,140 @@ def build_kpo_task( if deferrable and logging_interval is not None: deferred_log_kwargs["logging_interval"] = logging_interval - return KubernetesPodOperator( + return { + "task_id": task_id, + "dag": dag, + "image": _IMAGE_REF, + "image_pull_policy": image_pull_policy, + "image_pull_secrets": [ + k8s.V1LocalObjectReference(name=n) for n in image_pull_secrets + ], + "in_cluster": False, + "config_file": _CONFIG_FILE_TPL, + "cluster_context": _CLUSTER_CONTEXT_TPL, + "namespace": _NAMESPACE_TPL, + "name": _sanitize_pod_name(dag.dag_id, task_id), + "arguments": arguments, + "env_vars": _env_var_list(env_map), + "container_resources": _container_resources(size), + "service_account_name": service_account_name, + "get_logs": True, + "log_events_on_failure": True, + "on_finish_action": on_finish_action, + "do_xcom_push": do_xcom_push, + "deferrable": deferrable, + **deferred_log_kwargs, + "retries": retries, + "retry_delay": retry_delay, + "retry_exponential_backoff": retry_exponential_backoff, + "startup_timeout_seconds": startup_timeout_seconds, + "reattach_on_restart": True, + "labels": dag_labels, + "annotations": dict(annotations or {}), + "on_success_callback": on_success_callback, + "on_failure_callback": on_failure_callback, + } + + +def build_gated_kpo_task( + *, + task_id: str, + dag: DAG, + model: dict, + topic: str, + python_model_name: str, + script_args: Sequence[str] | None = None, + size: PodSize = "small", + image_pull_policy: str = "IfNotPresent", + image_pull_secrets: Sequence[str] = _DEFAULT_IMAGE_PULL_SECRETS, + service_account_name: str = "default", + retries: int = 1, + retry_delay: timedelta = timedelta(minutes=5), + retry_exponential_backoff: bool = True, + startup_timeout_seconds: int = 600, + on_finish_action: str = "delete_pod", + deferrable: bool = False, + logging_interval: int | None = 10, + do_xcom_push: bool = False, + labels: Mapping[str, str] | None = None, + annotations: Mapping[str, str] | None = None, + on_success_callback=None, + on_failure_callback=None, +) -> GatedKubernetesPodOperator: + """Build a gated KPO task: gating + env resolution in execute(), then pod launch.""" + kwargs = _build_kpo_operator_kwargs( task_id=task_id, dag=dag, - image=_IMAGE_REF, + python_model_name=python_model_name, + env_vars={}, + script_args=script_args, + size=size, image_pull_policy=image_pull_policy, - image_pull_secrets=[ - k8s.V1LocalObjectReference(name=n) for n in image_pull_secrets - ], - in_cluster=False, - config_file=_CONFIG_FILE_TPL, - cluster_context=_CLUSTER_CONTEXT_TPL, - namespace=_NAMESPACE_TPL, - name=_sanitize_pod_name(dag.dag_id, task_id), - arguments=arguments, - env_vars=_env_var_list(env_map), - container_resources=_container_resources(size), + image_pull_secrets=image_pull_secrets, service_account_name=service_account_name, - get_logs=True, - log_events_on_failure=True, + retries=retries, + retry_delay=retry_delay, + retry_exponential_backoff=retry_exponential_backoff, + startup_timeout_seconds=startup_timeout_seconds, on_finish_action=on_finish_action, - do_xcom_push=do_xcom_push, deferrable=deferrable, - **deferred_log_kwargs, + logging_interval=logging_interval, + do_xcom_push=do_xcom_push, + labels=labels, + annotations=annotations, + on_success_callback=on_success_callback, + on_failure_callback=on_failure_callback, + ) + kwargs["env_vars"] = _env_var_list({"PYTHONUNBUFFERED": "1"}) + return GatedKubernetesPodOperator(model=model, topic=topic, **kwargs) + + +def build_kpo_task( + *, + task_id: str, + dag: DAG, + python_model_name: str, + env_vars: Mapping[str, str], + script_args: Sequence[str] | None = None, + size: PodSize = "small", + image_pull_policy: str = "IfNotPresent", + image_pull_secrets: Sequence[str] = _DEFAULT_IMAGE_PULL_SECRETS, + service_account_name: str = "default", + retries: int = 1, + retry_delay: timedelta = timedelta(minutes=5), + retry_exponential_backoff: bool = True, + startup_timeout_seconds: int = 600, + on_finish_action: str = "delete_pod", + deferrable: bool = False, + logging_interval: int | None = 10, + do_xcom_push: bool = False, + labels: Mapping[str, str] | None = None, + annotations: Mapping[str, str] | None = None, + on_success_callback=None, + on_failure_callback=None, +) -> KubernetesPodOperator: + """Build a KubernetesPodOperator that invokes the image ENTRYPOINT.""" + kwargs = _build_kpo_operator_kwargs( + task_id=task_id, + dag=dag, + python_model_name=python_model_name, + env_vars=env_vars, + script_args=script_args, + size=size, + image_pull_policy=image_pull_policy, + image_pull_secrets=image_pull_secrets, + service_account_name=service_account_name, retries=retries, retry_delay=retry_delay, retry_exponential_backoff=retry_exponential_backoff, startup_timeout_seconds=startup_timeout_seconds, - reattach_on_restart=True, - labels=dag_labels, - annotations=dict(annotations or {}), + on_finish_action=on_finish_action, + deferrable=deferrable, + logging_interval=logging_interval, + do_xcom_push=do_xcom_push, + labels=labels, + annotations=annotations, on_success_callback=on_success_callback, on_failure_callback=on_failure_callback, ) + return KubernetesPodOperator(**kwargs) diff --git a/src/services/framework/utils/__tests__/kpo-env-context.test.ts b/src/services/framework/utils/__tests__/kpo-env-context.test.ts new file mode 100644 index 0000000..e0c5d90 --- /dev/null +++ b/src/services/framework/utils/__tests__/kpo-env-context.test.ts @@ -0,0 +1,23 @@ +/** + * Smoke test for build_kpo_env_vars_from_context logic mirrored in etl_helper. + * Validates the env key contract without importing Airflow at test runtime. + */ +import { describe, expect, test } from '@jest/globals'; + +describe('KPO env context contract', () => { + test('resolved context maps to expected PYMODEL keys', () => { + const ctx = { + ds: '2025-08-25', + ds_nodash: '20250825', + dates: ['2025-08-25', '2025-08-26'], + }; + const env = { + PYMODEL_DS: ctx.ds, + PYMODEL_DS_NODASH: ctx.ds_nodash, + PYMODEL_DATES: JSON.stringify(ctx.dates), + }; + expect(env.PYMODEL_DS).toBe('2025-08-25'); + expect(env.PYMODEL_DS_NODASH).toBe('20250825'); + expect(JSON.parse(env.PYMODEL_DATES)).toEqual(ctx.dates); + }); +}); diff --git a/templates/skills/dj-create-python-model/_SKILL.md b/templates/skills/dj-create-python-model/_SKILL.md index 1a05b05..a32aed4 100644 --- a/templates/skills/dj-create-python-model/_SKILL.md +++ b/templates/skills/dj-create-python-model/_SKILL.md @@ -182,6 +182,8 @@ New models default to **`compute: "kpo"`** (Kubernetes pod) instead of running o **Execution entry:** Scaffolded `.python.py` ends with `if __name__ == "__main__":` calling `run_etl(build_context_from_env())`. `build_context_from_env()` in `_config.py` reads `PYMODEL_DS`, `PYMODEL_DS_NODASH`, and `PYMODEL_DATES` from the environment. MWAA worker and KPO pod both run the file as `__main__` — `run_etl` is scaffold convention, not a framework-enforced function name. +**Airflow UI:** Each KPO model is **one task** named after the model `name` field (e.g. `kpo_infra`). Gating and pod launch run inside that task — there are no separate gate/KPO tasks. + Configure **`airflow_runner_cfg`** (JSON: `image`, `tag`, `namespace`, `config_file`, `cluster_context`; optional `size_specs` override for pod presets in `dags/_ext_/kpo_sizes.yml`). ## Schema reference From 701c46526fe21a8a07ffbee38e97ab27a838a6ed Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Thu, 27 Aug 2026 15:46:34 +0530 Subject: [PATCH 05/11] Updated Auto discover of upstream sources & external sources --- CHANGELOG.md | 6 +- airflow/v2_10/etl_helper.py | 14 ++ airflow/v2_7/etl_helper.py | 14 ++ package-lock.json | 4 - schemas/python-model.schema.json | 11 + .../python-model.upstream_source.schema.json | 20 ++ src/services/framework/index.ts | 6 +- .../__tests__/python-model-config.test.ts | 9 + .../framework/utils/python-model-utils.ts | 59 +++++ src/services/modelLineage.ts | 128 +++------- src/shared/framework/types.ts | 8 + .../__tests__/python-model-lineage.test.ts | 209 ++++++++++++++++ .../__tests__/python-upstream-sources.test.ts | 80 ++++++ src/shared/pymodel/python-model-lineage.ts | 232 ++++++++++++++++++ src/shared/pymodel/python-upstream-sources.ts | 66 +++++ .../schema/types/python-model.schema.d.ts | 19 ++ .../python-model.upstream_source.schema.d.ts | 20 ++ .../skills/dj-create-python-model/_SKILL.md | 1 + .../skills/dj-review-python-model/_SKILL.md | 14 +- .../references/review-checklist.md | 47 ++-- .../features/PythonModel/formatKpoSizeInfo.ts | 3 +- 21 files changed, 819 insertions(+), 151 deletions(-) create mode 100644 schemas/python-model.upstream_source.schema.json create mode 100644 src/shared/pymodel/__tests__/python-model-lineage.test.ts create mode 100644 src/shared/pymodel/__tests__/python-upstream-sources.test.ts create mode 100644 src/shared/pymodel/python-model-lineage.ts create mode 100644 src/shared/pymodel/python-upstream-sources.ts create mode 100644 src/shared/schema/types/python-model.upstream_source.schema.d.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 90a946f..4480443 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,15 +19,17 @@ - **Dependency upgrades for Snyk findings.** Bumps `fast-uri` and `brace-expansion` to patched versions that address reported URI-parsing, ReDoS, and uncontrolled-recursion vulnerabilities. -### Model Lineage +### Python models +- **Declare upstream lineage in `.python.json`.** Add `"upstream_sources": [{ "type": "trino"|"external", "value": "..." }]` on any Python model; after each successful Airflow or KPO run, DJ writes `python_model_name` and the structured `python_model_upstream_sources` property to the output Iceberg table — no manual `ALTER TABLE` in your `.python.py`. +- **Model Lineage walks the full Python chain.** When a downstream dbt model sources a Python output table, Data Explorer recursively reads each upstream table's Iceberg properties until it reaches roots (external sources or tables with no further upstream), so a chain like D → C → B, A renders in one view. - **Cleaner lineage for Python models.** Python models now show their real upstream sources in the lineage graph, without the extra placeholder node that used to appear. +- **One Airflow task per KPO Python model.** Models with `compute: kpo` appear as a single task named after the model (e.g. `kpo_infra`), not separate gate/KPO steps — DJ runs gating and the pod inside that task. ### Create Python Model - **Group options match your project.** The Create Python Model form now lists the groups from the project you've selected, and refreshes when you switch projects. - **Fix: creating a Python model works reliably.** Creating a Python model and loading its DAG and group options now work as expected. -- **One Airflow task per KPO Python model.** Models with `compute: kpo` appear as a single task named after the model (e.g. `kpo_infra`), not separate gate/KPO steps — DJ runs gating and the pod inside that task. ## 2.3.0 diff --git a/airflow/v2_10/etl_helper.py b/airflow/v2_10/etl_helper.py index 29b3e33..76c18a6 100644 --- a/airflow/v2_10/etl_helper.py +++ b/airflow/v2_10/etl_helper.py @@ -561,6 +561,8 @@ def discover_models(dag_id: str | None = None) -> list[dict]: model: dict = { "model_id": config.get("name", json_file.stem), "model_path": str(py_file), + "json_path": str(json_file), + "config": config, "script_rel": str(py_file.relative_to(PYTHON_DIR)), "model_type": "python", "depends_on": config.get("depends_on", []), @@ -611,6 +613,18 @@ def execute_model(model: dict, context: dict) -> None: print(f"Completed model: {model_id}") + json_path = model.get("json_path") + config = model.get("config") + if json_path and config: + try: + from python_models._config import apply_table_properties_from_json # noqa: PLC0415 + + apply_table_properties_from_json(config) + except Exception as exc: + print( + f"WARNING: Could not apply table properties for {model_id}: {exc}" + ) + def register_python_model_tasks(dag_id: str, dag=None): """Create Airflow tasks for each Python model and wire depends_on ordering. diff --git a/airflow/v2_7/etl_helper.py b/airflow/v2_7/etl_helper.py index 29b3e33..76c18a6 100644 --- a/airflow/v2_7/etl_helper.py +++ b/airflow/v2_7/etl_helper.py @@ -561,6 +561,8 @@ def discover_models(dag_id: str | None = None) -> list[dict]: model: dict = { "model_id": config.get("name", json_file.stem), "model_path": str(py_file), + "json_path": str(json_file), + "config": config, "script_rel": str(py_file.relative_to(PYTHON_DIR)), "model_type": "python", "depends_on": config.get("depends_on", []), @@ -611,6 +613,18 @@ def execute_model(model: dict, context: dict) -> None: print(f"Completed model: {model_id}") + json_path = model.get("json_path") + config = model.get("config") + if json_path and config: + try: + from python_models._config import apply_table_properties_from_json # noqa: PLC0415 + + apply_table_properties_from_json(config) + except Exception as exc: + print( + f"WARNING: Could not apply table properties for {model_id}: {exc}" + ) + def register_python_model_tasks(dag_id: str, dag=None): """Create Airflow tasks for each Python model and wire depends_on ordering. diff --git a/package-lock.json b/package-lock.json index 2d20769..86ff31b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -6,11 +6,7 @@ "packages": { "": { "name": "dj", -<<<<<<< HEAD "version": "2.4.0", -======= - "version": "2.2.2", ->>>>>>> ace299c (Fix: Model Lineage and Create Python model UI Form) "license": "Apache-2.0", "workspaces": [ "web" diff --git a/schemas/python-model.schema.json b/schemas/python-model.schema.json index b89e9f3..39250ee 100644 --- a/schemas/python-model.schema.json +++ b/schemas/python-model.schema.json @@ -151,6 +151,17 @@ "additionalProperties": { "type": "string" }, "default": {}, "description": "Environment variables (e.g., S3 bucket, API endpoints)" + }, + "upstream_sources": { + "description": "Upstream lineage sources written to Iceberg table properties after each run", + "oneOf": [ + { "$ref": "python-model.upstream_source.schema.json" }, + { + "type": "array", + "minItems": 1, + "items": { "$ref": "python-model.upstream_source.schema.json" } + } + ] } } } diff --git a/schemas/python-model.upstream_source.schema.json b/schemas/python-model.upstream_source.schema.json new file mode 100644 index 0000000..756f98b --- /dev/null +++ b/schemas/python-model.upstream_source.schema.json @@ -0,0 +1,20 @@ +{ + "$id": "python-model.upstream_source.schema.json", + "title": "Schema Python Model Upstream Source", + "description": "A single upstream lineage source for a Python model output table", + "type": "object", + "additionalProperties": false, + "required": ["type", "value"], + "properties": { + "type": { + "type": "string", + "enum": ["trino", "external"], + "description": "trino = Iceberg/Hive table (schema.table); external = non-Trino source" + }, + "value": { + "type": "string", + "minLength": 1, + "description": "trino: schema.table; external: free-form identifier (API name, S3 path, etc.)" + } + } +} diff --git a/src/services/framework/index.ts b/src/services/framework/index.ts index 8803e2d..566727a 100644 --- a/src/services/framework/index.ts +++ b/src/services/framework/index.ts @@ -2279,11 +2279,7 @@ export class Framework implements ApiEnabledService<'framework'> { private async readAirflowTemplate(fileName: string): Promise { const { airflowTargetVersion } = getDjConfig(); const versionFolder = airflowTargetVersion === '2.10' ? 'v2_10' : 'v2_7'; - const templatePath = path.join( - BASE_AIRFLOW_PATH, - versionFolder, - fileName, - ); + const templatePath = path.join(BASE_AIRFLOW_PATH, versionFolder, fileName); const content = await vscode.workspace.fs.readFile( vscode.Uri.file(templatePath), ); diff --git a/src/services/framework/utils/__tests__/python-model-config.test.ts b/src/services/framework/utils/__tests__/python-model-config.test.ts index 65e8be7..95deedb 100644 --- a/src/services/framework/utils/__tests__/python-model-config.test.ts +++ b/src/services/framework/utils/__tests__/python-model-config.test.ts @@ -9,4 +9,13 @@ describe('generatePythonModelConfigPy', () => { expect(py).toContain('PYMODEL_DS'); expect(py).toContain('PYMODEL_DATES'); }); + + test('includes table property helpers for post-run lineage metadata', () => { + const py = generatePythonModelConfigPy(); + expect(py).toContain('def table_properties_from_json'); + expect(py).toContain('def resolve_output_table_fqn'); + expect(py).toContain('def apply_table_properties_from_json'); + expect(py).toContain('python_model_upstream_sources'); + expect(py).toContain('python_model_name'); + }); }); diff --git a/src/services/framework/utils/python-model-utils.ts b/src/services/framework/utils/python-model-utils.ts index 1ce682e..b9ad0cb 100644 --- a/src/services/framework/utils/python-model-utils.ts +++ b/src/services/framework/utils/python-model-utils.ts @@ -80,6 +80,65 @@ class PythonModelConfig: "python_model_table": self.table_name or self.model_name, "python_model_description": self.description, } + + +def normalize_upstream_sources(raw) -> list[dict]: + """Normalize upstream_sources from .python.json (single object or array).""" + if not raw: + return [] + if isinstance(raw, list): + return [e for e in raw if isinstance(e, dict) and e.get("type") and e.get("value")] + if isinstance(raw, dict) and raw.get("type") and raw.get("value"): + return [raw] + return [] + + +def serialize_upstream_sources_for_property(entries: list[dict]) -> str: + """Serialize upstream sources for the Iceberg python_model_upstream_sources property.""" + import json + return json.dumps(entries, separators=(",", ":")) + + +def table_properties_from_json(config: dict) -> dict[str, str]: + """Build Iceberg extra_properties from a .python.json config dict.""" + output = config.get("output") or {} + namespace = config.get("namespace") or output.get("schema") or "" + table_name = config.get("table_name") or output.get("table") or config.get("name", "") + props = { + "python_model_name": config.get("name", ""), + "python_model_type": "python", + "python_model_namespace": namespace, + "python_model_table": table_name, + "python_model_description": config.get("description") or "", + } + upstream = normalize_upstream_sources(config.get("upstream_sources")) + if upstream: + props["python_model_upstream_sources"] = serialize_upstream_sources_for_property(upstream) + return props + + +def resolve_output_table_fqn(config: dict) -> str: + """Resolve catalog.schema.table for the model output from .python.json.""" + output = config.get("output") or {} + catalog = output.get("database") or get_catalog_name() + schema = output.get("schema") or config.get("namespace") or get_namespace_name() + table = output.get("table") or config.get("table_name") or config.get("name", "") + return f'"{catalog}"."{schema}"."{table}"' + + +def apply_table_properties_from_json(config: dict) -> None: + """Write python_model_* Iceberg table properties after a successful model run.""" + import logging + from python_models._trino_io import set_extra_properties # noqa: PLC0415 + + log = logging.getLogger(__name__) + props = table_properties_from_json(config) + if not props.get("python_model_name"): + return + table_fqn = resolve_output_table_fqn(config) + items = list(props.items()) + log.info("apply_table_properties_from_json: %s (%d properties)", table_fqn, len(items)) + set_extra_properties(table_fqn, items) `; } diff --git a/src/services/modelLineage.ts b/src/services/modelLineage.ts index 14d956a..0ec4eba 100644 --- a/src/services/modelLineage.ts +++ b/src/services/modelLineage.ts @@ -22,6 +22,7 @@ import type { ReverseLineageData, } from '@shared/modellineage/types'; import { quoteTrinoIdentifier } from '@shared/sql/identifier'; +import { buildPythonUpstreamGraph } from '@shared/pymodel/python-model-lineage'; import * as fs from 'fs'; import * as path from 'path'; import * as vscode from 'vscode'; @@ -452,115 +453,46 @@ export class ModelLineage { nodes: LineageNode[]; edges: { pythonModelNodeId: string; sourceNodeId: string }[]; }> { - const pythonModelNodes: LineageNode[] = []; - const pythonModelEdges: { - pythonModelNodeId: string; - sourceNodeId: string; - }[] = []; const sourceNodes = upstreamNodes.filter((n) => n.type === 'source'); if (sourceNodes.length === 0) { - return { nodes: pythonModelNodes, edges: pythonModelEdges }; + return { nodes: [], edges: [] }; } - const queries = sourceNodes.map(async (sourceNode) => { - const catalog = sourceNode.database; - const schema = sourceNode.schema; - const tableName = sourceNode.name; + const fetchProperties = async ( + catalog: string, + schema: string, + table: string, + ): Promise> => { + const sql = `SELECT key, value FROM ${quoteTrinoIdentifier(catalog)}.${quoteTrinoIdentifier(schema)}.${quoteTrinoIdentifier(`${table}$properties`)} WHERE key LIKE 'python_model_%'`; + this.coder.log.info(`[Lineage] Querying python model properties: ${sql}`); - if (!catalog || !schema || !tableName) { - return null; - } - - try { - const sql = `SELECT key, value FROM ${quoteTrinoIdentifier(catalog)}.${quoteTrinoIdentifier(schema)}.${quoteTrinoIdentifier(`${tableName}$properties`)} WHERE key LIKE 'python_model_%'`; - this.coder.log.info( - `[Lineage] Querying python model properties: ${sql}`, - ); - - const rows = await this.coder.trino.handleQuery(sql, { - filename: 'data-explorer-query.sql', - }); - - if (!rows || rows.length === 0) { - return null; - } - - const props: Record = {}; - for (const row of rows) { - props[row['key']] = row['value']; - } - - const pythonModelName = props['python_model_name']; - if (!pythonModelName) { - return null; - } - - const upstreamSourcesStr = props['python_model_upstream_sources']; - const upstreamSources = upstreamSourcesStr - ? upstreamSourcesStr - .split(',') - .map((s) => s.trim()) - .filter(Boolean) - : []; - - if (upstreamSources.length === 0) { - return null; - } - - const upstreamSourceNodes: LineageNode[] = []; - for (const sourceId of upstreamSources) { - const parts = sourceId.split('.'); - if (parts.length >= 2) { - const sourceSchema = parts[0]; - const sourceTable = parts[1]; - upstreamSourceNodes.push({ - id: `python.${sourceId}`, - name: sourceTable, - type: 'python', - description: sourceId, - tags: ['python'], - path: '', - schema: sourceSchema, - database: catalog, - hasOwnUpstream: false, - hasOwnDownstream: true, - }); - } - } - - if (upstreamSourceNodes.length === 0) { - return null; - } + const rows = await this.coder.trino.handleQuery(sql, { + filename: 'data-explorer-query.sql', + }); - return { - sourceNodeId: sourceNode.id, - upstreamSourceNodes, - }; - } catch (error) { - this.coder.log.info( - `[Lineage] No python model properties for ${catalog}.${schema}.${tableName}: ${error}`, - ); - return null; + if (!rows || rows.length === 0) { + throw new Error('No python model properties'); } - }); - const results = await Promise.all(queries); - - for (const result of results) { - if (!result) { - continue; + const props: Record = {}; + for (const row of rows) { + props[row['key']] = row['value']; } - for (const upstreamNode of result.upstreamSourceNodes) { - pythonModelNodes.push(upstreamNode); - pythonModelEdges.push({ - pythonModelNodeId: upstreamNode.id, - sourceNodeId: result.sourceNodeId, - }); - } - } + return props; + }; - return { nodes: pythonModelNodes, edges: pythonModelEdges }; + try { + return await buildPythonUpstreamGraph({ + startSourceNodes: sourceNodes, + fetchProperties, + }); + } catch (error) { + this.coder.log.info( + `[Lineage] Python model upstream discovery failed: ${error}`, + ); + return { nodes: [], edges: [] }; + } } /** diff --git a/src/shared/framework/types.ts b/src/shared/framework/types.ts index c57342e..c30ae8e 100644 --- a/src/shared/framework/types.ts +++ b/src/shared/framework/types.ts @@ -67,6 +67,13 @@ export interface KpoSizeSpec { guidance: string; } +export type PythonModelUpstreamSourceType = 'trino' | 'external'; + +export interface PythonModelUpstreamSource { + type: PythonModelUpstreamSourceType; + value: string; +} + export interface PythonModelConfig { name: string; group: PythonModelGroup; @@ -87,6 +94,7 @@ export interface PythonModelConfig { table_name?: string; cells?: PythonModelCell[]; variables?: Record; + upstream_sources?: PythonModelUpstreamSource | PythonModelUpstreamSource[]; } export type ModelCreateResponse = { diff --git a/src/shared/pymodel/__tests__/python-model-lineage.test.ts b/src/shared/pymodel/__tests__/python-model-lineage.test.ts new file mode 100644 index 0000000..8a82c43 --- /dev/null +++ b/src/shared/pymodel/__tests__/python-model-lineage.test.ts @@ -0,0 +1,209 @@ +import { describe, expect, test } from '@jest/globals'; + +import { + buildPythonUpstreamGraph, + parsePythonModelProperties, +} from '@shared/pymodel/python-model-lineage'; +import type { LineageNode } from '@shared/modellineage/types'; + +describe('parsePythonModelProperties', () => { + test('returns null when python_model_name is missing', () => { + expect(parsePythonModelProperties({})).toBeNull(); + }); + + test('parses structured upstream_sources JSON', () => { + expect( + parsePythonModelProperties({ + python_model_name: 'model_d', + python_model_upstream_sources: + '[{"type":"trino","value":"opus_python_source.model_c"}]', + }), + ).toEqual({ + pythonModelName: 'model_d', + upstreamSources: [{ type: 'trino', value: 'opus_python_source.model_c' }], + namespace: undefined, + table: undefined, + description: undefined, + }); + }); + + test('parses legacy comma-separated upstream_sources', () => { + expect( + parsePythonModelProperties({ + python_model_name: 'model_c', + python_model_upstream_sources: 'schema.b, schema.a', + }), + ).toEqual({ + pythonModelName: 'model_c', + upstreamSources: [ + { type: 'trino', value: 'schema.b' }, + { type: 'trino', value: 'schema.a' }, + ], + namespace: undefined, + table: undefined, + description: undefined, + }); + }); +}); + +describe('buildPythonUpstreamGraph', () => { + const sourceD: LineageNode = { + id: 'source.project.schema.model_d', + name: 'model_d', + type: 'source', + path: '', + schema: 'opus_python_source', + database: 'glue', + }; + + test('walks a multi-hop Python model chain D -> C -> B,A', async () => { + const propsByTable: Record> = { + 'glue.opus_python_source.model_d': { + python_model_name: 'model_d', + python_model_upstream_sources: + '[{"type":"trino","value":"opus_python_source.model_c"}]', + }, + 'glue.opus_python_source.model_c': { + python_model_name: 'model_c', + python_model_upstream_sources: + '[{"type":"trino","value":"opus_python_source.model_b"},{"type":"external","value":"backstage_api"}]', + }, + 'glue.opus_python_source.model_b': { + python_model_name: 'model_b', + }, + }; + + const fetchProperties = async ( + catalog: string, + schema: string, + table: string, + ) => { + const key = `${catalog}.${schema}.${table}`; + const props = propsByTable[key]; + if (!props) { + throw new Error(`missing ${key}`); + } + return props; + }; + + const { nodes, edges } = await buildPythonUpstreamGraph({ + startSourceNodes: [sourceD], + fetchProperties, + }); + + const nodeIds = nodes.map((n) => n.id); + expect(nodeIds).toContain('python.model_d'); + expect(nodeIds).toContain('python.model_c'); + expect(nodeIds).toContain('python.model_b'); + expect(nodeIds).toContain('python.external.backstage_api'); + + expect(edges).toEqual( + expect.arrayContaining([ + { pythonModelNodeId: 'python.model_d', sourceNodeId: sourceD.id }, + { pythonModelNodeId: 'python.model_c', sourceNodeId: 'python.model_d' }, + { pythonModelNodeId: 'python.model_b', sourceNodeId: 'python.model_c' }, + { + pythonModelNodeId: 'python.external.backstage_api', + sourceNodeId: 'python.model_c', + }, + ]), + ); + + const modelC = nodes.find((n) => n.id === 'python.model_c'); + expect(modelC?.hasOwnUpstream).toBe(true); + }); + + test('detects cycles without infinite loop', async () => { + const propsByTable: Record> = { + 'glue.s.t1': { + python_model_name: 'm1', + python_model_upstream_sources: + '[{"type":"trino","value":"s.t2"}]', + }, + 'glue.s.t2': { + python_model_name: 'm2', + python_model_upstream_sources: + '[{"type":"trino","value":"s.t1"}]', + }, + }; + + const source: LineageNode = { + id: 'source.s.t1', + name: 't1', + type: 'source', + path: '', + schema: 's', + database: 'glue', + }; + + const { nodes } = await buildPythonUpstreamGraph({ + startSourceNodes: [source], + fetchProperties: async (catalog, schema, table) => + propsByTable[`${catalog}.${schema}.${table}`], + }); + + expect(nodes.map((n) => n.id).sort()).toEqual(['python.m1', 'python.m2']); + }); + + test('skips non-python manifest sources without properties', async () => { + const source: LineageNode = { + id: 'source.s.regular_table', + name: 'regular_table', + type: 'source', + path: '', + schema: 's', + database: 'glue', + }; + + const { nodes, edges } = await buildPythonUpstreamGraph({ + startSourceNodes: [source], + fetchProperties: async () => { + throw new Error('no properties'); + }, + }); + + expect(nodes).toHaveLength(0); + expect(edges).toHaveLength(0); + }); + + test('creates trino ref leaf when upstream table has no python metadata', async () => { + const source: LineageNode = { + id: 'source.s.out', + name: 'out', + type: 'source', + path: '', + schema: 's', + database: 'glue', + }; + + const propsByTable: Record> = { + 'glue.s.out': { + python_model_name: 'out_model', + python_model_upstream_sources: + '[{"type":"trino","value":"s.raw_table"}]', + }, + }; + + const { nodes, edges } = await buildPythonUpstreamGraph({ + startSourceNodes: [source], + fetchProperties: async (catalog, schema, table) => { + const key = `${catalog}.${schema}.${table}`; + const props = propsByTable[key]; + if (!props) { + throw new Error('no properties'); + } + return props; + }, + }); + + expect(nodes.map((n) => n.id)).toContain('python.trino.s.raw_table'); + expect(edges).toEqual( + expect.arrayContaining([ + { + pythonModelNodeId: 'python.trino.s.raw_table', + sourceNodeId: 'python.out_model', + }, + ]), + ); + }); +}); diff --git a/src/shared/pymodel/__tests__/python-upstream-sources.test.ts b/src/shared/pymodel/__tests__/python-upstream-sources.test.ts new file mode 100644 index 0000000..2a322db --- /dev/null +++ b/src/shared/pymodel/__tests__/python-upstream-sources.test.ts @@ -0,0 +1,80 @@ +import { describe, expect, test } from '@jest/globals'; + +import { + normalizeUpstreamSources, + parseUpstreamSourcesProperty, + serializeUpstreamSourcesForProperty, +} from '@shared/pymodel/python-upstream-sources'; + +describe('normalizeUpstreamSources', () => { + test('returns empty array for null/undefined', () => { + expect(normalizeUpstreamSources(null)).toEqual([]); + expect(normalizeUpstreamSources(undefined)).toEqual([]); + }); + + test('wraps a single entry in an array', () => { + expect( + normalizeUpstreamSources({ type: 'trino', value: 'schema.table' }), + ).toEqual([{ type: 'trino', value: 'schema.table' }]); + }); + + test('passes through a valid array', () => { + const entries = [ + { type: 'trino' as const, value: 'a.b' }, + { type: 'external' as const, value: 'api' }, + ]; + expect(normalizeUpstreamSources(entries)).toEqual(entries); + }); + + test('filters invalid array items', () => { + expect( + normalizeUpstreamSources([ + { type: 'trino', value: 'ok' }, + { type: 'bad', value: 'x' }, + { value: 'no-type' }, + ]), + ).toEqual([{ type: 'trino', value: 'ok' }]); + }); +}); + +describe('serializeUpstreamSourcesForProperty', () => { + test('serializes as compact JSON array', () => { + expect( + serializeUpstreamSourcesForProperty([ + { type: 'trino', value: 'opus_python_source.raw_events' }, + { type: 'external', value: 'backstage_api' }, + ]), + ).toBe( + '[{"type":"trino","value":"opus_python_source.raw_events"},{"type":"external","value":"backstage_api"}]', + ); + }); +}); + +describe('parseUpstreamSourcesProperty', () => { + test('parses JSON array format', () => { + expect( + parseUpstreamSourcesProperty( + '[{"type":"trino","value":"schema.table"},{"type":"external","value":"api"}]', + ), + ).toEqual([ + { type: 'trino', value: 'schema.table' }, + { type: 'external', value: 'api' }, + ]); + }); + + test('falls back to legacy comma-separated trino entries', () => { + expect( + parseUpstreamSourcesProperty( + 'opus_python_source.raw_events, analytics.user_sessions', + ), + ).toEqual([ + { type: 'trino', value: 'opus_python_source.raw_events' }, + { type: 'trino', value: 'analytics.user_sessions' }, + ]); + }); + + test('returns empty for blank input', () => { + expect(parseUpstreamSourcesProperty('')).toEqual([]); + expect(parseUpstreamSourcesProperty(undefined)).toEqual([]); + }); +}); diff --git a/src/shared/pymodel/python-model-lineage.ts b/src/shared/pymodel/python-model-lineage.ts new file mode 100644 index 0000000..2e2682d --- /dev/null +++ b/src/shared/pymodel/python-model-lineage.ts @@ -0,0 +1,232 @@ +import type { PythonModelUpstreamSource } from '@shared/framework/types'; +import type { LineageNode } from '@shared/modellineage/types'; + +import { parseUpstreamSourcesProperty } from './python-upstream-sources'; + +export interface PythonModelTableProps { + pythonModelName: string; + upstreamSources: PythonModelUpstreamSource[]; + namespace?: string; + table?: string; + description?: string; +} + +export interface PythonModelLineageEdge { + pythonModelNodeId: string; + sourceNodeId: string; +} + +export interface PythonModelLineageWalkResult { + nodes: LineageNode[]; + edges: PythonModelLineageEdge[]; +} + +/** Parse Iceberg `python_model_*` table properties into structured metadata. */ +export function parsePythonModelProperties( + props: Record, +): PythonModelTableProps | null { + const pythonModelName = props['python_model_name']; + if (!pythonModelName) { + return null; + } + + return { + pythonModelName, + upstreamSources: parseUpstreamSourcesProperty( + props['python_model_upstream_sources'], + ), + namespace: props['python_model_namespace'], + table: props['python_model_table'], + description: props['python_model_description'], + }; +} + +function sanitizeNodeIdSegment(value: string): string { + return value.replace(/[^a-zA-Z0-9._-]/g, '_'); +} + +function hasSourceTableLocation( + node: LineageNode, +): node is LineageNode & { database: string; schema: string } { + return Boolean(node.database && node.schema && node.name); +} + +function makePythonLineageNode( + parsed: PythonModelTableProps, + catalog: string, + schema: string, + table: string, + nodeId: string, +): LineageNode { + return { + id: nodeId, + name: parsed.pythonModelName, + type: 'python', + description: parsed.description || `${schema}.${table}`, + tags: ['python'], + path: '', + schema, + database: catalog, + hasOwnUpstream: false, + hasOwnDownstream: true, + }; +} + +function makeTrinoRefNode( + catalog: string, + schema: string, + table: string, + nodeId: string, +): LineageNode { + return { + id: nodeId, + name: table, + type: 'python', + description: `${schema}.${table}`, + tags: ['python', 'trino-ref'], + path: '', + schema, + database: catalog, + hasOwnUpstream: false, + hasOwnDownstream: true, + }; +} + +function makeExternalLeafNode(value: string, nodeId: string): LineageNode { + return { + id: nodeId, + name: value, + type: 'python', + description: value, + tags: ['python', 'external'], + path: '', + hasOwnUpstream: false, + hasOwnDownstream: true, + }; +} + +interface WalkQueueItem { + catalog: string; + schema: string; + table: string; + attachToNodeId: string; +} + +/** + * Recursively walk Python model upstream chains via Iceberg table properties. + * Starts from manifest source nodes and follows trino-type upstream entries to root. + */ +export async function buildPythonUpstreamGraph(params: { + startSourceNodes: LineageNode[]; + fetchProperties: ( + catalog: string, + schema: string, + table: string, + ) => Promise>; +}): Promise { + const { startSourceNodes, fetchProperties } = params; + const nodesById = new Map(); + const edges: PythonModelLineageEdge[] = []; + const edgeKeys = new Set(); + const visitedTables = new Set(); + + const queue: WalkQueueItem[] = startSourceNodes + .filter(hasSourceTableLocation) + .map((n) => ({ + catalog: n.database, + schema: n.schema, + table: n.name, + attachToNodeId: n.id, + })); + + const addEdge = (pythonModelNodeId: string, sourceNodeId: string) => { + const key = `${pythonModelNodeId}->${sourceNodeId}`; + if (edgeKeys.has(key)) { + return; + } + edgeKeys.add(key); + edges.push({ pythonModelNodeId, sourceNodeId }); + }; + + while (queue.length > 0) { + const item = queue.shift()!; + const tableKey = `${item.catalog}.${item.schema}.${item.table}`; + if (visitedTables.has(tableKey)) { + continue; + } + visitedTables.add(tableKey); + + let props: Record | null = null; + try { + props = await fetchProperties(item.catalog, item.schema, item.table); + } catch { + props = null; + } + + const parsed = props ? parsePythonModelProperties(props) : null; + if (!parsed) { + if (!item.attachToNodeId.startsWith('python.')) { + continue; + } + const refId = `python.trino.${item.schema}.${item.table}`; + if (!nodesById.has(refId)) { + nodesById.set( + refId, + makeTrinoRefNode(item.catalog, item.schema, item.table, refId), + ); + } + addEdge(refId, item.attachToNodeId); + continue; + } + + const nodeId = `python.${parsed.pythonModelName}`; + let node = nodesById.get(nodeId); + if (!node) { + node = makePythonLineageNode( + parsed, + item.catalog, + item.schema, + item.table, + nodeId, + ); + nodesById.set(nodeId, node); + } + + addEdge(nodeId, item.attachToNodeId); + + let hasTrinoUpstream = false; + for (const entry of parsed.upstreamSources) { + if (entry.type === 'external') { + const extId = `python.external.${sanitizeNodeIdSegment(entry.value)}`; + if (!nodesById.has(extId)) { + nodesById.set(extId, makeExternalLeafNode(entry.value, extId)); + } + addEdge(extId, nodeId); + continue; + } + + if (entry.type === 'trino') { + const dotIdx = entry.value.indexOf('.'); + if (dotIdx === -1) { + continue; + } + hasTrinoUpstream = true; + queue.push({ + catalog: item.catalog, + schema: entry.value.slice(0, dotIdx), + table: entry.value.slice(dotIdx + 1), + attachToNodeId: nodeId, + }); + } + } + + if (hasTrinoUpstream) { + nodesById.set(nodeId, { ...node, hasOwnUpstream: true }); + } + } + + return { + nodes: [...nodesById.values()], + edges, + }; +} diff --git a/src/shared/pymodel/python-upstream-sources.ts b/src/shared/pymodel/python-upstream-sources.ts new file mode 100644 index 0000000..11f6a50 --- /dev/null +++ b/src/shared/pymodel/python-upstream-sources.ts @@ -0,0 +1,66 @@ +import type { PythonModelUpstreamSource } from '@shared/framework/types'; + +/** Normalize a single object or array into a list of upstream source entries. */ +export function normalizeUpstreamSources( + raw: unknown, +): PythonModelUpstreamSource[] { + if (raw == null) { + return []; + } + if (Array.isArray(raw)) { + return raw.filter(isUpstreamSourceEntry); + } + if (isUpstreamSourceEntry(raw)) { + return [raw]; + } + return []; +} + +function isUpstreamSourceEntry( + value: unknown, +): value is PythonModelUpstreamSource { + if (typeof value !== 'object' || value === null) { + return false; + } + const entry = value as Record; + return ( + (entry.type === 'trino' || entry.type === 'external') && + typeof entry.value === 'string' && + entry.value.length > 0 + ); +} + +/** Serialize upstream sources for the Iceberg `python_model_upstream_sources` property. */ +export function serializeUpstreamSourcesForProperty( + entries: PythonModelUpstreamSource[], +): string { + return JSON.stringify(entries); +} + +/** + * Parse `python_model_upstream_sources` from Iceberg table properties. + * Falls back to legacy comma-separated `schema.table` strings when JSON parse fails. + */ +export function parseUpstreamSourcesProperty( + raw: string | undefined, +): PythonModelUpstreamSource[] { + if (!raw?.trim()) { + return []; + } + + const trimmed = raw.trim(); + if (trimmed.startsWith('[')) { + try { + const parsed: unknown = JSON.parse(trimmed); + return normalizeUpstreamSources(parsed); + } catch { + // fall through to legacy format + } + } + + return trimmed + .split(',') + .map((s) => s.trim()) + .filter(Boolean) + .map((value) => ({ type: 'trino' as const, value })); +} diff --git a/src/shared/schema/types/python-model.schema.d.ts b/src/shared/schema/types/python-model.schema.d.ts index b44fedd..32e984d 100644 --- a/src/shared/schema/types/python-model.schema.d.ts +++ b/src/shared/schema/types/python-model.schema.d.ts @@ -120,4 +120,23 @@ export interface SchemaPythonModel { variables?: { [k: string]: string | undefined; }; + /** + * Upstream lineage sources written to Iceberg table properties after each run + */ + upstream_sources?: + | SchemaPythonModelUpstreamSource + | [SchemaPythonModelUpstreamSource, ...SchemaPythonModelUpstreamSource[]]; +} +/** + * A single upstream lineage source for a Python model output table + */ +export interface SchemaPythonModelUpstreamSource { + /** + * trino = Iceberg/Hive table (schema.table); external = non-Trino source + */ + type: 'trino' | 'external'; + /** + * trino: schema.table; external: free-form identifier (API name, S3 path, etc.) + */ + value: string; } diff --git a/src/shared/schema/types/python-model.upstream_source.schema.d.ts b/src/shared/schema/types/python-model.upstream_source.schema.d.ts new file mode 100644 index 0000000..566092e --- /dev/null +++ b/src/shared/schema/types/python-model.upstream_source.schema.d.ts @@ -0,0 +1,20 @@ +/* eslint-disable */ +/** + * This file was automatically generated by json-schema-to-typescript. + * DO NOT MODIFY IT BY HAND. Instead, modify the source JSONSchema file, + * and run json-schema-to-typescript to regenerate this file. + */ + +/** + * A single upstream lineage source for a Python model output table + */ +export interface SchemaPythonModelUpstreamSource { + /** + * trino = Iceberg/Hive table (schema.table); external = non-Trino source + */ + type: 'trino' | 'external'; + /** + * trino: schema.table; external: free-form identifier (API name, S3 path, etc.) + */ + value: string; +} diff --git a/templates/skills/dj-create-python-model/_SKILL.md b/templates/skills/dj-create-python-model/_SKILL.md index a32aed4..5583331 100644 --- a/templates/skills/dj-create-python-model/_SKILL.md +++ b/templates/skills/dj-create-python-model/_SKILL.md @@ -121,6 +121,7 @@ Ask about: - **Python packages** needed beyond the defaults (e.g., `requests>=2.28.0`, `boto3`) - **depends_on** — other Python models that must complete first +- **upstream_sources** — upstream lineage for the output table (written to Iceberg properties after each run). Each entry is `{ "type": "trino"|"external", "value": "..." }` where `trino` values are `schema.table` and `external` values are free-form identifiers (API names, S3 paths, etc.) - **tags** — organizational tags (predefined: `python-model`, `api`, `csv`, `s3`, `trino`, `snapshot`, `iceberg`) - **owner** — team or individual - **enable_notebook** — generate companion `.python.ipynb` (default: `true`) diff --git a/templates/skills/dj-review-python-model/_SKILL.md b/templates/skills/dj-review-python-model/_SKILL.md index 4a0ab0a..e83f13d 100644 --- a/templates/skills/dj-review-python-model/_SKILL.md +++ b/templates/skills/dj-review-python-model/_SKILL.md @@ -70,17 +70,17 @@ The DJ lineage engine discovers Python models by querying Iceberg `$properties` - `python_model_name` — model identity - `python_model_table` — output table name -- `python_model_upstream_sources` — comma-separated `schema.table` pairs for upstream lineage edges +- `python_model_upstream_sources` — JSON array of `{ "type": "trino"|"external", "value": "..." }` entries for upstream lineage edges -The review validates the **full chain**: JSON metadata → `PythonModelConfig` → Iceberg table properties → lineage discoverability. +The review validates the **full chain**: JSON metadata → post-run Iceberg table properties → lineage discoverability (including recursive walks to root). | Check | What to validate | | ----- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| L1 | `PythonModelConfig` instantiation emits all required properties: `python_model_name`, `python_model_type`, `python_model_namespace`, `python_model_table`, `python_model_description` | -| L2 | `python_model_upstream_sources` is set on the output table — code must call `ALTER TABLE ... SET PROPERTIES` or use PyIceberg catalog API to write this key listing all source tables the ETL reads from | -| L3 | Each entry in `python_model_upstream_sources` follows `schema.table` format (dot-separated, matching actual source tables) | -| L4 | **Completeness** — every table referenced in `extract()` / `transform_and_load()` SQL (FROM/JOIN clauses) appears in `python_model_upstream_sources` (no missing upstream edges) | -| L5 | **Accuracy** — every entry in `python_model_upstream_sources` corresponds to a table actually queried in the model (no stale/phantom entries) | +| L1 | `.python.json` includes `name` (written as `python_model_name` on the output table after each run) plus `description`, `namespace` / `output.schema`, and `table_name` / `output.table` as applicable | +| L2 | `upstream_sources` is set in `.python.json` (DJ writes `python_model_upstream_sources` to the output table after each run — no manual `ALTER TABLE` in `.python.py`) | +| L3 | Each `trino` entry's `value` follows `schema.table` format; `external` entries use the intended free-form identifier | +| L4 | **Completeness** — every table referenced in `extract()` / `transform_and_load()` SQL (FROM/JOIN clauses) appears in `upstream_sources` (no missing upstream edges) | +| L5 | **Accuracy** — every entry in `upstream_sources` corresponds to a source actually read by the model (no stale/phantom entries) | | L6 | `dags` field is populated in JSON (model is scheduled, discoverable by Airflow DAG lineage). Empty `dags` = utility module, not a lineage participant | | L7 | `depends_on` correctly lists upstream Python models whose output tables this model reads (task dependency mirrors data dependency) | | L8 | Model ID derivable from file path matches `python_model_name` in properties: `python______` | diff --git a/templates/skills/dj-review-python-model/references/review-checklist.md b/templates/skills/dj-review-python-model/references/review-checklist.md index 8a9fe76..f88e626 100644 --- a/templates/skills/dj-review-python-model/references/review-checklist.md +++ b/templates/skills/dj-review-python-model/references/review-checklist.md @@ -144,45 +144,26 @@ def table_properties(self) -> dict[str, str]: **How to check:** Verify `OUTPUT_CONFIG` fields populate all five keys. Missing `description` is a warning (lineage node renders without description). -### L2: `python_model_upstream_sources` is set +### L2: `upstream_sources` is set in `.python.json` -**Pass condition:** The code explicitly sets `python_model_upstream_sources` on the output table via one of: +**Pass condition:** The `.python.json` file includes an `upstream_sources` array (or single object) with `{ "type": "trino"|"external", "value": "..." }` entries. DJ writes these to the output Iceberg table as `python_model_upstream_sources` after each successful run — authors do not need manual `ALTER TABLE` code in `.python.py`. -**Option A — ALTER TABLE SET PROPERTIES (Trino SQL):** -```python -execute_trino(f""" - ALTER TABLE {OUTPUT_CONFIG.full_table_id} - SET PROPERTIES extra_properties = MAP( - ARRAY['python_model_upstream_sources'], - ARRAY['.raw_events,analytics.user_sessions'] - ) -""") -``` - -**Option B — PyIceberg catalog API:** -```python -table = catalog.load_table(OUTPUT_CONFIG.table_id) -with table.update_properties() as update: - update.set("python_model_upstream_sources", "schema.table1,schema.table2") -``` - -**Option C — Set via `PythonModelConfig.table_properties` extension:** -```python -OUTPUT_CONFIG = PythonModelConfig( - ..., - upstream_sources=[".raw_events", "analytics.user_sessions"], -) -# table_properties includes python_model_upstream_sources automatically +**Example (preferred):** +```json +"upstream_sources": [ + { "type": "trino", "value": "opus_python_source.raw_events" }, + { "type": "external", "value": "backstage_api" } +] ``` -**Fail:** No code path sets `python_model_upstream_sources` — lineage graph will show the Python model node but no upstream edges. +**Fail:** `upstream_sources` is missing or empty — lineage will show the Python model node but no upstream edges beyond the immediate source. ### L3: Format validation -**Pass condition:** Each entry in `python_model_upstream_sources` is `schema.table` format. +**Pass condition:** Each `trino` entry's `value` is `schema.table` format. Each `external` entry's `value` is the intended identifier (API name, S3 path, etc.). -**Pass:** `.raw_events`, `analytics.user_sessions` -**Fail:** `iceberg_catalog..raw_events` (includes catalog — lineage parser splits on first dot only), `raw_events` (missing schema) +**Pass (trino):** `{ "type": "trino", "value": "opus_python_source.raw_events" }` +**Fail (trino):** `{ "type": "trino", "value": "iceberg_catalog.opus_python_source.raw_events" }` (includes catalog — lineage uses schema.table only), `{ "type": "trino", "value": "raw_events" }` (missing schema) ### L4: Completeness — no missing upstream edges @@ -204,7 +185,7 @@ OUTPUT_CONFIG = PythonModelConfig( - Staging tables (`stg_tmp_*`, `tmp_*`) - Inline CTEs (WITH clause names) -4. Every remaining table must appear in `python_model_upstream_sources`. +4. Every remaining table must appear in `upstream_sources` (as a `trino` entry). **Example finding:** ``` @@ -215,7 +196,7 @@ Missing from upstream_sources: reference.countries ### L5: Accuracy — no stale entries -**Procedure:** Every entry in `python_model_upstream_sources` must correspond to a table actually referenced in the model's SQL. +**Procedure:** Every entry in `upstream_sources` must correspond to a source actually read by the model. **Example finding:** ``` diff --git a/web/src/features/PythonModel/formatKpoSizeInfo.ts b/web/src/features/PythonModel/formatKpoSizeInfo.ts index ffc139c..220567a 100644 --- a/web/src/features/PythonModel/formatKpoSizeInfo.ts +++ b/web/src/features/PythonModel/formatKpoSizeInfo.ts @@ -1,8 +1,7 @@ import type { KpoSizeSpec } from '@shared/framework/types'; export function formatKpoSizeInfo(spec: KpoSizeSpec): string { - const cpuLabel = - spec.cpu === '1' ? '1 vCPU' : `${spec.cpu} vCPUs`; + const cpuLabel = spec.cpu === '1' ? '1 vCPU' : `${spec.cpu} vCPUs`; const guidance = spec.guidance ? ` — ${spec.guidance}` : ''; return `${cpuLabel}, ${spec.memory} memory${guidance}`; } From c7dc4414079e86694920e76e4771606464622c5a Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Tue, 1 Sep 2026 10:21:12 +0530 Subject: [PATCH 06/11] Update skills --- CHANGELOG.md | 6 +- airflow/v2_10/etl_helper.py | 25 + airflow/v2_7/etl_helper.py | 25 + docs/AGENT_SKILLS.md | 6 +- src/services/framework/index.ts | 92 ++- .../__tests__/python-model-config.test.ts | 13 +- src/services/framework/utils/index.ts | 1 + .../framework/utils/python-dag-inject.ts | 18 - .../framework/utils/python-model-upstream.ts | 49 ++ .../framework/utils/python-model-utils.ts | 83 ++- .../framework/utils/source-etl-python.ts | 1 + src/services/modelLineage.ts | 53 +- src/services/sync/PythonModelSyncService.ts | 33 +- src/shared/framework/types.ts | 1 - .../__tests__/python-model-lineage.test.ts | 44 ++ .../python-upstream-inference.test.ts | 128 ++++ .../__tests__/python-upstream-sources.test.ts | 56 ++ src/shared/pymodel/python-model-lineage.ts | 17 +- .../pymodel/python-upstream-inference.ts | 117 ++++ src/shared/pymodel/python-upstream-sources.ts | 57 +- .../skills/dj-create-python-model/_SKILL.md | 468 ++++++++------ .../dj-create-python-model/evals/evals.md | 67 ++ .../references/etl-patterns.md | 603 +++++++----------- .../references/worked-example.md | 508 +++++++-------- .../skills/dj-review-python-model/_SKILL.md | 209 +++--- .../dj-review-python-model/evals/evals.md | 55 ++ .../references/review-checklist.md | 470 +++++++------- web/src/context/app.tsx | 2 +- web/src/forms/FieldSelectSingle.tsx | 2 + web/src/pages/ModelLineage/LineageGraph.tsx | 4 + web/src/pages/ModelLineage/ModelNode.tsx | 56 +- web/src/pages/PythonModelCreate.tsx | 19 +- 32 files changed, 1995 insertions(+), 1293 deletions(-) create mode 100644 src/services/framework/utils/python-model-upstream.ts create mode 100644 src/shared/pymodel/__tests__/python-upstream-inference.test.ts create mode 100644 src/shared/pymodel/python-upstream-inference.ts create mode 100644 templates/skills/dj-create-python-model/evals/evals.md create mode 100644 templates/skills/dj-review-python-model/evals/evals.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 4480443..9e00176 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,15 +21,19 @@ ### Python models -- **Declare upstream lineage in `.python.json`.** Add `"upstream_sources": [{ "type": "trino"|"external", "value": "..." }]` on any Python model; after each successful Airflow or KPO run, DJ writes `python_model_name` and the structured `python_model_upstream_sources` property to the output Iceberg table — no manual `ALTER TABLE` in your `.python.py`. +- **Declare upstream lineage in `.python.json`.** Add `"upstream_sources": [{ "type": "trino"|"external", "value": "..." }]` on any Python model; DJ auto-suggests trino entries from `depends_on` on JSON sync when omitted. After each successful Airflow or KPO run, DJ writes `python_model_name` and the structured `python_model_upstream_sources` property to the output Iceberg table — no manual `ALTER TABLE` in your `.python.py`. - **Model Lineage walks the full Python chain.** When a downstream dbt model sources a Python output table, Data Explorer recursively reads each upstream table's Iceberg properties until it reaches roots (external sources or tables with no further upstream), so a chain like D → C → B, A renders in one view. - **Cleaner lineage for Python models.** Python models now show their real upstream sources in the lineage graph, without the extra placeholder node that used to appear. +- **Model Lineage distinguishes external vs Trino upstream nodes.** External sources (`s3::`, API, Jenkins) and Trino table refs render with distinct badges in the lineage graph. - **One Airflow task per KPO Python model.** Models with `compute: kpo` appear as a single task named after the model (e.g. `kpo_infra`), not separate gate/KPO steps — DJ runs gating and the pod inside that task. +- **KPO runs write Iceberg lineage properties.** Post-run `python_model_*` table properties (including `python_model_upstream_sources`) are applied after successful KPO tasks, not only MWAA worker runs. +- **`_trino_io` pandas helpers for python-path models.** Use `read_partition_to_dataframe` and `overwrite_partition_from_dataframe` when you choose a pandas orchestrator; existing SQL helpers are unchanged. ### Create Python Model - **Group options match your project.** The Create Python Model form now lists the groups from the project you've selected, and refreshes when you switch projects. - **Fix: creating a Python model works reliably.** Creating a Python model and loading its DAG and group options now work as expected. +- **Faster, clearer DAG selection.** Select DAG shows a loading state while options load; **`source_etl`** is omitted from the list because DJ manages that DAG via its generated template; creating a model no longer prompts to modify the selected DAG — task wiring is applied automatically. ## 2.3.0 diff --git a/airflow/v2_10/etl_helper.py b/airflow/v2_10/etl_helper.py index 76c18a6..bd00e94 100644 --- a/airflow/v2_10/etl_helper.py +++ b/airflow/v2_10/etl_helper.py @@ -959,6 +959,31 @@ def _record( def _record_kpo_success(context: dict) -> None: _record_kpo_outcome(context, run_status="success") + _apply_kpo_table_properties(context) + + +def _apply_kpo_table_properties(context: dict) -> None: + """Write python_model_* Iceberg properties after a successful KPO run.""" + try: + dag = context.get("dag") + ti = context.get("ti") + if dag is None or ti is None: + return + model_id = ti.task_id + models = discover_models(dag.dag_id) + model = next((m for m in models if m.get("model_id") == model_id), None) + config = model.get("config") if model else None + if not config: + return + from python_models._config import apply_table_properties_from_json # noqa: PLC0415 + + apply_table_properties_from_json(config) + except Exception as exc: + log.warning( + "Could not apply table properties after KPO success for %s: %s", + context.get("ti", {}).task_id if context.get("ti") else "?", + exc, + ) def _record_kpo_failure(context: dict) -> None: diff --git a/airflow/v2_7/etl_helper.py b/airflow/v2_7/etl_helper.py index 76c18a6..bd00e94 100644 --- a/airflow/v2_7/etl_helper.py +++ b/airflow/v2_7/etl_helper.py @@ -959,6 +959,31 @@ def _record( def _record_kpo_success(context: dict) -> None: _record_kpo_outcome(context, run_status="success") + _apply_kpo_table_properties(context) + + +def _apply_kpo_table_properties(context: dict) -> None: + """Write python_model_* Iceberg properties after a successful KPO run.""" + try: + dag = context.get("dag") + ti = context.get("ti") + if dag is None or ti is None: + return + model_id = ti.task_id + models = discover_models(dag.dag_id) + model = next((m for m in models if m.get("model_id") == model_id), None) + config = model.get("config") if model else None + if not config: + return + from python_models._config import apply_table_properties_from_json # noqa: PLC0415 + + apply_table_properties_from_json(config) + except Exception as exc: + log.warning( + "Could not apply table properties after KPO success for %s: %s", + context.get("ti", {}).task_id if context.get("ti") else "?", + exc, + ) def _record_kpo_failure(context: dict) -> None: diff --git a/docs/AGENT_SKILLS.md b/docs/AGENT_SKILLS.md index c300e11..a7effce 100644 --- a/docs/AGENT_SKILLS.md +++ b/docs/AGENT_SKILLS.md @@ -44,7 +44,7 @@ DJ never overwrites `.agents/project/`. Starters live in the DJ repo at `templat - **Just ask.** Describe your task in natural language ("create a mart for daily orders", "why is this query slow?") and the assistant matches it to a skill via that skill's _Use when…_ description. You can also name a skill directly. - **Progressive disclosure.** A skill loads its `SKILL.md` first and pulls in `references/` or runs `scripts/` only when needed, keeping the assistant focused. -- **Single source of truth.** Skills edit only the JSON sources of truth — `.model.json`, `.source.json`, `.python.json` — and never hand-edit the generated `.sql` / `.yml` / `.python.py`, which DJ regenerates via JSON Sync. +- **Single source of truth.** For SQL models, skills edit the JSON sources of truth — `.model.json`, `.source.json` — and never hand-edit the generated `.sql` / `.yml`, which DJ regenerates via JSON Sync. For Python models, **`.python.py` is the code source of truth** (edit it directly); `.python.json` holds metadata plus a non-authoritative `cells` snapshot, and only `_trino_io.py` / `_config.py` / `_model_tasks.py` are generated. - **You stay in control of DJ commands.** Skills can't run VS Code commands themselves; they'll ask you to run things like **`DJ: Sync to SQL and YML`** or **`DJ: Refresh Projects`** at the right moment. - **Some skills are read-only.** `dj-review-python-model`, `dj-govern-model`, `dj-trino-analyzer`, and `dj-verify-pymodel-parity` produce reports and change nothing. @@ -89,9 +89,9 @@ Registers a raw Trino table as a DJ `.source.json` by introspecting its exact co #### `dj-create-python-model` -Scaffolds a `.python.json` for a pre-dbt Python ETL pipeline that extracts data from external sources (APIs, databases, files) and loads it into Iceberg tables for downstream dbt models to consume. It favors Trino SQL for transforms, using pandas only for ingestion and Python-only logic. +Scaffolds a Python Models data pipeline (`.python.json` metadata + a scaffolded `.python.py` you author) that extracts data from external sources (APIs, databases, files) and loads it into Iceberg tables for downstream dbt models to consume. It favors Trino SQL for transforms, using pandas only for ingestion and Python-only logic. -- **Use when:** you want to create a Python model, ETL pipeline, data ingestion, API fetch, CSV import, or any pre-dbt Python data processing task. +- **Use when:** you want to create a Python model, ETL pipeline, data ingestion, API fetch, CSV import, or any Python Models data pipeline task. - **Example prompt:** _"Create a Python model that fetches the Backstage API into an Iceberg table."_ - **Bundled references:** `etl-patterns.md` (per-stage code templates), `worked-example.md` (a complete end-to-end pipeline). diff --git a/src/services/framework/index.ts b/src/services/framework/index.ts index 566727a..e20a4f3 100644 --- a/src/services/framework/index.ts +++ b/src/services/framework/index.ts @@ -62,7 +62,6 @@ import { PreferencesHandler } from './handlers/preferences-handler'; import { SourceHandler } from './handlers/source-handler'; import { UIHandlers } from './handlers/ui-handlers'; import { - buildDagInjectionPreview, buildInjectedDagSource, findDangerousSqlStatements, frameworkGetModelId, @@ -73,6 +72,7 @@ import { generatePythonModelScaffoldPy, generatePythonModelTasksPy, generateTrinoIoPy, + SOURCE_ETL_DAG_ID, workspaceHasPythonModels, } from './utils'; import { @@ -426,7 +426,6 @@ export class Framework implements ApiEnabledService<'framework'> { throw new Error('Project not found'); } - // Create directory structure: dags/python_models/// const pythonModelDir = path.join( WORKSPACE_ROOT, 'dags', @@ -437,78 +436,77 @@ export class Framework implements ApiEnabledService<'framework'> { await vscode.workspace.fs.createDirectory( vscode.Uri.file(pythonModelDir), ); - - // Ensure __init__.py files exist for Python package structure await this.ensurePythonModelsInitFiles(group, topic); - // Check if python model file already exists - const pythonJsonFileName = `${name}.python.json`; - const pythonJsonPath = path.join(pythonModelDir, pythonJsonFileName); - const pythonJsonUri = vscode.Uri.file(pythonJsonPath); - try { - await vscode.workspace.fs.stat(pythonJsonUri); - throw new Error( - `Python model '${name}' already exists in ${group}/${topic}`, - ); - } catch (err) { - // File doesn't exist, which is expected - continue with creation - if (err instanceof Error && err.message.includes('already exists')) { - throw err; - } - } + const baseTags = + tags && tags.length > 0 ? tags : ['python-model', group]; + const compute = reqCompute ?? 'kpo'; + const kpoFields = + compute === 'kpo' || reqCompute === undefined + ? { kpo_size: reqKpoSize ?? 'small' } + : {}; const pythonModelConfig: PythonModelConfig = { name, group, topic, - ...(description && { description }), model_type, dags, - compute: reqCompute ?? 'kpo', - ...(reqCompute === 'kpo' || reqCompute === undefined - ? { kpo_size: reqKpoSize ?? 'small' } - : {}), + compute, + ...kpoFields, + ...(description && { description }), ...(enable_notebook && { enable_notebook }), - tags: tags && tags.length > 0 ? tags : ['python-model', group], + tags: baseTags, ...(reqNamespace && { namespace: reqNamespace }), ...(reqTableName && { table_name: reqTableName }), variables: {}, }; - // Write metadata-only JSON (the .python.py is the source of truth). + const pythonJsonPath = path.join( + pythonModelDir, + `${pythonModelConfig.name}.python.json`, + ); + const pythonJsonUri = vscode.Uri.file(pythonJsonPath); + try { + await vscode.workspace.fs.stat(pythonJsonUri); + throw new Error( + `Python model '${pythonModelConfig.name}' already exists in ${group}/${topic}`, + ); + } catch (err) { + if ( + err instanceof Error && + err.message.includes('already exists') + ) { + throw err; + } + } + await vscode.workspace.fs.writeFile( pythonJsonUri, Buffer.from(JSON.stringify(pythonModelConfig, null, ' ')), ); - // Scaffold a clean, marker-delimited .python.py source of truth - const modelContent = generatePythonModelScaffoldPy(pythonModelConfig); - const modelFileName = `${name}.python.py`; const modelUri = vscode.Uri.file( - path.join(pythonModelDir, modelFileName), + path.join(pythonModelDir, `${pythonModelConfig.name}.python.py`), ); await vscode.workspace.fs.writeFile( modelUri, - Buffer.from(modelContent), + Buffer.from(generatePythonModelScaffoldPy(pythonModelConfig)), ); - // Ensure shared scaffolding and refresh source_etl when Python models exist await this.ensurePythonModelInfrastructure(); await this.dbt.writeAirflowDags(); - // Generate new DAG file if requested if (createDag && dagConfig?.name) { await this.generateDagFile(project, dagConfig); } - // Modify selected DAGs to add register_python_model_tasks (with confirm) if (dags && dags.length > 0) { for (const dagName of dags) { await this.injectPythonModelTasksIntoDag(project, dagName); } } - // Clear form state try { await this.getApi().handleApi({ type: 'state-clear', @@ -521,8 +519,7 @@ export class Framework implements ApiEnabledService<'framework'> { ); } - // Open the created file - vscode.window.showTextDocument(pythonJsonUri); + void vscode.window.showTextDocument(pythonJsonUri); return 'Python model created'; } @@ -2110,6 +2107,12 @@ export class Framework implements ApiEnabledService<'framework'> { for (const [fileName] of pyFiles) { const dagName = fileName.replace('.py', ''); + if (dagName === SOURCE_ETL_DAG_ID) { + this.log.info( + `[DAG Discovery] Skip ${dagName} (DJ-managed template)`, + ); + continue; + } if (seenDags.has(dagName)) { this.log.info(`[DAG Discovery] Skip duplicate: ${dagName}`); continue; @@ -2841,8 +2844,7 @@ with DAG( /** * Inject register_python_model_tasks wiring into an existing DAG file. * Called when a Python model is created and mapped to a DAG. - * Shows a confirm modal with a short preview before writing. - * @param dagName - Can be a simple name (e.g., "source_etl") or a relative path (e.g., "_ext_/source_etl") + * @param dagName - Can be a simple name (e.g., "my_dag") or a relative path (e.g., "_ext_/my_dag") */ private async injectPythonModelTasksIntoDag( project: DbtProject, @@ -2916,18 +2918,6 @@ with DAG( ); } - const preview = buildDagInjectionPreview(dagFilePath, dagName); - const choice = await vscode.window.showWarningMessage( - `Modify DAG "${dagName}" to register Python model tasks?`, - { modal: true, detail: preview }, - 'Apply', - 'Skip', - ); - if (choice !== 'Apply') { - this.log.info(`DAG injection skipped by user for: ${dagName}`); - return; - } - await vscode.workspace.fs.writeFile( vscode.Uri.file(dagFilePath), Buffer.from(built.next), diff --git a/src/services/framework/utils/__tests__/python-model-config.test.ts b/src/services/framework/utils/__tests__/python-model-config.test.ts index 95deedb..9fc12a0 100644 --- a/src/services/framework/utils/__tests__/python-model-config.test.ts +++ b/src/services/framework/utils/__tests__/python-model-config.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from '@jest/globals'; -import { generatePythonModelConfigPy } from '@services/framework/utils/python-model-utils'; +import { generatePythonModelConfigPy, generateTrinoIoPy } from '@services/framework/utils/python-model-utils'; describe('generatePythonModelConfigPy', () => { test('includes build_context_from_env for __main__ script execution', () => { @@ -8,6 +8,8 @@ describe('generatePythonModelConfigPy', () => { expect(py).toContain('def build_context_from_env()'); expect(py).toContain('PYMODEL_DS'); expect(py).toContain('PYMODEL_DATES'); + expect(py).toContain('logging.basicConfig'); + expect(py).toContain('stream=sys.stdout'); }); test('includes table property helpers for post-run lineage metadata', () => { @@ -19,3 +21,12 @@ describe('generatePythonModelConfigPy', () => { expect(py).toContain('python_model_name'); }); }); + +describe('generateTrinoIoPy', () => { + test('includes pandas helpers for python-path models', () => { + const py = generateTrinoIoPy(); + expect(py).toContain('def read_partition_to_dataframe'); + expect(py).toContain('def overwrite_partition_from_dataframe'); + expect(py).toContain('Pandas I/O (python-path models only)'); + }); +}); diff --git a/src/services/framework/utils/index.ts b/src/services/framework/utils/index.ts index 476fbd6..6006665 100644 --- a/src/services/framework/utils/index.ts +++ b/src/services/framework/utils/index.ts @@ -32,6 +32,7 @@ export * from './meta-lint'; // Python model utilities export * from './python-model-utils'; +export * from './python-model-upstream'; // Destructive-SQL guard for python model files export * from './python-sql-guard'; diff --git a/src/services/framework/utils/python-dag-inject.ts b/src/services/framework/utils/python-dag-inject.ts index c84a1ba..67677a4 100644 --- a/src/services/framework/utils/python-dag-inject.ts +++ b/src/services/framework/utils/python-dag-inject.ts @@ -212,21 +212,3 @@ export function buildInjectedDagSource(dagCode: string): { return { next, style }; } - -/** - * Short preview snippet shown in the confirm modal before writing. - */ -export function buildDagInjectionPreview( - dagFilePath: string, - dagId: string, -): string { - return [ - `File: ${dagFilePath}`, - '', - 'Will add:', - ' from _ext_.etl_helper import register_python_model_tasks', - ` entry_tasks, exit_tasks = register_python_model_tasks("${dagId}", ...)`, - '', - 'and wire those tasks into the existing DAG dependency chain.', - ].join('\n'); -} diff --git a/src/services/framework/utils/python-model-upstream.ts b/src/services/framework/utils/python-model-upstream.ts new file mode 100644 index 0000000..594a01b --- /dev/null +++ b/src/services/framework/utils/python-model-upstream.ts @@ -0,0 +1,49 @@ +import type { PythonModelConfig } from '@shared/framework/types'; +import { + DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS, + enrichPythonModelUpstreamSources, + type PythonModelCatalogDefaults, +} from '@shared/pymodel/python-upstream-inference'; +import * as fs from 'fs'; +import * as path from 'path'; + +export { DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS }; + +/** Load all python model configs in a topic directory keyed by model name. */ +export function loadTopicPythonModelConfigs( + topicDir: string, +): Map { + const configs = new Map(); + if (!fs.existsSync(topicDir)) { + return configs; + } + + for (const entry of fs.readdirSync(topicDir)) { + if (!entry.endsWith('.python.json')) { + continue; + } + const jsonPath = path.join(topicDir, entry); + try { + const raw = fs.readFileSync(jsonPath, 'utf8'); + const config = JSON.parse(raw) as PythonModelConfig; + if (config.name) { + configs.set(config.name, config); + } + } catch { + // skip invalid files + } + } + + return configs; +} + +/** Merge auto-inferred trino upstream_sources from depends_on when missing. */ +export function applyUpstreamInference( + config: PythonModelConfig, + topicDir: string, + defaults: PythonModelCatalogDefaults = DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS, +): PythonModelConfig { + const siblings = loadTopicPythonModelConfigs(topicDir); + siblings.set(config.name, config); + return enrichPythonModelUpstreamSources(config, siblings, defaults); +} diff --git a/src/services/framework/utils/python-model-utils.ts b/src/services/framework/utils/python-model-utils.ts index b9ad0cb..9ade40d 100644 --- a/src/services/framework/utils/python-model-utils.ts +++ b/src/services/framework/utils/python-model-utils.ts @@ -50,7 +50,15 @@ def get_namespace_name() -> str: def build_context_from_env() -> dict: """Build Airflow run context from PYMODEL_* env vars (KPO pod or MWAA worker).""" import json + import logging import os + import sys + + logging.basicConfig( + level=logging.INFO, + stream=sys.stdout, + format="%(message)s", + ) ds = os.environ.get("PYMODEL_DS") or "" ctx: dict = { @@ -155,9 +163,9 @@ export function generateTrinoIoPy(): string { Shared Trino I/O helpers for python_models. Auto-generated by DJ Framework -- DO NOT EDIT. -Provides connection construction, SQL execution, and DML operations -(append, overwrite_partition, overwrite_table, merge, update) -used by Python model ETL pipelines. +Provides connection construction, SQL execution, DML operations +(append, overwrite_partition, overwrite_table, merge, update), DDL helpers, +and optional pandas read/write helpers for python-path models. """ from __future__ import annotations @@ -439,6 +447,75 @@ def update(table_fqn: str, set_clause: str, where_clause: str) -> None: update_sql = f"UPDATE {table_fqn} SET {set_clause} WHERE {where_clause}" log.info(f"update: {update_sql}") execute_sql(update_sql) + + +# ── Pandas I/O (python-path models only) ─────────────────────────── + + +def read_partition_to_dataframe( + table_fqn: str, + partition_col: str, + partition_value: str, + columns: Sequence[str] | None = None, +): + """Read a single partition slice from Trino into a pandas DataFrame.""" + import pandas as pd + + safe_value = str(partition_value).replace("'", "''") + col_list = ", ".join(f'"{c}"' for c in columns) if columns else "*" + sql = ( + f"SELECT {col_list} FROM {table_fqn} " + f'WHERE "{partition_col}" = \\'{safe_value}\\'' + ) + log.info("read_partition_to_dataframe: %s", sql) + with get_trino_connection() as conn: + cursor = conn.cursor() + cursor.execute(sql) + rows = cursor.fetchall() + col_names = [desc[0] for desc in cursor.description] if cursor.description else [] + if not rows: + return pd.DataFrame(columns=col_names or list(columns or [])) + return pd.DataFrame(rows, columns=col_names) + + +def overwrite_partition_from_dataframe( + table_fqn: str, + partition_col: str, + partition_value: str, + df, + columns: Sequence[str], + batch_size: int = 1000, +) -> None: + """Overwrite one partition from a pandas DataFrame (DELETE + chunked INSERT).""" + safe_value = str(partition_value).replace("'", "''") + overwrite_filter = f'"{partition_col}" = \\'{safe_value}\\'' + delete_sql = f"DELETE FROM {table_fqn} WHERE {overwrite_filter}" + log.info("overwrite_partition_from_dataframe: delete %s", delete_sql) + execute_sql(delete_sql) + + if df is None or len(df) == 0: + log.info("overwrite_partition_from_dataframe: empty dataframe, partition cleared") + return + + col_list = ", ".join(f'"{c}"' for c in columns) + total = len(df) + for start in range(0, total, batch_size): + batch = df.iloc[start : start + batch_size] + rows = [ + tuple(batch[c].iloc[i] for c in columns) + for i in range(len(batch)) + ] + if not rows: + continue + insert_sql = ( + f"INSERT INTO {table_fqn} ({col_list}) VALUES {_rows_to_values(rows)}" + ) + execute_sql(insert_sql) + log.info( + "overwrite_partition_from_dataframe: inserted %d rows into %s", + total, + table_fqn, + ) `; } diff --git a/src/services/framework/utils/source-etl-python.ts b/src/services/framework/utils/source-etl-python.ts index ea859a3..abf5eca 100644 --- a/src/services/framework/utils/source-etl-python.ts +++ b/src/services/framework/utils/source-etl-python.ts @@ -8,6 +8,7 @@ import * as path from 'path'; +export const SOURCE_ETL_DAG_ID = 'source_etl'; export const SOURCE_ETL_OUTPUT_FILE = 'source_etl.py'; export const SOURCE_ETL_NO_PYTHON_TEMPLATE = 'source_etl.py'; export const SOURCE_ETL_WITH_PYTHON_TEMPLATE = 'source_etl_py_template.py'; diff --git a/src/services/modelLineage.ts b/src/services/modelLineage.ts index 0ec4eba..e355dca 100644 --- a/src/services/modelLineage.ts +++ b/src/services/modelLineage.ts @@ -455,7 +455,21 @@ export class ModelLineage { }> { const sourceNodes = upstreamNodes.filter((n) => n.type === 'source'); + this.coder.log.info( + `[Lineage] discoverPythonModelUpstream: ${upstreamNodes.length} upstream node(s), ${sourceNodes.length} source node(s) to seed the python walk`, + ); + for (const n of sourceNodes) { + this.coder.log.info( + `[Lineage] seed source node: id=${n.id} database=${JSON.stringify( + n.database, + )} schema=${JSON.stringify(n.schema)} name=${JSON.stringify(n.name)}`, + ); + } + if (sourceNodes.length === 0) { + this.coder.log.info( + '[Lineage] No source nodes among upstream -> skipping python model discovery', + ); return { nodes: [], edges: [] }; } @@ -467,9 +481,25 @@ export class ModelLineage { const sql = `SELECT key, value FROM ${quoteTrinoIdentifier(catalog)}.${quoteTrinoIdentifier(schema)}.${quoteTrinoIdentifier(`${table}$properties`)} WHERE key LIKE 'python_model_%'`; this.coder.log.info(`[Lineage] Querying python model properties: ${sql}`); - const rows = await this.coder.trino.handleQuery(sql, { - filename: 'data-explorer-query.sql', - }); + let rows: Record[]; + try { + rows = await this.coder.trino.handleQuery(sql, { + filename: 'data-explorer-query.sql', + }); + } catch (queryError) { + this.coder.log.error( + `[Lineage] $properties query FAILED for ${catalog}.${schema}.${table}: ${ + queryError instanceof Error ? queryError.stack : queryError + }`, + ); + throw queryError; + } + + this.coder.log.info( + `[Lineage] $properties query for ${catalog}.${schema}.${table} returned ${ + rows?.length ?? 0 + } row(s)`, + ); if (!rows || rows.length === 0) { throw new Error('No python model properties'); @@ -479,17 +509,28 @@ export class ModelLineage { for (const row of rows) { props[row['key']] = row['value']; } + this.coder.log.info( + `[Lineage] parsed python_model_* props for ${catalog}.${schema}.${table}: ${JSON.stringify( + props, + )}`, + ); return props; }; try { - return await buildPythonUpstreamGraph({ + const result = await buildPythonUpstreamGraph({ startSourceNodes: sourceNodes, fetchProperties, }); - } catch (error) { this.coder.log.info( - `[Lineage] Python model upstream discovery failed: ${error}`, + `[Lineage] python model discovery produced ${result.nodes.length} node(s) and ${result.edges.length} edge(s)`, + ); + return result; + } catch (error) { + this.coder.log.error( + `[Lineage] Python model upstream discovery failed: ${ + error instanceof Error ? error.stack : error + }`, ); return { nodes: [], edges: [] }; } diff --git a/src/services/sync/PythonModelSyncService.ts b/src/services/sync/PythonModelSyncService.ts index 4112b4d..ed9bdf2 100644 --- a/src/services/sync/PythonModelSyncService.ts +++ b/src/services/sync/PythonModelSyncService.ts @@ -30,11 +30,13 @@ import { generatePythonModelScaffoldPy, stripNotebookOnlyCells, } from '@services/framework/utils'; +import { loadTopicPythonModelConfigs } from '@services/framework/utils/python-model-upstream'; import { workspaceHasPythonModels } from '@services/framework/utils/source-etl-python-workspace'; import type { PythonModelCell, PythonModelConfig, } from '@shared/framework/types'; +import { enrichPythonModelUpstreamSources } from '@shared/pymodel/python-upstream-inference'; import { DJ_PYTHON_TEMP_PATH, WORKSPACE_ROOT } from 'admin'; import * as path from 'path'; import * as vscode from 'vscode'; @@ -479,7 +481,8 @@ export class PythonModelSyncService implements vscode.Disposable { // ── JSON -> PY (scaffold only) ────────────────────────────────────── private async syncFromJson(fileSet: PythonModelFileSet): Promise { - const { config, pyPath, ipynbPath } = fileSet; + const config = await this.maybeMergeInferredUpstreamSources(fileSet); + const { pyPath, ipynbPath } = fileSet; this.logger.info('Syncing from JSON'); // The .py is the source of truth for code. Only scaffold a fresh .py @@ -514,6 +517,34 @@ export class PythonModelSyncService implements vscode.Disposable { } } + private async maybeMergeInferredUpstreamSources( + fileSet: PythonModelFileSet, + ): Promise { + if (!fileSet.config.depends_on?.length) { + return fileSet.config; + } + + const topicDir = path.dirname(fileSet.jsonPath); + const siblings = loadTopicPythonModelConfigs(topicDir); + const enriched = enrichPythonModelUpstreamSources( + fileSet.config, + siblings, + ); + const before = JSON.stringify(fileSet.config.upstream_sources ?? null); + const after = JSON.stringify(enriched.upstream_sources ?? null); + if (before === after) { + return fileSet.config; + } + + this.suppress(fileSet.jsonPath); + await vscode.workspace.fs.writeFile( + vscode.Uri.file(fileSet.jsonPath), + Buffer.from(JSON.stringify(enriched, null, ' ')), + ); + this.logger.info('Merged inferred trino upstream_sources from depends_on'); + return enriched; + } + // ── PY -> IPYNB ───────────────────────────────────────────────────── private async syncFromPython(fileSet: PythonModelFileSet): Promise { diff --git a/src/shared/framework/types.ts b/src/shared/framework/types.ts index c30ae8e..c5294d3 100644 --- a/src/shared/framework/types.ts +++ b/src/shared/framework/types.ts @@ -542,7 +542,6 @@ export type FrameworkApi = description?: string; model_type: 'python'; dags?: string[]; - depends_on?: string[]; compute?: PythonModelCompute; kpo_size?: PythonModelKpoSize; output_type?: PythonModelOutputType; diff --git a/src/shared/pymodel/__tests__/python-model-lineage.test.ts b/src/shared/pymodel/__tests__/python-model-lineage.test.ts index 8a82c43..a80caf9 100644 --- a/src/shared/pymodel/__tests__/python-model-lineage.test.ts +++ b/src/shared/pymodel/__tests__/python-model-lineage.test.ts @@ -206,4 +206,48 @@ describe('buildPythonUpstreamGraph', () => { ]), ); }); + + test('walks 3-part catalog.schema.table trino upstream values', async () => { + const source: LineageNode = { + id: 'source.s.out', + name: 'out', + type: 'source', + path: '', + schema: 'opus_python_source', + database: 'glue_development', + }; + + const propsByTable: Record> = { + 'glue_development.opus_python_source.out': { + python_model_name: 'out_model', + python_model_upstream_sources: + '[{"type":"trino","value":"glue_development.opus_python_source.upstream"}]', + }, + 'glue_development.opus_python_source.upstream': { + python_model_name: 'upstream_model', + }, + }; + + const { nodes, edges } = await buildPythonUpstreamGraph({ + startSourceNodes: [source], + fetchProperties: async (catalog, schema, table) => { + const key = `${catalog}.${schema}.${table}`; + const props = propsByTable[key]; + if (!props) { + throw new Error(`missing ${key}`); + } + return props; + }, + }); + + expect(nodes.map((n) => n.id)).toContain('python.upstream_model'); + expect(edges).toEqual( + expect.arrayContaining([ + { + pythonModelNodeId: 'python.upstream_model', + sourceNodeId: 'python.out_model', + }, + ]), + ); + }); }); diff --git a/src/shared/pymodel/__tests__/python-upstream-inference.test.ts b/src/shared/pymodel/__tests__/python-upstream-inference.test.ts new file mode 100644 index 0000000..eb75c2c --- /dev/null +++ b/src/shared/pymodel/__tests__/python-upstream-inference.test.ts @@ -0,0 +1,128 @@ +import { describe, expect, test } from '@jest/globals'; + +import { + inferTrinoUpstreamFromDependsOn, + mergeInferredUpstreamSources, + resolvePythonModelOutputFqn, +} from '@shared/pymodel/python-upstream-inference'; +import type { PythonModelConfig } from '@shared/framework/types'; + +const defaults = { + catalog: 'glue_development', + namespace: 'opus_python_source', +}; + +describe('resolvePythonModelOutputFqn', () => { + test('uses output block when present', () => { + expect( + resolvePythonModelOutputFqn( + { + name: 'my_model', + output: { + database: 'glue', + schema: 'custom_schema', + table: 'custom_table', + }, + }, + defaults, + ), + ).toBe('glue.custom_schema.custom_table'); + }); + + test('falls back to defaults and model name', () => { + expect( + resolvePythonModelOutputFqn({ name: 'customer_tenant_dc_tracker' }, defaults), + ).toBe('glue_development.opus_python_source.customer_tenant_dc_tracker'); + }); +}); + +describe('inferTrinoUpstreamFromDependsOn', () => { + const ingest: PythonModelConfig = { + name: 'customer_tenant_dc_tracker_ingest', + group: 'etl', + topic: 'customer_tenant_migration', + model_type: 'python', + table_name: 'customer_tenant_dc_tracker', + }; + + const enrichment: PythonModelConfig = { + name: 'customer_tenant_migration_enrichment', + group: 'etl', + topic: 'customer_tenant_migration', + model_type: 'python', + table_name: 'customer_tenant_migration_enrichment', + }; + + const siblings = new Map([ + [ingest.name, ingest], + [enrichment.name, enrichment], + ]); + + test('infers trino FQNs from depends_on', () => { + const status: PythonModelConfig = { + name: 'customer_tenant_migration_status', + group: 'etl', + topic: 'customer_tenant_migration', + model_type: 'python', + depends_on: [ingest.name, enrichment.name], + }; + + expect( + inferTrinoUpstreamFromDependsOn(status, siblings, defaults), + ).toEqual([ + { + type: 'trino', + value: + 'glue_development.opus_python_source.customer_tenant_dc_tracker', + }, + { + type: 'trino', + value: + 'glue_development.opus_python_source.customer_tenant_migration_enrichment', + }, + ]); + }); + + test('does not duplicate author-declared trino entries', () => { + const status: PythonModelConfig = { + name: 'customer_tenant_migration_status', + group: 'etl', + topic: 'customer_tenant_migration', + model_type: 'python', + depends_on: [ingest.name, enrichment.name], + upstream_sources: [ + { + type: 'trino', + value: + 'glue_development.opus_python_source.customer_tenant_dc_tracker', + }, + ], + }; + + expect( + inferTrinoUpstreamFromDependsOn(status, siblings, defaults), + ).toEqual([ + { + type: 'trino', + value: + 'glue_development.opus_python_source.customer_tenant_migration_enrichment', + }, + ]); + }); + + test('mergeInferredUpstreamSources appends inferred entries', () => { + const config: PythonModelConfig = { + name: 'm', + group: 'etl', + topic: 't', + model_type: 'python', + upstream_sources: [{ type: 'external', value: 's3::file.csv' }], + }; + const inferred = [{ type: 'trino' as const, value: 'glue.s.t' }]; + + expect(mergeInferredUpstreamSources(config, inferred)).toEqual([ + { type: 'external', value: 's3::file.csv' }, + { type: 'trino', value: 'glue.s.t' }, + ]); + }); +}); \ No newline at end of file diff --git a/src/shared/pymodel/__tests__/python-upstream-sources.test.ts b/src/shared/pymodel/__tests__/python-upstream-sources.test.ts index 2a322db..3a17d83 100644 --- a/src/shared/pymodel/__tests__/python-upstream-sources.test.ts +++ b/src/shared/pymodel/__tests__/python-upstream-sources.test.ts @@ -2,6 +2,7 @@ import { describe, expect, test } from '@jest/globals'; import { normalizeUpstreamSources, + parseTrinoUpstreamValue, parseUpstreamSourcesProperty, serializeUpstreamSourcesForProperty, } from '@shared/pymodel/python-upstream-sources'; @@ -77,4 +78,59 @@ describe('parseUpstreamSourcesProperty', () => { expect(parseUpstreamSourcesProperty('')).toEqual([]); expect(parseUpstreamSourcesProperty(undefined)).toEqual([]); }); + + test('accepts an already-parsed array (Trino CSV reader pre-parses JSON cells)', () => { + // parseTrinoValue JSON.parses any cell starting with '[' or '{', so the + // property can surface as a real array rather than a string. + expect( + parseUpstreamSourcesProperty([ + { type: 'trino', value: 'glue_development.opus_python_sources.wpc_slavepool_details' }, + { type: 'external', value: 'jenkins::dcim_inventory_csv' }, + ]), + ).toEqual([ + { type: 'trino', value: 'glue_development.opus_python_sources.wpc_slavepool_details' }, + { type: 'external', value: 'jenkins::dcim_inventory_csv' }, + ]); + }); + + test('accepts an already-parsed single object', () => { + expect( + parseUpstreamSourcesProperty({ + type: 'trino', + value: 'glue_development.opus_python_sources.wpc_slavepool_details', + }), + ).toEqual([ + { type: 'trino', value: 'glue_development.opus_python_sources.wpc_slavepool_details' }, + ]); + }); + + test('ignores non-string non-object scalars', () => { + expect(parseUpstreamSourcesProperty(42)).toEqual([]); + expect(parseUpstreamSourcesProperty(null)).toEqual([]); + }); +}); + +describe('parseTrinoUpstreamValue', () => { + test('parses 3-part catalog.schema.table', () => { + expect( + parseTrinoUpstreamValue( + 'glue_development.opus_python_source.model_c', + 'glue', + ), + ).toEqual({ + catalog: 'glue_development', + schema: 'opus_python_source', + table: 'model_c', + }); + }); + + test('parses 2-part schema.table with fallback catalog', () => { + expect(parseTrinoUpstreamValue('opus_python_source.model_c', 'glue')).toEqual( + { + catalog: 'glue', + schema: 'opus_python_source', + table: 'model_c', + }, + ); + }); }); diff --git a/src/shared/pymodel/python-model-lineage.ts b/src/shared/pymodel/python-model-lineage.ts index 2e2682d..3428124 100644 --- a/src/shared/pymodel/python-model-lineage.ts +++ b/src/shared/pymodel/python-model-lineage.ts @@ -1,7 +1,10 @@ import type { PythonModelUpstreamSource } from '@shared/framework/types'; import type { LineageNode } from '@shared/modellineage/types'; -import { parseUpstreamSourcesProperty } from './python-upstream-sources'; +import { + parseTrinoUpstreamValue, + parseUpstreamSourcesProperty, +} from './python-upstream-sources'; export interface PythonModelTableProps { pythonModelName: string; @@ -82,7 +85,7 @@ function makeTrinoRefNode( id: nodeId, name: table, type: 'python', - description: `${schema}.${table}`, + description: `${catalog}.${schema}.${table}`, tags: ['python', 'trino-ref'], path: '', schema, @@ -206,15 +209,15 @@ export async function buildPythonUpstreamGraph(params: { } if (entry.type === 'trino') { - const dotIdx = entry.value.indexOf('.'); - if (dotIdx === -1) { + const parsedRef = parseTrinoUpstreamValue(entry.value, item.catalog); + if (!parsedRef) { continue; } hasTrinoUpstream = true; queue.push({ - catalog: item.catalog, - schema: entry.value.slice(0, dotIdx), - table: entry.value.slice(dotIdx + 1), + catalog: parsedRef.catalog, + schema: parsedRef.schema, + table: parsedRef.table, attachToNodeId: nodeId, }); } diff --git a/src/shared/pymodel/python-upstream-inference.ts b/src/shared/pymodel/python-upstream-inference.ts new file mode 100644 index 0000000..d7b0596 --- /dev/null +++ b/src/shared/pymodel/python-upstream-inference.ts @@ -0,0 +1,117 @@ +import type { + PythonModelConfig, + PythonModelUpstreamSource, +} from '@shared/framework/types'; + +import { normalizeUpstreamSources } from './python-upstream-sources'; + +export interface PythonModelCatalogDefaults { + catalog: string; + namespace: string; +} + +/** Resolve catalog.schema.table for a python model output table. */ +export function resolvePythonModelOutputFqn( + config: Pick< + PythonModelConfig, + 'name' | 'namespace' | 'table_name' | 'output' + >, + defaults: PythonModelCatalogDefaults, +): string { + const output = config.output ?? {}; + const catalog = output.database ?? defaults.catalog; + const schema = output.schema ?? config.namespace ?? defaults.namespace; + const table = output.table ?? config.table_name ?? config.name; + return `${catalog}.${schema}.${table}`; +} + +function existingTrinoValueKeys( + entries: PythonModelUpstreamSource[], +): Set { + const keys = new Set(); + for (const entry of entries) { + if (entry.type !== 'trino') { + continue; + } + keys.add(entry.value); + const parts = entry.value.split('.'); + if (parts.length >= 2) { + keys.add(parts.slice(-2).join('.')); + } + } + return keys; +} + +/** + * Infer trino upstream_sources entries from depends_on model names. + * Does not overwrite author-declared entries; returns only new trino refs. + */ +export function inferTrinoUpstreamFromDependsOn( + config: PythonModelConfig, + siblingConfigs: Map, + defaults: PythonModelCatalogDefaults, +): PythonModelUpstreamSource[] { + const dependsOn = config.depends_on ?? []; + if (dependsOn.length === 0) { + return []; + } + + const existingKeys = existingTrinoValueKeys( + normalizeUpstreamSources(config.upstream_sources), + ); + const inferred: PythonModelUpstreamSource[] = []; + + for (const depName of dependsOn) { + const upstream = siblingConfigs.get(depName); + if (!upstream) { + continue; + } + const fqn = resolvePythonModelOutputFqn(upstream, defaults); + const schemaTable = fqn.split('.').slice(-2).join('.'); + if (existingKeys.has(fqn) || existingKeys.has(schemaTable)) { + continue; + } + inferred.push({ type: 'trino', value: fqn }); + existingKeys.add(fqn); + existingKeys.add(schemaTable); + } + + return inferred; +} + +export const DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS: PythonModelCatalogDefaults = { + catalog: 'glue_development', + namespace: 'opus_python_source', +}; + +/** Apply depends_on inference and return a config with merged upstream_sources. */ +export function enrichPythonModelUpstreamSources( + config: PythonModelConfig, + siblingConfigs: Map, + defaults: PythonModelCatalogDefaults = DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS, +): PythonModelConfig { + const inferred = inferTrinoUpstreamFromDependsOn( + config, + siblingConfigs, + defaults, + ); + if (inferred.length === 0) { + return config; + } + return { + ...config, + upstream_sources: mergeInferredUpstreamSources(config, inferred), + }; +} + +/** Merge inferred trino entries onto existing upstream_sources (inferred appended). */ +export function mergeInferredUpstreamSources( + config: PythonModelConfig, + inferred: PythonModelUpstreamSource[], +): PythonModelUpstreamSource[] { + const existing = normalizeUpstreamSources(config.upstream_sources); + if (inferred.length === 0) { + return existing; + } + return [...existing, ...inferred]; +} diff --git a/src/shared/pymodel/python-upstream-sources.ts b/src/shared/pymodel/python-upstream-sources.ts index 11f6a50..7ca9d05 100644 --- a/src/shared/pymodel/python-upstream-sources.ts +++ b/src/shared/pymodel/python-upstream-sources.ts @@ -39,12 +39,27 @@ export function serializeUpstreamSourcesForProperty( /** * Parse `python_model_upstream_sources` from Iceberg table properties. - * Falls back to legacy comma-separated `schema.table` strings when JSON parse fails. + * + * The raw value may arrive as a JSON string (e.g. read verbatim from the + * property) OR as an already-parsed array/object: the Trino CSV reader + * (`parseTrinoValue`) eagerly `JSON.parse`s any cell starting with `[`/`{`, + * so the property surfaces as a structured value rather than a string. + * Both shapes are handled. Falls back to legacy comma-separated + * `schema.table` strings when a string fails JSON parsing. */ export function parseUpstreamSourcesProperty( - raw: string | undefined, + raw: unknown, ): PythonModelUpstreamSource[] { - if (!raw?.trim()) { + if (raw == null) { + return []; + } + + // Already parsed into an array/object by the Trino CSV reader. + if (typeof raw !== 'string') { + return normalizeUpstreamSources(raw); + } + + if (!raw.trim()) { return []; } @@ -64,3 +79,39 @@ export function parseUpstreamSourcesProperty( .filter(Boolean) .map((value) => ({ type: 'trino' as const, value })); } + +export interface ParsedTrinoUpstreamRef { + catalog: string; + schema: string; + table: string; +} + +/** + * Parse a trino upstream_sources value as catalog.schema.table (3-part) + * or schema.table (2-part, using fallbackCatalog). + */ +export function parseTrinoUpstreamValue( + value: string, + fallbackCatalog: string, +): ParsedTrinoUpstreamRef | null { + const trimmed = value.trim(); + if (!trimmed) { + return null; + } + const parts = trimmed.split('.'); + if (parts.length >= 3) { + return { + catalog: parts[0]!, + schema: parts[1]!, + table: parts.slice(2).join('.'), + }; + } + if (parts.length === 2) { + return { + catalog: fallbackCatalog, + schema: parts[0]!, + table: parts[1]!, + }; + } + return null; +} diff --git a/templates/skills/dj-create-python-model/_SKILL.md b/templates/skills/dj-create-python-model/_SKILL.md index 5583331..ec61c2a 100644 --- a/templates/skills/dj-create-python-model/_SKILL.md +++ b/templates/skills/dj-create-python-model/_SKILL.md @@ -3,9 +3,9 @@ name: dj-create-python-model description: >- Create a DJ .python.json file for a new Python ETL model. Use when the user wants to create a Python model, ETL pipeline, data ingestion, API fetch, - CSV import, or any pre-dbt Python data processing task. Not for dbt SQL models - (-> dj-create-new-model) or registering a raw table as a source (-> - dj-create-source). + CSV import, multi-model topic pipeline, or any Python Models data pipeline + task. Not for dbt SQL models (-> dj-create-new-model) or registering a raw + table as a source (-> dj-create-source). compatibility: DJ (Data JSON) Framework extension workspace with .dj/schemas/ metadata: dj-skill: '1.0' @@ -13,26 +13,23 @@ metadata: # Create DJ Python Model -**Create** new **`.python.json`** files for Python ETL models. These are **not** dbt SQL models — they are pre-dbt Python pipelines that extract data from external sources (APIs, databases, files) and load it into Iceberg tables. Downstream dbt SQL models then consume these tables as sources. +**Create** new **`.python.json`** + **`.python.py`** files for Python ETL models. These are **not** dbt SQL models — they are Python Models data pipelines that extract data from external sources (APIs, files, Jenkins jobs) or Trino tables and load results into Iceberg (or S3). Downstream dbt SQL models consume the output tables as sources. -**SQL-first principle:** Use **Trino SQL** for transformations and loading wherever possible. Trino manages compute, memory, partitioning, and table metadata — let it do the heavy lifting. Python's role is orchestration and external data ingestion; Trino's role is transformation and storage. +**Author choice:** Pick one orchestration style and stay consistent unless the user explicitly asks for mixed logic: -**Authoring mode** (from Step 0): In **`cells`** mode, only edit **`.python.json`** — DJ regenerates `.python.py` from the `cells` array. In **`hand-written`** mode, **`.python.py` is source of truth** and `.python.json` is metadata-only (empty or absent `cells` is OK). Never hand-edit **`_trino_io.py`** — it is auto-generated and always overwritten. +| Style | Orchestrator | Typical use | +| ----- | ------------ | ----------- | +| **Pandas** | `extract → transform → load → post_load` | External ingest, nested JSON, ML, notebook-style logic | +| **SQL / Trino** | `extract → transform_and_load → cleanup` | Warehouse-native transforms, large joins, default scaffold | +| **Mixed** | Combine stages when needed | Valid when the user requests it — not required | -## Step 0 — Resolve project defaults (before gathering) +There is **no SQL-first gate** and **no enforced hybrid** (stage-then-SQL) unless the author chooses that path. -1. Read `.agents/project/skills/dj-create-python-model/project-defaults.md` if it exists. -2. If missing, discover defaults from sibling models under `dags/python_models/`, python-source `.source.json` files, and `dbt_project.yml` vars. -3. Resolve and state: **authoring_mode** (`hand-written` | `cells`), **dev/prod catalogs**, **output schema**, **default DAG**, **source registration path**, optional **KPO image** / helper APIs. -4. Use placeholders below only when project defaults are unknown — confirm with the user before writing files. - -## DJ CLI (preferred when DJ is running) - -When `.dj/bin/dj system.ping` succeeds, use `trino.columns` and `trino.tables` to introspect warehouse tables before guessing column types; after the model has written output, use `source.create` to register the Iceberg table for downstream dbt (instead of hand-authoring `.source.json` from memory). Invocation patterns → `dj-cli`. Routing → `dj-cli-registry`. +**`.python.py`** is the **source of truth for code** — it is scaffolded once when the model is created, then authored/edited directly (cells are delimited with jupytext `# %%` markers). **`.python.json`** holds model metadata plus a non-authoritative `cells` snapshot of the last notebook save; it is **not** used to regenerate `.python.py`. **Never** hand-edit **`_trino_io.py`**, **`_config.py`**, or **`_model_tasks.py`** — these are auto-generated by DJ Framework and reverted on edit. ## When this skill applies -Use this skill when the user mentions: python model, ETL, data ingestion, API fetch, CSV import, pre-dbt processing, Python pipeline, data extraction, or creating a model that pulls data from an external system. +Use this skill when the user mentions: python model, ETL, data ingestion, API fetch, CSV import, data pipeline, Python pipeline, data extraction, **python model topic**, **multi-model pipeline**, or **dependency chain**. **Out of scope** — delegate to sibling skills: @@ -40,18 +37,25 @@ Use this skill when the user mentions: python model, ETL, data ingestion, API fe - `.source.json` files (registering a raw table as a source) → `dj-create-source` - Lightdash YAML → `dj-edit-lightdash-yaml` - Refactoring existing models → `dj-review-and-refactor-model` -- Migrating a legacy Jupyter notebook into a new python model → `dj-migrate-notebook-to-pymodel` (it hands off to this skill's workflow once the migration plan is approved) +- Migrating a legacy Jupyter notebook into a new python model → `dj-migrate-notebook-to-pymodel` - Verifying a python model's output table against a legacy/reference table → `dj-verify-pymodel-parity` - Documenting a topic's python models in a README → `dj-document-pymodels` +## Authoring mode + +| Mode | How | When | +| ---- | --- | ---- | +| **Cells (UI)** | Create Python Model webview, then write code in the notebook / `.python.py` | Default for new models; the webview scaffolds a stub `.python.py` | +| **Developer** | Author `.python.py` directly; edit `.python.json` for metadata; add topic-level `_helpers.py` for shared code | Multi-model topics, shared S3/API utilities | + +Shared topic helpers live at `dags/python_models///_helpers.py` (one per topic). Import from sibling models as needed. + ## Interactive gathering workflow -Ask the user the following questions in order. Batch related questions together to minimize round-trips. Skip questions when the answer is already clear from context. +Ask the following in order. Batch related questions to minimize round-trips. **Skip the Q&A entirely** when the user supplies a complete spec (name, group, topic, DAG, transform style, output, dependencies). ### Step 1: Identity -Ask for: - | Field | Rule | Example | | --------------- | ------------------------------------------------------------------------- | ----------------------------------------- | | **name** | `^[a-z][a-z0-9_]*$` | `backstage_catalogs` | @@ -59,293 +63,339 @@ Ask for: | **topic** | `^[a-z][a-z0-9_]*$` | `api_data` | | **description** | Free text (optional) | `Fetches catalog data from Backstage API` | -### Step 2: DAG assignment +### Step 2: Single model vs multi-model topic + +| Approach | How | Behavior | +| -------- | --- | -------- | +| **Single model (UI)** | Default **Create Python Model** form | One `.python.json` + scaffolded `.python.py` per submit | +| **Multi-model topic (agent / manual)** | Author `.python.json` files under the same `group` / `topic`, or use this skill to batch-write configs | Set per-model `name`, `description`, `depends_on`, and `upstream_sources`; DJ auto-infers trino `upstream_sources` from `depends_on` on JSON sync when omitted | + +`source_etl` is excluded from the Create Python Model DAG picker (DJ manages it via the generated template). DAG wiring is applied silently on create. + +### Step 3: DAG assignment + +Ask which Airflow DAG(s) to attach the model to. Look at existing DAGs in the project (`dags/` directory). Utility/library modules may use `"dags": []`. -Ask which Airflow DAG(s) to attach the model to. Look at existing DAGs in the project (`dags/` directory). If the model is a utility/library module, `dags` can be empty. +### Step 4: Transform strategy (author choice) -### Step 3: Data source type +Ask which orchestration style the author wants. If undecided, offer brief heuristics only — do not block on SQL: -Ask what kind of data source the model will extract from: +| Often pandas | Often SQL | +| ------------ | --------- | +| Nested JSON / API pagination | Filtering, joins, aggregations on warehouse data | +| ML preprocessing | Deduplication (`ROW_NUMBER`), window functions | +| Small external files | Large Trino-to-Trino transforms | + +### Step 5: Data source type | Source type | Typical packages | Extract pattern | | -------------------- | ---------------- | -------------------------------- | | **REST API** | `requests` | HTTP GET/POST with pagination | -| **Database / Trino** | `trino` | SQL query via Trino client | +| **Database / Trino** | `_trino_io` | `read_partition_to_dataframe` or SQL via `execute_sql` | | **CSV / file** | `pandas` | `pd.read_csv()` from local or S3 | | **S3 objects** | `boto3` | List/download from S3 bucket | | **Custom** | User-specified | User provides extract logic | -### Step 4: Transformation needs — SQL-first decision - -Ask what transformations are needed. Then apply this decision tree: - -**Can the transformation be expressed as SQL?** - -| If YES (use Trino SQL) | If NO (use pandas DataFrame) | -| ----------------------------- | ------------------------------------- | -| Filtering / WHERE clauses | Nested JSON flattening (dicts/lists) | -| Column renaming / aliasing | API response parsing / pagination | -| Type casting (CAST) | ML preprocessing (scikit-learn, etc.) | -| Deduplication (ROW_NUMBER) | Complex string parsing not in SQL | -| Aggregation (GROUP BY) | External service calls mid-transform | -| Joins with other Trino tables | Binary/image data processing | -| Date partitioning | | -| Window functions | | -| CASE expressions | | - -**Always ask:** "Can this transformation be done in Trino SQL?" before defaulting to pandas. SQL transformations are preferred because: +### Step 6: Output configuration -- Trino manages compute and memory — no Python OOM on large datasets -- Trino handles partitioned writes natively — no manual partition management -- SQL transformations are declarative and easier to reason about -- Trino pushes predicates down to storage — reads only needed data - -### Step 5: Output configuration - -Present the defaults and ask if the user wants to override any: +Present defaults and confirm overrides: | Field | Default | Override? | | --------------------- | ---------------------------- | ----------------------- | -| `output.database` | `` (from Step 0) | Yes | -| `output.schema` | `` (from Step 0) | Yes | +| `output.database` | `glue_development` | Yes | +| `output.schema` | `opus_python_source` | Yes | | `output_type` | `iceberg` | Yes (alternative: `s3`) | | `output.write_mode` | `overwrite_partitions` | Yes | | `output.partition_by` | `["portal_partition_daily"]` | Yes | | `namespace` | `python` | Yes | +| `compute` | `kpo` | Yes (`airflow` for MWAA worker) | +| `kpo_size` | `small` | Yes (when `compute` is `kpo`) | -**Confirm the write target before proceeding.** The output database/schema is where this model _writes_ data, and `overwrite_partitions` is destructive (it replaces existing partitions). **State the resolved `output.database` / `output.schema` / `write_mode` from Step 0 and get explicit confirmation** before writing files. Never point a new model at a production catalog without the user confirming it. - -### Step 6: Dependencies & optional fields - -Ask about: +**Confirm the write target before proceeding.** `overwrite_partitions` is destructive for affected partitions. State resolved `output.database` / `output.schema` / `write_mode` and get explicit confirmation when not using defaults. Never point a new model at production without confirmation. -- **Python packages** needed beyond the defaults (e.g., `requests>=2.28.0`, `boto3`) -- **depends_on** — other Python models that must complete first -- **upstream_sources** — upstream lineage for the output table (written to Iceberg properties after each run). Each entry is `{ "type": "trino"|"external", "value": "..." }` where `trino` values are `schema.table` and `external` values are free-form identifiers (API names, S3 paths, etc.) -- **tags** — organizational tags (predefined: `python-model`, `api`, `csv`, `s3`, `trino`, `snapshot`, `iceberg`) -- **owner** — team or individual -- **enable_notebook** — generate companion `.python.ipynb` (default: `true`) +### Step 7: Dependencies, lineage, and optional fields -**Governance metadata is optional — offer, never require.** Alongside `owner`, offer to tag `owner_slack`, `pii`, `classification`, and `compliance` in `meta` (see `.agents/dj/reference/meta-and-governance.md`), matching the keys sibling models in the project already use. If the user skips, write nothing and do not re-ask. - -## Optimization guidance - -**Proactively suggest these optimizations** to the user while building the model. Do not wait for the user to ask — surface relevant advice based on the model's source type, transformation needs, and data volume. - -### Partitioning strategy - -- **Default:** `portal_partition_daily` — suits most daily ingestion workloads -- **Low-volume data** (< 1000 rows/day): consider `portal_partition_monthly` to avoid small-file overhead -- **High-cardinality or large tables**: consider composite partitions (e.g., `["portal_partition_daily", "region"]`) for better predicate pushdown -- **Event streams**: partition by hour if sub-daily granularity is needed +| Field | Purpose | +| ----- | ------- | +| **depends_on** | Upstream **Python model names** — wires Airflow task order (model B waits for model A) | +| **upstream_sources** | Upstream **lineage** — written to Iceberg `python_model_upstream_sources` after each successful run | +| **dependencies** | Python pip packages (e.g. `requests>=2.28.0`) | +| **tags** | `python-model`, `api`, `csv`, `s3`, `trino`, `snapshot`, `iceberg` | +| **owner** | Team or individual | +| **enable_notebook** | Generate companion `.python.ipynb` (default: `true`) | -### SQL transform best practices +**`depends_on` vs `upstream_sources` — both matter:** -- **Never `SELECT *`** in production transforms — always enumerate needed columns explicitly -- **Always filter on partition columns** in WHERE clauses to enable predicate pushdown -- **Use `CAST(... AS type)`** in the SQL SELECT instead of pandas type conversion -- **Deduplicate with `ROW_NUMBER()`** window function instead of pandas `drop_duplicates` -- **Aggregations belong in SQL** — `GROUP BY` in Trino uses distributed compute, unlike single-node pandas -- **Joins with Trino tables** — join against existing Iceberg/Trino tables in SQL rather than merging DataFrames +- **`depends_on`** — scheduling only. Tells Airflow which sibling Python models must finish first. +- **`upstream_sources`** — lineage only. Tells Model Lineage which Trino tables and external systems this model reads. -### Write mode selection +When a model reads another Python model's output table, set **both**: `depends_on` for task order and a `trino` entry in `upstream_sources` for the graph edge. DJ auto-infers trino `upstream_sources` from `depends_on` on JSON sync when the author omits them (appended, never overwriting declared entries). -| Write mode | Use when | -| ---------------------- | ---------------------------------------------------------------------- | -| `overwrite_partitions` | Idempotent daily loads — rerun-safe, replaces only affected partitions | -| `append` | Event streams or append-only logs — never overwrites existing data | -| `overwrite` | Full table refresh — replaces entire table on each run | +**Governance metadata is optional — offer, never require.** Offer `owner_slack`, `pii`, `classification`, and `compliance` in `meta` (see `.agents/dj/reference/meta-and-governance.md`). If the user skips, write nothing. -### Resource management +### `upstream_sources` reference -- **Stage then transform**: for API/CSV sources, stage raw data into a temporary Trino table, then run SQL transforms. Drop the staging table after the final write -- **Batch large API responses**: chunk into batches (e.g., 10,000 rows) before staging to Trino to manage memory -- **Avoid loading entire datasets into pandas**: if data is already in Trino, query it with SQL — do not pull it into a DataFrame just to push it back - -## File path and naming +Declare in `.python.json` only — DJ writes properties after each run; **do not** call `ALTER TABLE` for lineage in `.python.py`. -``` -dags/python_models///.python.json +```json +"upstream_sources": [ + { "type": "trino", "value": "glue_development.opus_python_source.raw_events" }, + { "type": "external", "value": "api::backstage_catalog" }, + { "type": "external", "value": "s3::my-bucket/path/data.csv" }, + { "type": "external", "value": "jenkins::my-job-name" } +] ``` -Model ID convention: `python______` +| `type` | `value` format | Notes | +| ------ | -------------- | ----- | +| **trino** | **`catalog.schema.table`** (canonical) | e.g. `glue_development.opus_python_source.model_a` | +| **trino** | `schema.table` (legacy) | Still parsed; prefer 3-part FQN for new models | +| **external** | Prefix convention | `s3::…`, `api::…`, `jenkins::…`, or other stable identifier | -## Compute and KPO image naming +## KPO and execution requirements -New models default to **`compute: "kpo"`** (Kubernetes pod) instead of running on the MWAA worker. Set **`compute: "airflow"`** explicitly when the KPO image pipeline is not ready. +New models default to **`compute: "kpo"`** (Kubernetes pod). Set **`compute: "airflow"`** when the KPO image pipeline is not ready. | Field | Purpose | | ----- | ------- | | `compute` | `kpo` (default) or `airflow` | | `kpo_size` | Pod preset when `compute` is `kpo`: `small`, `medium`, `large`, `xlarge` | -**KPO script path:** The `mwaa-opus-kpo-runner` image syncs `s3:///dags/python_models/` into the pod and runs `python //.python.py`. Airflow passes that relative path as the container argument. Set Airflow Variable **`python_models_s3_bucket`** (DAG bucket name) on KPO tasks; optional **`python_model_s3_region`** (default `us-west-2`). +**KPO script path:** The `mwaa-opus-kpo-runner` image syncs `s3:///dags/python_models/` into the pod and runs `python //.python.py`. Set Airflow Variable **`python_models_s3_bucket`**; optional **`python_model_s3_region`** (default `us-west-2`). Configure **`airflow_runner_cfg`** (image, tag, namespace, cluster). + +**`PYMODEL_*` environment:** Airflow sets `PYMODEL_DS`, `PYMODEL_DS_NODASH`, and optionally `PYMODEL_DATES` (JSON array) before execution. `build_context_from_env()` in `_config.py` reads them and configures `logging.basicConfig` to stdout (required for KPO log visibility). + +**`__main__` block (required):** Scaffolded `.python.py` must end with: -**Execution entry:** Scaffolded `.python.py` ends with `if __name__ == "__main__":` calling `run_etl(build_context_from_env())`. `build_context_from_env()` in `_config.py` reads `PYMODEL_DS`, `PYMODEL_DS_NODASH`, and `PYMODEL_DATES` from the environment. MWAA worker and KPO pod both run the file as `__main__` — `run_etl` is scaffold convention, not a framework-enforced function name. +```python +if __name__ == "__main__": + from python_models._config import build_context_from_env + run_etl(build_context_from_env()) +``` -**Airflow UI:** Each KPO model is **one task** named after the model `name` field (e.g. `kpo_infra`). Gating and pod launch run inside that task — there are no separate gate/KPO tasks. +MWAA worker and KPO pod both execute the file as `__main__`. `run_etl` is scaffold convention, not a framework-enforced name. -Configure **`airflow_runner_cfg`** (JSON: `image`, `tag`, `namespace`, `config_file`, `cluster_context`; optional `size_specs` override for pod presets in `dags/_ext_/kpo_sizes.yml`). +**Airflow UI:** Each KPO model is **one task** named after the model `name` field. -## Schema reference +## File path and naming + +``` +dags/python_models///.python.json +``` -Before writing, read `.dj/schemas/python-model.schema.json` to validate field shapes. Required fields: `name`, `group`, `topic`. +Model ID: `python______` ## Workflow checklist -- [ ] Step 1: Gather identity (name, group, topic, description) -- [ ] Step 2: Determine DAG assignment -- [ ] Step 3: Identify data source type and extraction pattern -- [ ] Step 4: Determine transformation needs (SQL-first decision tree) -- [ ] Step 0: Read project defaults (authoring mode, catalog, schema, DAG) -- [ ] Step 5: Confirm output config (`` / `` from Step 0) -- [ ] Step 6: Collect dependencies and optional fields -- [ ] Step 7: Suggest performance optimizations based on source type and data volume -- [ ] Step 8: Read `.dj/schemas/python-model.schema.json` for validation +- [ ] Step 1: Gather identity (or accept complete spec) +- [ ] Step 2: Single model vs topic workflow +- [ ] Step 3: Determine DAG assignment +- [ ] Step 4: Confirm transform strategy (pandas, SQL, or mixed) +- [ ] Step 5: Identify data sources +- [ ] Step 6: Confirm output config and compute (`kpo` default) +- [ ] Step 7: Set `depends_on`, `upstream_sources`, and optional fields +- [ ] Step 8: Read `.dj/schemas/python-model.schema.json` - [ ] Step 9: Write `.python.json` at `dags/python_models///.python.json` -- [ ] Step 10: Verify the file matches the schema +- [ ] Step 10: Verify schema compliance and `__main__` entry in the scaffolded `.python.py` ## ETL cell structure -The `cells` array in `.python.json` uses notebook format. Each cell has `cell_type` (`code` or `markdown`), `source` (array of strings, each ending with `\n`), and optional `metadata`, `outputs`, `execution_count`. - -### SQL-first flow (preferred) +The `cells` array uses notebook format: `cell_type` (`code` | `markdown`), `source` (array of strings ending with `\n`), optional `metadata`, `outputs`, `execution_count`. -Use this when data comes from an external source (API, CSV, S3) but all transformations can be expressed in SQL. All Trino I/O uses the shared `_trino_io.py` helpers — **do not write inline Trino connection code**. +### Pandas orchestrator 1. **Markdown header** — model name and DAG info -2. **Imports + context** — standard imports, `from python_models._trino_io import execute_trino, overwrite_partition`, and Airflow context dict -3. **Output config** — `PythonModelConfig` from `python_models/_config.py` +2. **Imports + logging** — `pandas`, domain packages, `from python_models._trino_io import read_partition_to_dataframe, overwrite_partition_from_dataframe` +3. **Output config** — `PythonModelConfig` from `python_models._config` 4. **Input variables** — model-specific constants -5. **Extract** — `def extract(context)` fetches external data (API/CSV/S3) and stages it into a temp Trino table -6. **Transform + Load** — `def transform_and_load(context)` runs SQL `INSERT INTO ... SELECT` with all transformations, using `overwrite_partition()` or other DML helpers from `_trino_io` -7. **Cleanup** — `def cleanup(context)` drops the temporary staging table via `execute_trino()` -8. **Main** — `def run_etl(context)` calls extract -> transform_and_load -> cleanup -9. **Script entry** — `if __name__ == "__main__":` calls `run_etl(build_context_from_env())` (included in `.python.py`) -10. **Runner** — notebook-only `run_etl(context)` call for interactive use +5. **`extract(context)`** — fetch external data or read Trino partitions into a DataFrame +6. **`transform(df, context)`** — pandas transforms; set `portal_partition_daily` +7. **`load(df, context)`** — `overwrite_partition_from_dataframe(table_fqn, …)` +8. **`post_load(context)`** — optional validation via `fetch_value` +9. **`run_etl(context)`** — orchestrate the stages +10. **Script entry** — `if __name__ == "__main__": run_etl(build_context_from_env())` +11. **Runner cell (optional, notebook only)** — a bare `run_etl(context)` for interactive notebook runs; keep it out of `.python.py` module scope so the file stays import-safe (the `__main__` block is the script entry) + +### SQL / Trino orchestrator + +1. **Markdown header** +2. **Imports** — `from python_models._trino_io import execute_sql, fetch_value, overwrite_partition` +3. **Output config** — `PythonModelConfig`, `get_catalog_name` +4. **Input variables** +5. **`extract(context)`** — optional staging (external → temp table via `create_iceberg_table` + batched inserts, or skip when source is already in Trino) +6. **`transform_and_load(context)`** — SQL transforms; `overwrite_partition(table_fqn, …)` +7. **`cleanup(context)`** — `execute_sql("DROP TABLE IF EXISTS …")` for staging tables +8. **`run_etl(context)`** +9. **Script entry** — `if __name__ == "__main__": run_etl(build_context_from_env())` +10. **Runner cell (optional, notebook only)** — bare `run_etl(context)` for interactive runs; not part of the `.python.py` module scope + +### Table FQN (`table_fqn`) + +All `_trino_io` helpers take a fully-qualified table string. Build it from Airflow variables and output config: -### `_trino_io.py` DML operations reference +```python +table_fqn = f'{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name or OUTPUT_CONFIG.model_name}' +``` -The extension auto-generates `dags/python_models/_trino_io.py` with these functions. **Do not edit this file** — it is always overwritten by the extension. +Use the same pattern for upstream tables declared in `upstream_sources`. -| Function | Use when | -| ---------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------- | -| `execute_trino(sql)` | DDL/DML with no return value (CREATE, DROP, INSERT, etc.) | -| `execute_trino(sql, return_result=True)` | Scalar reads (SELECT COUNT(\*), SELECT MAX(...), etc.) — returns first column of first row | -| `append(table_fqn, insert_sql)` | Append-only INSERT — runs the user's INSERT SQL as-is | -| `overwrite_partition(table_fqn, partition_col, partition_value, *, insert_sql=...)` | Idempotent daily loads — DELETE one partition, then INSERT (pass a full `INSERT INTO ...` statement) | -| `overwrite_partition(table_fqn, partition_col, partition_value, source_query, *, columns=...)` | Same as above, but DJ builds `INSERT INTO ... (cols) ` from a SELECT / WITH ... SELECT | -| `overwrite(table_fqn, overwrite_filter, *, insert_sql=...)` | Overwrite rows matching any SQL filter (use `"true"` for full-table refresh) | -| `merge(table_fqn, merge_sql)` | Upsert — runs the user's full MERGE INTO ... USING ... SQL | -| `delete(table_fqn, where_clause)` | Standalone conditional DELETE | -| `update(table_fqn, set_clause, where_clause)` | Standalone conditional UPDATE | +### `_trino_io.py` API reference -`insert_sql` is keyword-only. Prefer `insert_sql=...` when you already have a full INSERT; use `source_query=` (+ optional `columns=`) when you want the helper to assemble the INSERT. +The extension auto-generates `dags/python_models/_trino_io.py`. **Do not edit it.** -**Example — overwrite partition with full INSERT (most common):** +**SQL path:** -```python -from python_models._trino_io import overwrite_partition +| Function | Use when | +| -------- | -------- | +| `execute_sql(sql)` | DDL/DML with no return value (CREATE, DROP, INSERT, …) | +| `fetch_value(sql)` | Scalar reads (`SELECT COUNT(*)`, `SELECT MAX(…)`) — first column of first row | +| `overwrite_partition(table_fqn, partition_col, partition_value, *, insert_sql=…)` | Idempotent partition load — DELETE partition, then INSERT | +| `overwrite_partition(…, source_query=…, columns=[…])` | Same, but helper assembles `INSERT INTO … (cols) ` | +| `overwrite(table_fqn, filter, *, insert_sql=…)` | Conditional overwrite (`"true"` for full refresh) | +| `append(table_fqn, insert_sql)` | Append-only INSERT | +| `merge(table_fqn, merge_sql)` | Upsert via `MERGE INTO …` | +| `update(table_fqn, set_clause, where_clause)` | Standalone UPDATE | +| `create_iceberg_table(…)`, `create_schema(…)`, `set_extra_properties(…)` | DDL / metadata | -def transform_and_load(context): - ds = context["ds"] - table = OUTPUT_CONFIG.full_table_id - insert_sql = f""" - INSERT INTO {table} - SELECT col_a, col_b, '{ds}' AS portal_partition_daily - FROM iceberg_catalog.stg_tmp.my_staging_table - """ - overwrite_partition(table, "portal_partition_daily", ds, insert_sql=insert_sql) -``` +**Pandas path (python-path models only):** + +| Function | Use when | +| -------- | -------- | +| `read_partition_to_dataframe(table_fqn, partition_col, partition_value, columns=None)` | Read one partition slice into a DataFrame | +| `overwrite_partition_from_dataframe(table_fqn, partition_col, partition_value, df, columns, batch_size=1000)` | Idempotent partition write from a DataFrame | -**Example — overwrite partition with source_query:** +**Do not use:** `execute_trino`, inline `trino.dbapi.connect`, raw HTTP to Trino, or standalone `delete()` — partition overwrites go through `overwrite_partition` / `overwrite_partition_from_dataframe`. + +**Example — SQL partition load:** ```python from python_models._trino_io import overwrite_partition def transform_and_load(context): ds = context["ds"] - table = OUTPUT_CONFIG.full_table_id - source_query = f""" + table_fqn = f"{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" + insert_sql = f""" + INSERT INTO {table_fqn} SELECT col_a, col_b, '{ds}' AS portal_partition_daily - FROM iceberg_catalog.stg_tmp.my_staging_table + FROM {get_catalog_name()}.stg_tmp.my_staging_table """ - overwrite_partition( - table, - "portal_partition_daily", - ds, - source_query=source_query, - columns=["col_a", "col_b", "portal_partition_daily"], - ) + overwrite_partition(table_fqn, "portal_partition_daily", ds, insert_sql=insert_sql) ``` -**Example — merge (upsert):** +**Example — pandas partition load:** ```python -from python_models._trino_io import merge +from python_models._trino_io import overwrite_partition_from_dataframe -def transform_and_load(context): - table = OUTPUT_CONFIG.full_table_id - merge_sql = f""" - MERGE INTO {table} AS t - USING iceberg_catalog.stg_tmp.customers_staging AS s - ON t.customer_id = s.customer_id - WHEN MATCHED THEN UPDATE SET name = s.name, email = s.email - WHEN NOT MATCHED THEN INSERT (customer_id, name, email) - VALUES (s.customer_id, s.name, s.email) - """ - merge(table, merge_sql) +def load(df, context): + ds = context["ds"] + table_fqn = f"{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" + columns = ["col_a", "col_b", "portal_partition_daily"] + overwrite_partition_from_dataframe( + table_fqn, "portal_partition_daily", ds, df, columns + ) ``` **Example — scalar validation:** ```python -from python_models._trino_io import execute_trino +from python_models._trino_io import fetch_value -def validate(context): +def post_load(context): ds = context["ds"] - count = execute_trino(f""" - SELECT COUNT(*) FROM {OUTPUT_CONFIG.full_table_id} + table_fqn = f"{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" + count = fetch_value(f""" + SELECT COUNT(*) FROM {table_fqn} WHERE portal_partition_daily = '{ds}' - """, return_result=True) - if count == 0: + """) + if not count: raise RuntimeError(f"No rows loaded for {ds}") ``` ### Code templates and worked example -For complete code templates for each ETL stage (extract, stage, transform+load, cleanup, orchestrators), see [references/etl-patterns.md](references/etl-patterns.md). +- Stage-level patterns: [references/etl-patterns.md](references/etl-patterns.md) +- Topic demo (pandas models + trino upstream reads): [references/worked-example.md](references/worked-example.md) +- Evaluation scenarios (prompt → expected behavior): [evals/evals.md](evals/evals.md) + +## After creating — developer checklist + +1. **Code** — author the ETL directly in `.python.py` (scaffolded on create); confirm it has the `__main__` block and the chosen `_trino_io` imports. +2. **`depends_on`** — every upstream Python model this one waits on is listed. +3. **`upstream_sources`** — every Trino table and external system read by the model is declared (`catalog.schema.table` for trino). +4. **Output target** — `output.database` / `output.schema` / `write_mode` match the intended environment. +5. **Partition column** — `portal_partition_daily` is set in output (SQL `'{ds}' AS portal_partition_daily` or pandas assignment) when using partition writes. +6. **DAG wiring** — model appears in the selected DAG; KPO variables (`python_models_s3_bucket`, `airflow_runner_cfg`) are set. +7. **Register downstream source** — after first successful run, create `.source.json` via `dj-create-source` so dbt models can reference the table. +8. **Topic README (optional)** — run `dj-document-pymodels` for multi-model topics. + +### Lineage expectations + +After a successful Airflow or KPO run, DJ writes Iceberg table properties on the output table: -For a complete `.python.json` file showing the SQL-first pattern end-to-end, see [references/worked-example.md](references/worked-example.md). +- `python_model_name`, `python_model_table`, `python_model_namespace`, `python_model_description` +- `python_model_upstream_sources` — JSON from `.python.json` `upstream_sources` (plus any auto-inferred trino entries) + +Model Lineage reads these properties and walks upstream recursively: + +- **Python model nodes** — output tables with `python_model_name` +- **Trino ref nodes** — `type: trino` upstreams (`catalog.schema.table`) +- **External nodes** — `type: external` upstreams (`s3::…`, `api::…`, etc.) + +Chains like downstream dbt → model C → model B + model A render in one graph when each hop declares `upstream_sources`. ## Downstream dbt integration -After the Python model writes data to `..`, downstream dbt SQL models reference it as a source using the DJ source notation: +Downstream dbt SQL models reference the output table using DJ source notation: ``` -__. +glue_development__opus_python_source. ``` -(Double underscore separates catalog from schema, dot separates schema from table.) +(Double underscore separates catalog from schema.) + +Register the output table as a **`.source.json`** before a downstream model can read it — use **`dj-create-source`** or the `DJ: Create Source` webview, then sync. + +## Optimization guidance + +Surface relevant advice based on source type and volume; do not mandate SQL. + +### Partitioning + +- **Default:** `portal_partition_daily` for daily loads +- **Low volume:** `portal_partition_monthly` to reduce small files +- **High cardinality:** composite partitions for predicate pushdown + +### Write modes + +| Write mode | Use when | +| ---------- | -------- | +| `overwrite_partitions` | Idempotent daily loads | +| `append` | Append-only event streams | +| `overwrite` | Full table refresh | + +### SQL transforms (when using SQL path) + +- Enumerate columns explicitly — avoid `SELECT *` +- Filter on partition columns in WHERE clauses +- Use `CAST(… AS type)` for type safety +- Prefer `ROW_NUMBER()` for deduplication over pandas `drop_duplicates` + +### Pandas transforms (when using pandas path) -The output table must be **registered as a `.source.json`** before a downstream model can read it — use the **`dj-create-source`** skill (it introspects the Iceberg table's columns with `SHOW COLUMNS`) or the `DJ: Create Source` webview, then sync. +- Batch large writes via `overwrite_partition_from_dataframe(…, batch_size=…)` +- Read upstream partitions with `read_partition_to_dataframe` instead of pulling full tables +- Chunk API pagination before building large DataFrames ## Conventions and gotchas -- **SQL-first** — prefer Trino SQL for all transformations that SQL can express. Use DataFrames only for external data ingestion and Python-only logic -- **`cells` mode:** only edit `.python.json` — the extension regenerates `.python.py` (and `.python.ipynb`) from the `cells` array. **`hand-written` mode:** edit `.python.py` directly; keep JSON metadata in sync -- **`if __name__ == "__main__"`** — scaffold includes this block in `.python.py`; Airflow and KPO execute the script as `__main__` (use `build_context_from_env()` from `_config.py` for partition dates) -- **`run_etl()`** — scaffold convention for orchestrating extract/transform/cleanup; not required by discovery or execution -- **`OUTPUT_CONFIG`** must use `PythonModelConfig` from `python_models/_config.py` — the extension auto-generates this file -- **`portal_partition_daily`** must be set as a column (either via SQL `'{ds}' AS portal_partition_daily` or pandas assignment) for partition-based writes -- **Always suggest optimizations** — partitioning strategy, column pruning, predicate pushdown, write mode selection -- **Stage then SQL-transform** — for external data sources, stage raw data into a temp Trino table, run SQL transforms against it, then drop the staging table -- **Never `SELECT *`** in production SQL — enumerate columns explicitly for clarity and performance -- **Drop staging tables** — always clean up `stg_tmp_*` tables in the cleanup step -- **Batch INSERT for staging** — chunk large DataFrames into batches (e.g., 1000 rows) when staging to Trino -- **`depends_on`** ordering is enforced by Airflow task dependency wiring -- **`_config.py`, `_trino_io.py`, and `etl_helper.py`** are auto-created by the extension — do not manually create or edit them -- **Use `_trino_io` helpers** — import `execute_trino`, `overwrite_partition`, etc. from `python_models._trino_io` instead of writing inline Trino connection code -- **Cell `source` format** — each element in the `source` array is a string ending with `\n`; the array represents lines -- **The notebook runner cell** (`run_etl(context)`) is stripped when generating `.python.py` — keep it as the last cell for interactive notebook use -- **Name uniqueness** — model names must be unique across all groups/topics within the project +- **`.python.py` is the code source of truth** — edit it directly; `.python.json` holds metadata + a non-authoritative `cells` snapshot; `.python.ipynb` is an ephemeral view (in `.dj/.python_temp/`, auto-deleted on close) +- **`get_catalog_name()` / `get_namespace_name()`** — resolve catalog and schema from Airflow Variables; do not hardcode in production paths +- **Lineage in JSON only** — `upstream_sources` on `.python.json`, not manual property DDL in cells +- **`portal_partition_daily`** — required for partition-based writes +- **`depends_on`** — Airflow task ordering; **`upstream_sources`** — lineage graph edges; set both when reading sibling Python outputs +- **`_config.py`, `_trino_io.py`, `_model_tasks.py`, `etl_helper.py`** — auto-generated; do not hand-edit (edits are reverted) +- **Use `_trino_io` helpers** — no inline Trino connection code +- **Cell `source` format** — each element ends with `\n` +- **Notebook header** — the `# Python Model: …` cell is derived/notebook-only and never written back to `.python.py` +- **Name uniqueness** — model names unique across the project diff --git a/templates/skills/dj-create-python-model/evals/evals.md b/templates/skills/dj-create-python-model/evals/evals.md new file mode 100644 index 0000000..8879b45 --- /dev/null +++ b/templates/skills/dj-create-python-model/evals/evals.md @@ -0,0 +1,67 @@ +# Evals — dj-create-python-model + +Prompt → expected-behavior cases for authoring Python Models data pipelines. Each case lists the user prompt, the artifacts the skill should produce, and the checks that must hold. Use these to sanity-check the skill after edits. + +Reference the correct file model throughout: + +- `.python.py` — code source of truth (scaffolded once, then authored directly). +- `.python.json` — metadata + a non-authoritative `cells` snapshot. +- `_trino_io.py` / `_config.py` / `_model_tasks.py` — generated & guarded; never hand-edited. + +--- + +## E1 — Single pandas API-ingest model + +**Prompt:** "Create a Python model that fetches the Backstage catalog API into an Iceberg table." + +**Expected:** +- Files at `dags/python_models///backstage_catalogs.python.json` + `.python.py`. +- Pandas orchestrator: `extract → transform → load → post_load`, with `run_etl` and the `__main__` entry. +- `output.database` / `output.schema` / `write_mode` stated and confirmed before writing. +- `upstream_sources` contains an `external` entry (`api::…`); `dependencies` lists `requests`. +- `portal_partition_daily` set in `transform`/`load`. + +**Fails if:** it claims `.python.py` is regenerated from `cells`; hardcodes catalog/schema instead of `get_catalog_name()`; omits the `__main__` block. + +## E2 — SQL / Trino warehouse transform + +**Prompt:** "Create a Python model that deduplicates warehouse events with ROW_NUMBER and writes daily partitions." + +**Expected:** +- SQL orchestrator: `extract` (optional) → `transform_and_load` → `cleanup`. +- Uses `_trino_io` helpers (`execute_sql`, `overwrite_partition`) — no inline `trino.dbapi.connect`. +- Explicit column list (no `SELECT *`); partition predicate on the source read. +- `upstream_sources` lists each Trino table read as `catalog.schema.table`. + +**Fails if:** it emits inline Trino connection code; uses a 2-part FQN for new models without noting it is legacy. + +## E3 — Multi-model topic with dependency chain + +**Prompt:** "Create two models in the same topic where `tenant_status` reads the output of `tenant_events`." + +**Expected:** +- Two `.python.json` + `.python.py` pairs under the same `group`/`topic`. +- `tenant_status` sets **both** `depends_on: ["tenant_events"]` (task order) **and** a `trino` `upstream_sources` entry for `tenant_events`'s output table (lineage edge). +- Notes that DJ auto-infers the trino upstream from `depends_on` on sync when omitted (appended, never overwriting declared entries). + +**Fails if:** it sets only one of `depends_on` / `upstream_sources` when a sibling output is read. + +## E4 — Compute / KPO configuration + +**Prompt:** "Create a model that runs on a large KPO pod." + +**Expected:** +- `compute: "kpo"` (default) with `kpo_size: "large"`. +- Mentions the required Airflow Variable `python_models_s3_bucket` and `airflow_runner_cfg`. +- One Airflow task named after the model `name`. + +**Fails if:** it invents a per-cell compute field or omits `kpo_size` when `compute` is `kpo`. + +## E5 — Destructive-write confirmation + +**Prompt:** "Point this model at `glue_production.opus_python_source`." + +**Expected:** +- The skill states the resolved `output.database` / `output.schema` / `write_mode` and asks for explicit confirmation before writing to a non-default (production) target, because `overwrite_partitions` is destructive. + +**Fails if:** it writes to production without confirmation. diff --git a/templates/skills/dj-create-python-model/references/etl-patterns.md b/templates/skills/dj-create-python-model/references/etl-patterns.md index d40490c..a8d66ed 100644 --- a/templates/skills/dj-create-python-model/references/etl-patterns.md +++ b/templates/skills/dj-create-python-model/references/etl-patterns.md @@ -1,59 +1,51 @@ # ETL Pattern Reference -Detailed patterns for each ETL stage. **SQL-first:** all transformation and loading patterns use Trino SQL by default. DataFrame patterns are provided only for cases where SQL cannot express the logic. +Stage-level patterns for both orchestrators. All Trino I/O goes through **`python_models._trino_io`** — no inline `trino.dbapi.connect`, no `execute_trino`, no standalone `delete()`. -All patterns target `.` by default. +Default output target: `glue_development.opus_python_source` (override via `output` in `.python.json` and `get_catalog_name()` / `get_namespace_name()` at runtime). ---- - -## Trino Connection Helper - -Include this in all models that use Trino for staging or transformation. All patterns below assume this function is defined: +Build `table_fqn` for every helper call: ```python -def get_trino_conn(): - """Reusable Trino connection.""" - from trino.dbapi import connect - return connect( - host=INPUT_VARIABLES.get("trino_host", "localhost"), - port=int(INPUT_VARIABLES.get("trino_port", "8080")), - user=INPUT_VARIABLES.get("trino_user", "etl"), - catalog="", - schema="", - ) +from python_models._config import PythonModelConfig, get_catalog_name + +table_fqn = f"{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name or OUTPUT_CONFIG.model_name}" ``` --- -## Orchestrator Patterns +## Orchestrator patterns -### SQL-first flow (preferred) +### Pandas orchestrator ```python def run_etl(context: dict): - """Main ETL orchestrator — SQL-first flow.""" + """Pandas end-to-end flow.""" df = extract(context) - stage(df, context) - transform_and_load(context) - cleanup(context) + df = transform(df, context) + load(df, context) + post_load(context) ``` -### DataFrame fallback flow +### SQL / Trino orchestrator ```python def run_etl(context: dict): - """Main ETL orchestrator — DataFrame fallback for Python-only transforms.""" - df = extract(context) - df = transform(df, context) - load(df, context) - post_load(context) + """SQL end-to-end flow.""" + extract(context) # optional — stage external data or no-op + transform_and_load(context) + cleanup(context) ``` +### Mixed (author choice) + +Combine stages when the user requests it — e.g. pandas `extract` + SQL `transform_and_load` after staging. Not required by the framework. + --- -## Extract Patterns +## Extract patterns -Extract is always Python — it pulls data from external systems into a DataFrame (or skips extraction if data is already in Trino). +Extract is Python — it pulls data from external systems or reads Trino partitions into a DataFrame. Skip when the SQL path reads directly from existing Trino tables. ### REST API with pagination @@ -94,7 +86,6 @@ def extract(context: dict) -> pd.DataFrame: log.info(f"Page {page}: fetched {len(records)} records (total: {len(all_records)})") page += 1 - log.info(f"Total records extracted: {len(all_records)}") return pd.DataFrame(all_records) ``` @@ -178,468 +169,338 @@ def extract(context: dict) -> pd.DataFrame: return pd.concat(frames, ignore_index=True) ``` -### Data already in Trino (skip extract) +### Read upstream Trino partition (pandas path) + +When `upstream_sources` declares a trino table, read via `_trino_io`: + +```python +from python_models._trino_io import read_partition_to_dataframe + +def extract(context: dict) -> pd.DataFrame: + """Read an upstream Iceberg partition declared in upstream_sources.""" + ds = context["ds"] + upstream_fqn = f"{get_catalog_name()}.opus_python_source.upstream_model_a" + return read_partition_to_dataframe( + upstream_fqn, + "portal_partition_daily", + ds, + columns=["id", "name", "amount"], + ) +``` + +Declare in `.python.json`: + +```json +"upstream_sources": [ + { "type": "trino", "value": "glue_development.opus_python_source.upstream_model_a" } +] +``` -When the source data is already available in Trino, skip extraction entirely and go straight to SQL transform + load: +### Data already in Trino (SQL path — skip extract) ```python -def extract(context: dict): - """No extraction needed — source data is already in Trino.""" - log.info("Source data is in Trino — skipping extract, will transform via SQL") - return None +def extract(context: dict) -> None: + """No extraction — transform_and_load reads directly from Trino.""" + log.info("Source data is in Trino — skipping extract") ``` --- -## Stage Pattern +## Pandas transform patterns -Stage a raw DataFrame into a temporary Trino table so all subsequent transformation happens in SQL. **Always batch INSERTs** to manage memory. +```python +def transform(df: pd.DataFrame, context: dict) -> pd.DataFrame: + """Pandas transforms — set partition column before load.""" + if df is None or df.empty: + return pd.DataFrame() + + # Example: nested JSON flattening + if "metadata" in df.columns: + meta_df = pd.json_normalize(df["metadata"].apply( + lambda x: x if isinstance(x, dict) else {} + )) + meta_df.columns = [f"meta_{c}" for c in meta_df.columns] + df = pd.concat([df.drop(columns=["metadata"]), meta_df], axis=1) + + df["portal_partition_daily"] = context["ds"] + return df +``` + +### Merge multiple upstream DataFrames ```python -def stage(df: pd.DataFrame, context: dict) -> None: - """Stage raw data into a temporary Trino table.""" +def transform(dfs: dict[str, pd.DataFrame], context: dict) -> pd.DataFrame: + """Join upstream partitions in pandas (integration model pattern).""" + merged = dfs["model_a"].merge(dfs["model_b"], on="tenant_id", how="outer") + merged["portal_partition_daily"] = context["ds"] + return merged +``` + +--- + +## Pandas load pattern + +```python +from python_models._trino_io import overwrite_partition_from_dataframe + +def load(df: pd.DataFrame, context: dict) -> None: + """Write partition output from a DataFrame.""" if df is None or df.empty: - log.warning("No data to stage") + log.warning("No data to load") return - conn = get_trino_conn() - cursor = conn.cursor() - table_name = OUTPUT_CONFIG.table_name or OUTPUT_CONFIG.model_name - staging_table = f"stg_tmp_{table_name}" + ds = context["ds"] + table_fqn = f"{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" + columns = ["tenant_id", "status", "portal_partition_daily"] + overwrite_partition_from_dataframe( + table_fqn, + "portal_partition_daily", + ds, + df, + columns, + batch_size=1000, + ) + log.info(f"Loaded {len(df)} rows into {table_fqn} for {ds}") +``` - cursor.execute(f"DROP TABLE IF EXISTS ..{staging_table}") +--- - columns = df.columns.tolist() - col_defs = ", ".join(f"{c} VARCHAR" for c in columns) - cursor.execute( - f"CREATE TABLE ..{staging_table} ({col_defs})" - ) +## SQL stage pattern (optional — SQL path with external extract) - batch_size = 1000 - for i in range(0, len(df), batch_size): - batch = df.iloc[i:i + batch_size] - values_list = [] - for _, row in batch.iterrows(): - vals = ", ".join( - f"'{str(v).replace(chr(39), chr(39)+chr(39))}'" for v in row - ) - values_list.append(f"({vals})") - cursor.execute( - f"INSERT INTO ..{staging_table} " - f"VALUES {', '.join(values_list)}" - ) +Stage raw external data into a temporary table before SQL transforms. Use `create_iceberg_table` + batched `execute_sql` INSERTs, or stage in pandas and load via `overwrite_partition_from_dataframe` to a `stg_tmp_*` table. + +```python +from python_models._trino_io import create_iceberg_table, overwrite_partition_from_dataframe + +def stage(df: pd.DataFrame, context: dict) -> None: + """Stage raw API data into a temporary Iceberg table for SQL transforms.""" + if df is None or df.empty: + log.warning("No data to stage") + return - log.info(f"Staged {len(df)} rows into {staging_table}") + catalog = get_catalog_name() + staging_fqn = f"{catalog}.{OUTPUT_CONFIG.namespace}.stg_tmp_{OUTPUT_CONFIG.table_name}" + columns = [(c, "VARCHAR") for c in df.columns.tolist()] + + create_iceberg_table(staging_fqn, columns) + # Write all rows to a single staging partition (ds-scoped for cleanup) + df = df.copy() + df["_stage_partition"] = context["ds"] + staging_columns = list(df.columns) + overwrite_partition_from_dataframe( + staging_fqn, "_stage_partition", context["ds"], df, staging_columns + ) + log.info(f"Staged {len(df)} rows into {staging_fqn}") ``` --- -## Transform + Load Patterns (SQL-first) +## SQL transform + load patterns -All transformations are SQL expressions executed by Trino. Trino manages the Iceberg table writes. +All partition loads use `overwrite_partition` — not manual `DELETE` + `INSERT`. ### Column renaming and type casting ```python +from python_models._trino_io import overwrite_partition + def transform_and_load(context: dict) -> None: - conn = get_trino_conn() - cursor = conn.cursor() ds = context["ds"] - target = f"..{OUTPUT_CONFIG.table_name}" - staging = f"..stg_tmp_{OUTPUT_CONFIG.table_name}" - - cursor.execute(f""" - CREATE TABLE IF NOT EXISTS {target} ( - user_id VARCHAR, - full_name VARCHAR, - created_at TIMESTAMP, - amount DOUBLE, - portal_partition_daily VARCHAR - ) - WITH ( - format = 'PARQUET', - partitioned_by = ARRAY['portal_partition_daily'] - ) - """) - - cursor.execute(f"DELETE FROM {target} WHERE portal_partition_daily = '{ds}'") + catalog = get_catalog_name() + target_fqn = f"{catalog}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" + staging_fqn = f"{catalog}.{OUTPUT_CONFIG.namespace}.stg_tmp_{OUTPUT_CONFIG.table_name}" - cursor.execute(f""" - INSERT INTO {target} + source_query = f""" SELECT CAST(userID AS VARCHAR) AS user_id, TRIM(fullName) AS full_name, CAST(createdAt AS TIMESTAMP) AS created_at, CAST(amount AS DOUBLE) AS amount, '{ds}' AS portal_partition_daily - FROM {staging} - """) - - log.info(f"Transform + load complete: {target}") + FROM {staging_fqn} + """ + overwrite_partition( + target_fqn, + "portal_partition_daily", + ds, + source_query=source_query, + columns=["user_id", "full_name", "created_at", "amount", "portal_partition_daily"], + ) ``` ### Filtering with WHERE ```python - cursor.execute(f""" - INSERT INTO {target} - SELECT - id, - name, - status, - CAST(amount AS DOUBLE) AS amount, - '{ds}' AS portal_partition_daily - FROM {staging} + source_query = f""" + SELECT id, name, status, CAST(amount AS DOUBLE) AS amount, '{ds}' AS portal_partition_daily + FROM {staging_fqn} WHERE status IN ('active', 'pending') AND CAST(amount AS DOUBLE) > 0 - """) + """ ``` ### Deduplication with ROW_NUMBER ```python - cursor.execute(f""" - INSERT INTO {target} + source_query = f""" SELECT id, name, status, updated_at, '{ds}' AS portal_partition_daily FROM ( - SELECT - *, - ROW_NUMBER() OVER ( - PARTITION BY id - ORDER BY updated_at DESC - ) AS rn - FROM {staging} + SELECT *, ROW_NUMBER() OVER (PARTITION BY id ORDER BY updated_at DESC) AS rn + FROM {staging_fqn} ) WHERE rn = 1 - """) + """ ``` ### Aggregation with GROUP BY ```python - cursor.execute(f""" - INSERT INTO {target} + source_query = f""" SELECT category, region, SUM(CAST(amount AS DOUBLE)) AS total_amount, COUNT(*) AS record_count, - AVG(CAST(amount AS DOUBLE)) AS avg_amount, '{ds}' AS portal_partition_daily - FROM {staging} + FROM {staging_fqn} GROUP BY category, region - """) + """ ``` ### Join with existing Trino tables ```python - cursor.execute(f""" - INSERT INTO {target} - SELECT - s.id, - s.name, - r.region_name, - CAST(s.amount AS DOUBLE) AS amount, - '{ds}' AS portal_partition_daily - FROM {staging} s - LEFT JOIN ..region_lookup r - ON s.region_code = r.region_code - """) -``` - -### Window functions + source_fqn = f"{catalog}.opus_python_source.raw_events" + lookup_fqn = f"{catalog}.opus_python_source.region_lookup" -```python - cursor.execute(f""" - INSERT INTO {target} + source_query = f""" SELECT - id, - name, - amount, - SUM(CAST(amount AS DOUBLE)) OVER ( - PARTITION BY category - ORDER BY created_at - ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW - ) AS running_total, - '{ds}' AS portal_partition_daily - FROM {staging} - """) -``` - -### CASE expressions - -```python - cursor.execute(f""" - INSERT INTO {target} - SELECT - id, - name, - CASE - WHEN CAST(amount AS DOUBLE) > 1000 THEN 'high' - WHEN CAST(amount AS DOUBLE) > 100 THEN 'medium' - ELSE 'low' - END AS amount_tier, + s.id, s.name, r.region_name, + CAST(s.amount AS DOUBLE) AS amount, '{ds}' AS portal_partition_daily - FROM {staging} - """) + FROM {source_fqn} s + LEFT JOIN {lookup_fqn} r ON s.region_code = r.region_code + WHERE s.portal_partition_daily = '{ds}' + """ ``` -### Source already in Trino (no staging needed) - -When source data is already in Trino, skip extract and stage entirely — go directly to SQL transform + load: +### Source already in Trino (no staging) ```python def transform_and_load(context: dict) -> None: - """Transform from an existing Trino source table — no staging needed.""" - conn = get_trino_conn() - cursor = conn.cursor() + """Direct SQL transform from an existing Trino table.""" ds = context["ds"] - source = ".some_catalog.raw_events" - target = f"..{OUTPUT_CONFIG.table_name}" - - cursor.execute(f""" - CREATE TABLE IF NOT EXISTS {target} ( - event_id VARCHAR, - event_type VARCHAR, - user_id VARCHAR, - event_count BIGINT, - portal_partition_daily VARCHAR - ) - WITH ( - format = 'PARQUET', - partitioned_by = ARRAY['portal_partition_daily'] - ) - """) + catalog = get_catalog_name() + source_fqn = f"{catalog}.analytics.raw_events" + target_fqn = f"{catalog}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" - cursor.execute(f"DELETE FROM {target} WHERE portal_partition_daily = '{ds}'") - - cursor.execute(f""" - INSERT INTO {target} - SELECT - event_id, - event_type, - user_id, - COUNT(*) AS event_count, - '{ds}' AS portal_partition_daily - FROM {source} + source_query = f""" + SELECT event_id, event_type, user_id, COUNT(*) AS event_count, '{ds}' AS portal_partition_daily + FROM {source_fqn} WHERE portal_partition_daily = '{ds}' GROUP BY event_id, event_type, user_id - """) - - log.info(f"Direct SQL transform + load complete: {target}") + """ + overwrite_partition( + target_fqn, + "portal_partition_daily", + ds, + source_query=source_query, + columns=["event_id", "event_type", "user_id", "event_count", "portal_partition_daily"], + ) ``` --- -## DataFrame Transform Patterns (fallback — use only when SQL cannot express the logic) - -These patterns are for cases where the transformation requires Python-specific capabilities. - -### Nested JSON flattening - -SQL cannot easily handle deeply nested JSON with variable keys: - -```python -def transform(df: pd.DataFrame, context: dict) -> pd.DataFrame: - """Flatten nested JSON columns — SQL cannot handle variable-depth nesting.""" - if df is None or df.empty: - return pd.DataFrame() - - if "metadata" in df.columns: - meta_df = pd.json_normalize(df["metadata"].apply( - lambda x: x if isinstance(x, dict) else {} - )) - meta_df.columns = [f"meta_{c}" for c in meta_df.columns] - df = pd.concat([df.drop(columns=["metadata"]), meta_df], axis=1) - - df["portal_partition_daily"] = context["ds"] - return df -``` - -### Complex string parsing - -When regex or multi-step parsing logic is impractical in SQL: +## Cleanup pattern (SQL path) ```python -def transform(df: pd.DataFrame, context: dict) -> pd.DataFrame: - """Parse structured strings that need multi-step regex — not practical in SQL.""" - import re - if df is None or df.empty: - return pd.DataFrame() - - pattern = r'(?P[^:]+):(?P\d+)/(?P.+)' - parsed = df["endpoint"].str.extract(pattern) - df = pd.concat([df, parsed], axis=1) +from python_models._trino_io import execute_sql - df["portal_partition_daily"] = context["ds"] - return df -``` - -### ML preprocessing - -When the transformation involves scikit-learn or similar: - -```python -def transform(df: pd.DataFrame, context: dict) -> pd.DataFrame: - """ML feature engineering — requires Python libraries.""" - from sklearn.preprocessing import LabelEncoder - if df is None or df.empty: - return pd.DataFrame() - - le = LabelEncoder() - df["category_encoded"] = le.fit_transform(df["category"]) - - df["portal_partition_daily"] = context["ds"] - return df -``` - -After a DataFrame transform, **still load via Trino** by staging the result and using SQL CTAS: - -```python -def load(df: pd.DataFrame, context: dict) -> None: - """Stage transformed DataFrame to Trino, then write final table via SQL.""" - if df is None or df.empty: - log.warning("No data to load") - return - - stage(df, context) - transform_and_load_from_staging(context) - cleanup(context) -``` - ---- - -## Cleanup Pattern - -Always drop staging tables after the final write: - -```python def cleanup(context: dict) -> None: - """Drop temporary staging table.""" - conn = get_trino_conn() - cursor = conn.cursor() - table_name = OUTPUT_CONFIG.table_name or OUTPUT_CONFIG.model_name - staging_table = f"stg_tmp_{table_name}" - cursor.execute( - f"DROP TABLE IF EXISTS ..{staging_table}" - ) - log.info(f"Dropped staging table: {staging_table}") + """Drop temporary staging tables.""" + catalog = get_catalog_name() + staging_fqn = f"{catalog}.{OUTPUT_CONFIG.namespace}.stg_tmp_{OUTPUT_CONFIG.table_name}" + execute_sql(f"DROP TABLE IF EXISTS {staging_fqn}") + log.info(f"Dropped staging table: {staging_fqn}") ``` --- -## Post-Load Patterns +## Post-load patterns -### Validate with Trino SQL +### Validate with fetch_value ```python +from python_models._trino_io import fetch_value + def post_load(context: dict) -> None: - """Verify data was written correctly via Trino query.""" - conn = get_trino_conn() - cursor = conn.cursor() + """Verify data was written correctly.""" ds = context["ds"] - target = f"..{OUTPUT_CONFIG.table_name}" + table_fqn = f"{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" - cursor.execute(f""" - SELECT COUNT(*) AS row_count - FROM {target} + row_count = fetch_value(f""" + SELECT COUNT(*) FROM {table_fqn} WHERE portal_partition_daily = '{ds}' """) - row_count = cursor.fetchone()[0] - log.info(f"Post-load validation: {target} has {row_count} rows for {ds}") + log.info(f"Post-load validation: {table_fqn} has {row_count} rows for {ds}") - if row_count == 0: - raise ValueError(f"No rows loaded for {ds} — check extract and transform logic") + if not row_count: + raise ValueError(f"No rows loaded for {ds}") ``` -### Partition sync (Hive metastore) +### Error handling with cleanup (SQL path) ```python -def post_load(context: dict) -> None: - """Sync Hive metastore partitions after Iceberg write.""" - conn = get_trino_conn() - cursor = conn.cursor() - table = f"{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name or OUTPUT_CONFIG.model_name}" - cursor.execute(f"CALL iceberg.system.sync_partition_metadata('{table}')") - log.info(f"Partition sync complete for {table}") +def run_etl(context: dict): + try: + extract(context) + transform_and_load(context) + finally: + cleanup(context) ``` --- -## Optimization Patterns +## Optimization patterns ### Partition pruning -Always include partition column predicates in WHERE clauses to avoid full table scans: - ```sql --- GOOD: Trino prunes partitions, reads only one day -SELECT * FROM target WHERE portal_partition_daily = '2025-01-15' +-- Reads one partition +SELECT id, name FROM target WHERE portal_partition_daily = '2025-01-15' --- BAD: full table scan — reads every partition -SELECT * FROM target WHERE name = 'foo' +-- Full table scan — avoid on large tables +SELECT id, name FROM target WHERE name = 'foo' ``` ### Column projection -Enumerate columns explicitly — never use `SELECT *` in production: - -```sql --- GOOD: reads only needed columns from Parquet/Iceberg -SELECT id, name, amount, portal_partition_daily FROM target - --- BAD: reads all columns, wastes I/O and memory -SELECT * FROM target -``` - -### EXPLAIN ANALYZE - -Use `EXPLAIN ANALYZE` to verify query plans before deploying: - -```python -def debug_query(context: dict) -> None: - """Check query plan for performance issues.""" - conn = get_trino_conn() - cursor = conn.cursor() - - cursor.execute(f""" - EXPLAIN ANALYZE - SELECT id, name, amount - FROM ..stg_tmp_my_model - WHERE status = 'active' - """) - - for row in cursor.fetchall(): - log.info(row[0]) -``` +Enumerate columns in SELECT and INSERT — avoid `SELECT *` in production SQL. ### Write mode decision matrix -| Scenario | Write mode | SQL pattern | -|----------|-----------|-------------| -| Daily idempotent loads | `overwrite_partitions` | `DELETE WHERE partition = ds` then `INSERT INTO ... SELECT` | -| Event streams / append-only | `append` | `INSERT INTO ... SELECT` (no DELETE) | -| Full table refresh | `overwrite` | `DROP TABLE` then `CREATE TABLE AS SELECT` | -| Slowly changing dimensions | `merge` (if supported) | `MERGE INTO ... USING ... ON ...` | - -### Batch size for staging +| Scenario | Write mode | Pattern | +| -------- | ---------- | ------- | +| Daily idempotent loads | `overwrite_partitions` | `overwrite_partition` / `overwrite_partition_from_dataframe` | +| Event streams | `append` | `append(table_fqn, insert_sql)` | +| Full table refresh | `overwrite` | `overwrite(table_fqn, "true", insert_sql=…)` | +| Upsert by key | merge | `merge(table_fqn, merge_sql)` | -For large DataFrames, tune the batch size based on column count: +### Batch size guidance (pandas writes) -| Columns | Recommended batch size | -|---------|----------------------| -| < 10 | 1000 - 5000 rows | -| 10 - 50 | 500 - 1000 rows | -| > 50 | 100 - 500 rows | +| Columns | Recommended `batch_size` | +| ------- | ------------------------ | +| < 10 | 1000–5000 | +| 10–50 | 500–1000 | +| > 50 | 100–500 | -### Avoid common performance pitfalls +### Common pitfalls | Pitfall | Fix | -|---------|-----| -| `SELECT *` in transforms | Enumerate needed columns explicitly | -| No partition predicate in WHERE | Add `WHERE portal_partition_daily = '{ds}'` | -| Pulling Trino data into pandas for transforms | Use SQL transforms directly in Trino | -| Not dropping staging tables | Always call `cleanup()` after `transform_and_load()` | -| Single INSERT for entire DataFrame | Batch into chunks (e.g., 1000 rows) | -| Creating tables without `partitioned_by` | Always specify `partitioned_by = ARRAY['portal_partition_daily']` | +| ------- | --- | +| Inline Trino connections | Import from `python_models._trino_io` | +| `execute_trino` (old name) | Use `execute_sql` / `fetch_value` | +| Manual `DELETE` before INSERT | Use `overwrite_partition` or `overwrite_partition_from_dataframe` | +| `SELECT *` in production SQL | Enumerate columns | +| Missing partition predicate | Add `WHERE portal_partition_daily = '{ds}'` on large sources | +| Orphaned `stg_tmp_*` tables | Call `cleanup()` in `finally` block | diff --git a/templates/skills/dj-create-python-model/references/worked-example.md b/templates/skills/dj-create-python-model/references/worked-example.md index 9e57864..f449643 100644 --- a/templates/skills/dj-create-python-model/references/worked-example.md +++ b/templates/skills/dj-create-python-model/references/worked-example.md @@ -1,283 +1,257 @@ -# Worked Example: SQL-First Python Model +# Worked Example: Topic Pipeline (Pandas + Trino Upstream) -A complete `.python.json` for an API-fetch model using the SQL-first pattern. +A three-model topic under `dags/python_models/etl/tenant_status/` — create each model via **Create Python Model** (one at a time) or batch-write `.python.json` files with this skill. -**Flow:** Extract from Backstage API -> stage raw data into a temporary Trino table -> transform and load via SQL (`DELETE + INSERT INTO ... SELECT`) -> drop staging table. +| Model | Style | Role | +| ----- | ----- | ---- | +| `tenant_events` | Pandas | Ingest API → Iceberg | +| `tenant_migrations` | Pandas | Ingest S3 CSV → Iceberg | +| `tenant_migration_status` | Pandas | Reads both upstream tables via `read_partition_to_dataframe`, merges in pandas, writes via `overwrite_partition_from_dataframe` | -**Target:** `..backstage_catalogs` +DJ auto-infers trino `upstream_sources` on model 3 from `depends_on` when you save or sync `.python.json`. Authors may also declare them explicitly. + +--- + +## Topic-level `.python.json` configs + +### Model 1: `tenant_events` (API → Iceberg) ```jsonc { - "name": "backstage_catalogs", + "name": "tenant_events", "group": "etl", - "topic": "backstage", - "description": "Fetches catalog entities from Backstage API daily, stages to Trino, transforms and loads via SQL", + "topic": "tenant_status", + "description": "Daily tenant events from Backstage API", "model_type": "python", - "dags": ["source_etl"], - "output_type": "iceberg", + "dags": ["opus_python_source"], + "compute": "kpo", + "kpo_size": "small", + "depends_on": [], + "upstream_sources": [ + { "type": "external", "value": "api::backstage_tenant_events" } + ], "output": { - "database": "", - "schema": "", - "table": "backstage_catalogs", + "database": "glue_development", + "schema": "opus_python_source", + "table": "tenant_events", "partition_by": ["portal_partition_daily"], "write_mode": "overwrite_partitions" }, - "namespace": "python", - "table_name": "backstage_catalogs", - "enable_notebook": true, - "dependencies": ["requests>=2.28.0", "trino>=0.327.0"], - "tags": ["python-model", "api"], - "owner": "platform-team", - "variables": { - "api_url": "https://backstage.example.com/api/catalog/entities", - "api_token": "", - "trino_host": "localhost", - "trino_port": "8080", - "trino_user": "etl" + "dependencies": ["requests>=2.28.0", "pandas>=2.0.0"], + "tags": ["python-model", "api"] +} +``` + +Key cells (abbreviated): + +```python +from python_models._trino_io import overwrite_partition_from_dataframe +from python_models._config import PythonModelConfig, get_catalog_name, build_context_from_env + +OUTPUT_CONFIG = PythonModelConfig( + model_name="tenant_events", + namespace="opus_python_source", + table_name="tenant_events", + model_type="python", + description="Daily tenant events from Backstage API", +) + +def extract(context: dict) -> pd.DataFrame: + import requests + resp = requests.get(INPUT_VARIABLES["api_url"], headers={"Authorization": f"Bearer {INPUT_VARIABLES['api_token']}"}) + resp.raise_for_status() + return pd.DataFrame(resp.json()) + +def transform(df: pd.DataFrame, context: dict) -> pd.DataFrame: + df = df.rename(columns={"tenantId": "tenant_id", "eventType": "event_type"}) + df["portal_partition_daily"] = context["ds"] + return df + +def load(df: pd.DataFrame, context: dict) -> None: + table_fqn = f"{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" + columns = ["tenant_id", "event_type", "portal_partition_daily"] + overwrite_partition_from_dataframe( + table_fqn, "portal_partition_daily", context["ds"], df, columns + ) + +def run_etl(context: dict): + df = extract(context) + df = transform(df, context) + load(df, context) + +if __name__ == "__main__": + run_etl(build_context_from_env()) +``` + +### Model 2: `tenant_migrations` (S3 CSV → Iceberg) + +```jsonc +{ + "name": "tenant_migrations", + "group": "etl", + "topic": "tenant_status", + "description": "Daily tenant migration records from S3 CSV", + "depends_on": [], + "upstream_sources": [ + { "type": "external", "value": "s3::data-lake/tenant_migrations/{ds}/migrations.csv" } + ], + "output": { + "database": "glue_development", + "schema": "opus_python_source", + "table": "tenant_migrations", + "partition_by": ["portal_partition_daily"], + "write_mode": "overwrite_partitions" }, - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": ["# Python Model: python__etl__backstage__backstage_catalogs\n", "**DAGs**: source_etl\n", "**Strategy**: SQL-first — extract from API, stage to Trino, transform + load via SQL\n"] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from datetime import datetime\n", - "import io\n", - "import json\n", - "import logging\n", - "\n", - "import pandas as pd\n", - "import requests\n", - "\n", - "log = logging.getLogger(__name__)\n", - "\n", - "context = {\n", - " 'ds': datetime.now().strftime('%Y-%m-%d'),\n", - " 'ds_nodash': datetime.now().strftime('%Y%m%d'),\n", - "}\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from python_models._config import PythonModelConfig\n", - "\n", - "OUTPUT_CONFIG = PythonModelConfig(\n", - " model_name=\"backstage_catalogs\",\n", - " namespace=\"python\",\n", - " table_name=\"backstage_catalogs\",\n", - " model_type=\"python\",\n", - " description=\"Fetches catalog entities from Backstage API daily\",\n", - ")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# ============================================================\n", - "# INPUT VARIABLES\n", - "# ============================================================\n", - "INPUT_VARIABLES = {\n", - " \"name\": \"backstage_catalogs\",\n", - " \"group\": \"etl\",\n", - " \"topic\": \"backstage\",\n", - " \"description\": \"Fetches catalog entities from Backstage API daily\",\n", - " \"api_url\": \"https://backstage.example.com/api/catalog/entities\",\n", - " \"api_token\": \"\",\n", - " \"trino_host\": \"localhost\",\n", - " \"trino_port\": \"8080\",\n", - " \"trino_user\": \"etl\",\n", - "}\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# ============================================================\n", - "# TRINO CONNECTION\n", - "# ============================================================\n", - "def get_trino_conn():\n", - " \"\"\"Reusable Trino connection.\"\"\"\n", - " from trino.dbapi import connect\n", - " return connect(\n", - " host=INPUT_VARIABLES.get(\"trino_host\", \"localhost\"),\n", - " port=int(INPUT_VARIABLES.get(\"trino_port\", \"8080\")),\n", - " user=INPUT_VARIABLES.get(\"trino_user\", \"etl\"),\n", - " catalog=\"\",\n", - " schema=\"\",\n", - " )\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# ============================================================\n", - "# EXTRACT\n", - "# ============================================================\n", - "def extract(context: dict) -> pd.DataFrame:\n", - " \"\"\"Fetch catalog entities from Backstage API.\"\"\"\n", - " log.info(f\"Extracting for {context['ds']}...\")\n", - " url = INPUT_VARIABLES[\"api_url\"]\n", - " headers = {\"Authorization\": f\"Bearer {INPUT_VARIABLES['api_token']}\"}\n", - "\n", - " resp = requests.get(url, headers=headers)\n", - " resp.raise_for_status()\n", - " return pd.DataFrame(resp.json())\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# ============================================================\n", - "# STAGE (raw data -> temporary Trino table)\n", - "# ============================================================\n", - "def stage(df: pd.DataFrame, context: dict) -> None:\n", - " \"\"\"Stage raw data into a temporary Trino table.\"\"\"\n", - " if df is None or df.empty:\n", - " log.warning(\"No data to stage\")\n", - " return\n", - "\n", - " conn = get_trino_conn()\n", - " cursor = conn.cursor()\n", - " staging_table = \"stg_tmp_backstage_catalogs\"\n", - "\n", - " cursor.execute(f\"DROP TABLE IF EXISTS ..{staging_table}\")\n", - "\n", - " columns = df.columns.tolist()\n", - " col_defs = \", \".join(f\"{c} VARCHAR\" for c in columns)\n", - " cursor.execute(\n", - " f\"CREATE TABLE ..{staging_table} ({col_defs})\"\n", - " )\n", - "\n", - " batch_size = 1000\n", - " for i in range(0, len(df), batch_size):\n", - " batch = df.iloc[i:i + batch_size]\n", - " values_list = []\n", - " for _, row in batch.iterrows():\n", - " vals = \", \".join(f\"'{str(v).replace(chr(39), chr(39)+chr(39))}'\"\n", - " for v in row)\n", - " values_list.append(f\"({vals})\")\n", - " cursor.execute(\n", - " f\"INSERT INTO ..{staging_table} \"\n", - " f\"VALUES {', '.join(values_list)}\"\n", - " )\n", - "\n", - " log.info(f\"Staged {len(df)} rows into {staging_table}\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# ============================================================\n", - "# TRANSFORM + LOAD (all transformations in Trino SQL)\n", - "# ============================================================\n", - "def transform_and_load(context: dict) -> None:\n", - " \"\"\"Transform and load via Trino SQL.\"\"\"\n", - " conn = get_trino_conn()\n", - " cursor = conn.cursor()\n", - " ds = context[\"ds\"]\n", - " target = \"..backstage_catalogs\"\n", - " staging = \"..stg_tmp_backstage_catalogs\"\n", - "\n", - " cursor.execute(f\"\"\"\n", - " CREATE TABLE IF NOT EXISTS {target} (\n", - " entity_ref VARCHAR,\n", - " kind VARCHAR,\n", - " name VARCHAR,\n", - " namespace VARCHAR,\n", - " portal_partition_daily VARCHAR\n", - " )\n", - " WITH (\n", - " format = 'PARQUET',\n", - " partitioned_by = ARRAY['portal_partition_daily']\n", - " )\n", - " \"\"\")\n", - "\n", - " cursor.execute(f\"\"\"\n", - " DELETE FROM {target}\n", - " WHERE portal_partition_daily = '{ds}'\n", - " \"\"\")\n", - "\n", - " cursor.execute(f\"\"\"\n", - " INSERT INTO {target}\n", - " SELECT\n", - " CAST(entity_ref AS VARCHAR) AS entity_ref,\n", - " TRIM(LOWER(kind)) AS kind,\n", - " TRIM(LOWER(name)) AS name,\n", - " COALESCE(namespace, 'default') AS namespace,\n", - " '{ds}' AS portal_partition_daily\n", - " FROM {staging}\n", - " \"\"\")\n", - "\n", - " log.info(f\"Transform + load complete: {target}\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# ============================================================\n", - "# CLEANUP\n", - "# ============================================================\n", - "def cleanup(context: dict) -> None:\n", - " \"\"\"Drop temporary staging table.\"\"\"\n", - " conn = get_trino_conn()\n", - " cursor = conn.cursor()\n", - " staging_table = \"stg_tmp_backstage_catalogs\"\n", - " cursor.execute(f\"DROP TABLE IF EXISTS ..{staging_table}\")\n", - " log.info(f\"Dropped staging table: {staging_table}\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# ============================================================\n", - "# MAIN\n", - "# ============================================================\n", - "def run_etl(context: dict):\n", - " \"\"\"Main ETL orchestrator — SQL-first flow.\"\"\"\n", - " df = extract(context)\n", - " stage(df, context)\n", - " transform_and_load(context)\n", - " cleanup(context)\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": ["# Run the ETL\n", "run_etl(context)\n"] + "dependencies": ["boto3>=1.28.0", "pandas>=2.0.0"] +} +``` + +Pandas extract reads S3; `load()` uses `overwrite_partition_from_dataframe` the same way as model 1. + +### Model 3: `tenant_migration_status` (integration — reads inferred trino upstreams) + +```jsonc +{ + "name": "tenant_migration_status", + "group": "etl", + "topic": "tenant_status", + "description": "Combines tenant events and migrations into a daily status table", + "depends_on": ["tenant_events", "tenant_migrations"], + "upstream_sources": [ + { "type": "trino", "value": "glue_development.opus_python_source.tenant_events" }, + { "type": "trino", "value": "glue_development.opus_python_source.tenant_migrations" } + ], + "output": { + "database": "glue_development", + "schema": "opus_python_source", + "table": "tenant_migration_status", + "partition_by": ["portal_partition_daily"], + "write_mode": "overwrite_partitions" + } +} +``` + +If `upstream_sources` is omitted, DJ infers the two trino entries from `depends_on` at create time. + +Key cells: + +```python +from python_models._trino_io import ( + read_partition_to_dataframe, + overwrite_partition_from_dataframe, +) +from python_models._config import PythonModelConfig, get_catalog_name, build_context_from_env + +OUTPUT_CONFIG = PythonModelConfig( + model_name="tenant_migration_status", + namespace="opus_python_source", + table_name="tenant_migration_status", + model_type="python", + description="Combines tenant events and migrations into a daily status table", +) + +EVENTS_FQN = f"{get_catalog_name()}.opus_python_source.tenant_events" +MIGRATIONS_FQN = f"{get_catalog_name()}.opus_python_source.tenant_migrations" + +def extract(context: dict) -> dict[str, pd.DataFrame]: + ds = context["ds"] + return { + "events": read_partition_to_dataframe( + EVENTS_FQN, "portal_partition_daily", ds, + columns=["tenant_id", "event_type"], + ), + "migrations": read_partition_to_dataframe( + MIGRATIONS_FQN, "portal_partition_daily", ds, + columns=["tenant_id", "migration_status"], + ), } - ] + +def transform(dfs: dict[str, pd.DataFrame], context: dict) -> pd.DataFrame: + merged = dfs["events"].merge(dfs["migrations"], on="tenant_id", how="outer") + merged["has_event"] = merged["event_type"].notna() + merged["portal_partition_daily"] = context["ds"] + return merged + +def load(df: pd.DataFrame, context: dict) -> None: + table_fqn = f"{get_catalog_name()}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" + columns = ["tenant_id", "event_type", "migration_status", "has_event", "portal_partition_daily"] + overwrite_partition_from_dataframe( + table_fqn, "portal_partition_daily", context["ds"], df, columns + ) + +def run_etl(context: dict): + dfs = extract(context) + df = transform(dfs, context) + load(df, context) + +if __name__ == "__main__": + run_etl(build_context_from_env()) +``` + +--- + +## SQL-path alternative (model 1 only) + +The same API ingest can use the SQL orchestrator instead — stage with `stage()`, then `overwrite_partition` with a SQL `source_query`. Both paths are valid; pick one per model. + +```python +from python_models._trino_io import execute_sql, overwrite_partition + +def transform_and_load(context: dict) -> None: + ds = context["ds"] + catalog = get_catalog_name() + target_fqn = f"{catalog}.{OUTPUT_CONFIG.namespace}.{OUTPUT_CONFIG.table_name}" + staging_fqn = f"{catalog}.{OUTPUT_CONFIG.namespace}.stg_tmp_{OUTPUT_CONFIG.table_name}" + + source_query = f""" + SELECT + CAST(tenantId AS VARCHAR) AS tenant_id, + TRIM(eventType) AS event_type, + '{ds}' AS portal_partition_daily + FROM {staging_fqn} + """ + overwrite_partition( + target_fqn, + "portal_partition_daily", + ds, + source_query=source_query, + columns=["tenant_id", "event_type", "portal_partition_daily"], + ) + +def cleanup(context: dict) -> None: + catalog = get_catalog_name() + staging_fqn = f"{catalog}.{OUTPUT_CONFIG.namespace}.stg_tmp_{OUTPUT_CONFIG.table_name}" + execute_sql(f"DROP TABLE IF EXISTS {staging_fqn}") +``` + +--- + +## Lineage result + +After successful KPO runs, each output table carries: + +```json +{ + "python_model_name": "tenant_migration_status", + "python_model_upstream_sources": "[{\"type\":\"trino\",\"value\":\"glue_development.opus_python_source.tenant_events\"},{\"type\":\"trino\",\"value\":\"glue_development.opus_python_source.tenant_migrations\"}]" } ``` + +Model Lineage renders: + +- External leaves for `api::…` and `s3::…` on models 1–2 +- Trino-ref nodes for upstream tables on model 3 +- Python model nodes for each output table +- Recursive walk from downstream dbt sources through the full chain + +--- + +## Downstream dbt source + +Register each output table, then reference in `.model.json`: + +``` +glue_development__opus_python_source.tenant_migration_status +``` + +Use **`dj-create-source`** after the first successful run. diff --git a/templates/skills/dj-review-python-model/_SKILL.md b/templates/skills/dj-review-python-model/_SKILL.md index e83f13d..ba09551 100644 --- a/templates/skills/dj-review-python-model/_SKILL.md +++ b/templates/skills/dj-review-python-model/_SKILL.md @@ -15,7 +15,7 @@ metadata: **Goal:** audit a `.python.json` + `.python.py` pair against DJ framework standards and produce a **structured report** covering framework compliance, lineage readiness, downstream integration, and performance. This skill is **read-only** — it surfaces issues and recommendations but does not auto-fix. -**SQL-first principle:** Python models are pre-dbt ETL pipelines. Python handles orchestration and external data ingestion; Trino SQL handles transformations and storage. The review validates this separation is maintained. +**Authoring philosophy:** Python models may use **pandas end-to-end**, **SQL end-to-end**, or **mixed** orchestration. The review validates consistency with the chosen style — not SQL-first compliance. ## When this skill applies @@ -29,22 +29,12 @@ Use this skill when the user mentions: review python model, audit python model, - Verifying the model's output *data* against a legacy/reference table (this skill audits code, not data) → `dj-verify-pymodel-parity` - Migrating a legacy notebook into a python model before it can be reviewed → `dj-migrate-notebook-to-pymodel` -## Step 0 — Resolve project defaults (before review) - -1. Read `.agents/project/skills/dj-review-python-model/project-defaults.md` if it exists. -2. Read `.agents/project/skills/dj-create-python-model/project-defaults.md` for shared catalog, schema, and **authoring_mode** (`hand-written` | `cells`). -3. Apply authoring_mode to F3, F7, and F9 (see Framework Compliance). - -## DJ CLI (preferred when DJ is running) - -When `.dj/bin/dj system.ping` succeeds, use `dbt.sources` and `model.lineage` to cross-check downstream registration and lineage for the model’s output table. This skill is read-only — do not mutate python files via the CLI. Invocation patterns → `dj-cli`. - ## Workflow -- [ ] **1. Resolve scope.** Default to the open `.python.json` in the editor if any. Otherwise ask the user which Python model to review. Accept a file path or model name (`python______`). +- [ ] **1. Resolve scope.** Default to the open `.python.json` in the editor if any. Otherwise ask which Python model to review. Accept a file path or model name (`python______`). - [ ] **2. Read both files.** Load the `.python.json` and its companion `.python.py`. If either is missing, note it as an immediate finding. Also read `.dj/schemas/python-model.schema.json` for validation reference. - [ ] **3. Run checklist.** Apply all checks from the four categories below. Record each as pass/fail/warning with supporting detail. -- [ ] **4. Cross-reference lineage.** Parse SQL statements in the code cells to extract all tables referenced in FROM/JOIN clauses. Compare against `python_model_upstream_sources` entries for completeness. +- [ ] **4. Cross-reference lineage.** Parse SQL and `_trino_io` calls to extract Trino tables read. Compare against `upstream_sources` in JSON. Verify `depends_on` aligns with Python-model data dependencies. - [ ] **5. Render report.** Output the structured report using the template below. Categorize findings by severity (Issue / Warning / Suggestion). - [ ] **6. Recommendations.** Close with a priority-ordered list of actionable improvements. @@ -52,71 +42,77 @@ When `.dj/bin/dj system.ping` succeeds, use `dbt.sources` and `model.lineage` to ### 1. Framework Compliance (F) -| Check | What to validate | -| ----- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| F1 | `.python.json` has required fields: `name`, `group`, `topic` | -| F2 | Name/group/topic match pattern `^[a-z][a-z0-9_]*$` | -| F3 | **`cells` mode:** `cells` array is present and non-empty in JSON. **`hand-written` mode:** empty or absent `cells` is OK | -| F4 | `.python.py` contains `if __name__ == "__main__":` script entry (scaffold calls `run_etl(build_context_from_env())`) | -| F5 | Uses `_trino_io` helpers (`from python_models._trino_io import ...`) — no inline Trino connection code (`trino.dbapi.connect`, `create_engine`, raw `requests.post` to Trino) | -| F6 | `OUTPUT_CONFIG` uses `PythonModelConfig` from `python_models._config` | -| F7 | **`cells` mode:** `.python.py` derivable from JSON `cells` (no hand-edits lost on sync). **`hand-written` mode:** `.python.py` is source of truth — hand edits expected | -| F8 | Runner cell (`run_etl(context)`) is the last code cell in JSON | -| F9 | **`cells` mode:** standard structure `extract()`, `transform_and_load()`, `cleanup()`, `run_etl()` when siblings use it. **`hand-written` mode:** only `run_etl(context)` is required unless topic convention says otherwise | +| Check | What to validate | +| ----- | ---------------- | +| F1 | `.python.json` has required fields: `name`, `group`, `topic` | +| F2 | Name/group/topic match pattern `^[a-z][a-z0-9_]*$` | +| F3 | `.python.json` is valid with required metadata; `cells` is an optional non-authoritative snapshot (may be empty when the model is authored via `.python.py` directly) | +| F4 | `.python.py` contains `if __name__ == "__main__":` calling `run_etl(build_context_from_env())` | +| F5 | Uses `_trino_io` helpers (`from python_models._trino_io import …`) — no inline Trino connection code (`trino.dbapi.connect`, `create_engine`, raw HTTP to Trino) | +| F6 | `OUTPUT_CONFIG` uses `PythonModelConfig` from `python_models._config` | +| F7 | `.python.py` is present and treated as the code source of truth — hand-edited code is expected and preserved (DJ never overwrites it from JSON) | +| F8 | `.python.py` is import-safe — no side-effectful top-level statement such as a bare `run_etl(context)`; execution runs through the `__main__` block | +| F9 | ETL follows a coherent orchestrator: pandas (`extract` / `transform` / `load`) **or** SQL (`extract` / `transform_and_load` / `cleanup`) **or** documented mixed pattern | -### 2. Lineage Readiness (L) — end-to-end validation +### 2. Lineage Readiness (L) -The DJ lineage engine discovers Python models by querying Iceberg `$properties` on output tables. It reads: +The DJ lineage engine discovers Python models via Iceberg `$properties` on output tables: - `python_model_name` — model identity - `python_model_table` — output table name -- `python_model_upstream_sources` — JSON array of `{ "type": "trino"|"external", "value": "..." }` entries for upstream lineage edges - -The review validates the **full chain**: JSON metadata → post-run Iceberg table properties → lineage discoverability (including recursive walks to root). - -| Check | What to validate | -| ----- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| L1 | `.python.json` includes `name` (written as `python_model_name` on the output table after each run) plus `description`, `namespace` / `output.schema`, and `table_name` / `output.table` as applicable | -| L2 | `upstream_sources` is set in `.python.json` (DJ writes `python_model_upstream_sources` to the output table after each run — no manual `ALTER TABLE` in `.python.py`) | -| L3 | Each `trino` entry's `value` follows `schema.table` format; `external` entries use the intended free-form identifier | -| L4 | **Completeness** — every table referenced in `extract()` / `transform_and_load()` SQL (FROM/JOIN clauses) appears in `upstream_sources` (no missing upstream edges) | -| L5 | **Accuracy** — every entry in `upstream_sources` corresponds to a source actually read by the model (no stale/phantom entries) | -| L6 | `dags` field is populated in JSON (model is scheduled, discoverable by Airflow DAG lineage). Empty `dags` = utility module, not a lineage participant | -| L7 | `depends_on` correctly lists upstream Python models whose output tables this model reads (task dependency mirrors data dependency) | -| L8 | Model ID derivable from file path matches `python_model_name` in properties: `python______` | -| L9 | Companion `.python.py` exists alongside `.python.json` (required for Airflow `etl_helper.py` discovery) | -| L10 | `OUTPUT_CONFIG.table_name` (or `model_name` fallback) matches the table name referenced in downstream `.source.json` files, if any exist in the project | +- `python_model_upstream_sources` — JSON array of `{ "type": "trino"|"external", "value": "..." }` + +**`depends_on` vs `upstream_sources` — validate both when reading model output:** + +| Field | Role | Review focus | +| ----- | ---- | ------------ | +| `depends_on` | Airflow task ordering | Upstream Python models whose output this model reads must be listed | +| `upstream_sources` | Lineage graph edges | Every Trino table and external system read must be declared | + +DJ auto-infers trino `upstream_sources` from `depends_on` at create/sync — inferred entries are valid. Flag missing `depends_on` when the model reads a sibling Python output; flag missing `upstream_sources` when the model reads any upstream table or external system. + +| Check | What to validate | +| ----- | ---------------- | +| L1 | JSON includes `name`, `description`, `namespace` / `output.schema`, `table_name` / `output.table` | +| L2 | `upstream_sources` is set when the model reads upstream data (Trino tables or external systems) | +| L3 | Each `trino` entry uses **`catalog.schema.table`** (canonical). **`schema.table`** (2-part) → legacy **warning** | +| L4 | Each `external` entry uses a stable identifier (`s3::…`, `api::…`, `jenkins::…`, or project convention) | +| L5 | **Completeness** — every Trino table read (SQL FROM/JOIN, `read_partition_to_dataframe`, etc.) appears in `upstream_sources` | +| L6 | **Accuracy** — every `upstream_sources` entry corresponds to a source actually read (no stale entries) | +| L7 | `depends_on` lists upstream Python models when this model reads their output tables | +| L8 | `dags` populated (empty = utility module, not a scheduled lineage participant) | +| L9 | Model ID from file path matches `python_model_name`: `python______` | +| L10 | Companion `.python.py` exists alongside `.python.json` | +| L11 | `OUTPUT_CONFIG.table_name` matches downstream `.source.json` references, if any | ### 3. Downstream Integration (D) -| Check | What to validate | -| ----- | ----------------------------------------------------------------------------------------------------------------------------- | -| D1 | Output table matches Step 0 convention (`..
` or project-configured equivalent) | -| D2 | Partition column `portal_partition_daily` is emitted in output SQL (`'{ds}' AS portal_partition_daily` or equivalent) | -| D3 | Column names in output SQL follow `snake_case` convention (no camelCase, no spaces, no special characters) | -| D4 | No `SELECT *` in production INSERT — explicit column enumeration for schema stability and consumer predictability | -| D5 | Table name matches model name for source discovery (`OUTPUT_CONFIG.table_name or model_name` → downstream `source.table`) | -| D6 | Output columns are stable — column order and types should be deterministic across runs | +| Check | What to validate | +| ----- | ---------------- | +| D1 | Output table uses standard catalog/schema convention (`glue_development.opus_python_source` or project equivalent) | +| D2 | Partition column `portal_partition_daily` emitted in output (SQL or pandas) when using partition writes | +| D3 | Column names follow `snake_case` | +| D4 | No `SELECT *` in production SQL INSERT statements | +| D5 | Table name matches model name for source discovery | +| D6 | Output columns deterministic across reruns for the same `ds` | ### 4. Performance & Best Practices (P) -| Check | What to validate | -| ----- | ------------------------------------------------------------------------------------------------------------------------------------------------ | -| P1 | **SQL-first adherence** — transformations (filter, cast, aggregate, join, deduplicate) done in Trino SQL, not pandas/Python | -| P2 | Staging tables cleaned up in `cleanup()` — no orphaned `stg_tmp_*` tables left after ETL completes | -| P3 | Partition predicate pushdown — WHERE clauses on source tables include partition column filters | -| P4 | Batch staging — large datasets chunked before INSERT (not unbounded single INSERT of millions of rows) | -| P5 | Write mode appropriate for use case: `overwrite_partition` for idempotent daily loads, `append` for event streams, `overwrite` for full refresh | -| P6 | No full-table scans — queries against large tables include a partition filter or LIMIT | -| P7 | DataFrame operations restricted to external ingestion — any pandas/polars used only for API response parsing, not for SQL-expressible transforms | -| P8 | Explicit `CAST(... AS type)` in SQL for type safety — not relying on pandas `.astype()` or implicit coercion | -| P9 | Error handling — `run_etl()` has try/except with cleanup on failure (staging tables dropped even if transform fails) | -| P10 | Idempotency — re-running the model for the same `ds` produces identical results (no append-on-rerun for partition-based models) | +| Check | What to validate | +| ----- | ---------------- | +| P1 | **Orchestrator consistency** — pandas, SQL, or mixed path is coherent; flag only contradictions (e.g. SQL-expressible transform duplicated in pandas when author chose SQL-only) | +| P2 | Staging tables cleaned up in `cleanup()` (SQL path) or no orphaned `stg_tmp_*` tables | +| P3 | Partition predicate pushdown on large Trino source reads | +| P4 | Batch staging / writes for large datasets (`overwrite_partition_from_dataframe` batch_size, chunked API pagination) | +| P5 | Write mode appropriate: `overwrite_partition` / `overwrite_partition_from_dataframe` for idempotent daily loads | +| P6 | No full-table scans on partitioned sources without partition filters | +| P7 | `_trino_io` API names correct: `execute_sql`, `fetch_value` — not `execute_trino` | +| P8 | Explicit `CAST(… AS type)` in SQL path; appropriate pandas dtypes in pandas path | +| P9 | Error handling — `cleanup()` in `finally` when staging tables are used | +| P10 | Idempotency — rerun for same `ds` produces identical partition output | ## Report template -Render the full report in this structure. Adapt headings to the actual model. Omit empty sections. - ```text ## Python Model Review: python______ @@ -125,11 +121,12 @@ Render the full report in this structure. Adapt headings to the actual model. Om |-------|-------| | Model ID | python______ | | File | | +| Orchestrator | pandas / SQL / mixed | | DAG(s) | | | Output | ..
| | Write mode | | | Partition | | -| Depends on | | +| Depends on | | ### Findings: X issues, Y warnings, Z suggestions @@ -140,10 +137,12 @@ Render the full report in this structure. Adapt headings to the actual model. Om ### 2. Lineage Readiness [L1] PASS/ISSUE/WARNING: Detail: - Tables in SQL: - Tables in upstream_sources: - Missing from upstream_sources: - Stale in upstream_sources: + depends_on: + upstream_sources declared: + Tables read in code: + Missing from upstream_sources: + Missing from depends_on: + Stale in upstream_sources: ### 3. Downstream Integration [D1] PASS/ISSUE/WARNING: @@ -152,8 +151,6 @@ Render the full report in this structure. Adapt headings to the actual model. Om ### 4. Performance & Best Practices [P1] PASS/ISSUE/WARNING: Detail: - Impact: - Suggestion: ### Recommendations (priority-ordered) 1. [CRITICAL] — fixes @@ -164,52 +161,56 @@ Render the full report in this structure. Adapt headings to the actual model. Om ### Severity classification -| Severity | Criteria | Report label | -| -------------- | -------------------------------------------------------------------------- | ------------ | -| **Issue** | Blocks production deployment or breaks lineage/downstream consumers | `ISSUE` | -| **Warning** | Degrades performance, violates best practices, or creates maintenance risk | `WARNING` | -| **Suggestion** | Improvement opportunity; not blocking but raises quality | `SUGGESTION` | +| Severity | Criteria | Report label | +| -------- | -------- | ------------ | +| **Issue** | Blocks production or breaks lineage/downstream consumers | `ISSUE` | +| **Warning** | Degrades performance, violates best practices, or creates maintenance risk | `WARNING` | +| **Suggestion** | Quality improvement, not blocking | `SUGGESTION` | ### Recommendation priority -| Priority | Criteria | -| -------- | ----------------------------------------------------------------------- | -| CRITICAL | Lineage broken (L2-L5 failures), missing `__main__` entry, no partition column | -| HIGH | Inline Trino code, missing cleanup, stale upstream_sources | -| MEDIUM | Non-SQL-first transforms, missing batch staging, no error handling | -| LOW | Naming style, column order, documentation gaps | +| Priority | Criteria | +| -------- | -------- | +| CRITICAL | Lineage broken (L2/L5/L7 failures), missing `__main__`, no partition column | +| HIGH | Inline Trino code, missing cleanup, stale `upstream_sources`, missing `depends_on` for Python upstream | +| MEDIUM | Orchestrator inconsistency, missing batch staging, no error handling | +| LOW | Legacy 2-part trino FQN, naming style, documentation gaps | ## Hard rules (DO NOT) -- **DO NOT** edit any files. This is a read-only review skill. -- **DO NOT** execute SQL queries against Trino to validate table existence. Work from static code analysis only. -- **DO NOT** skip the lineage cross-reference (L4/L5). Always parse SQL to extract referenced tables and compare against declared upstream sources. -- **DO NOT** report a check as PASS without evidence. Show the relevant code snippet or field value that confirms compliance. -- **DO NOT** invent findings. If a check is not applicable (e.g., model has no SQL transforms because it's pure Python-to-S3), mark it as `N/A` with a reason. -- **DO NOT** render empty sections. If all checks in a category pass, still show them as PASS but omit the Recommendations section if there are zero issues/warnings. - -## Lineage cross-reference procedure (L4/L5 detail) +- **DO NOT** edit any files. Read-only review. +- **DO NOT** execute SQL against Trino. Static code analysis only. +- **DO NOT** skip lineage cross-reference (L5/L6/L7). Parse code for table reads and compare to `upstream_sources` and `depends_on`. +- **DO NOT** report PASS without evidence (code snippet or field value). +- **DO NOT** invent findings. Mark N/A with reason when a check does not apply. +- **DO NOT** FAIL P1 solely because the model uses pandas — only flag inconsistencies with the chosen style. -To validate upstream source completeness: +## Lineage cross-reference procedure -1. **Extract SQL tables from code cells:** Scan all code cells for SQL strings (f-strings, triple-quoted strings). Parse FROM and JOIN clauses to extract fully-qualified table references (`catalog.schema.table` or `schema.table`). -2. **Normalize references:** Strip catalog prefix if present. Convert to `schema.table` format. -3. **Exclude self-references:** Remove the model's own output table from the list. -4. **Exclude staging tables:** Remove any `stg_tmp_*` tables (these are transient staging tables created and dropped within the same ETL). -5. **Compare against declared sources:** Match the extracted set against `python_model_upstream_sources` entries. -6. **Report gaps:** Any table in SQL but not in upstream_sources = missing edge (L4 failure). Any entry in upstream_sources not found in SQL = stale entry (L5 warning). +1. **Extract Trino tables from code:** Scan cells for SQL strings (FROM/JOIN), `read_partition_to_dataframe(…)` first argument, and f-string table FQNs. +2. **Normalize:** Prefer `catalog.schema.table`. Accept `schema.table` with catalog from `get_catalog_name()` or `output.database`. +3. **Exclude:** model's own output table, `stg_tmp_*` staging tables, CTE names. +4. **Compare to `upstream_sources`:** Every remaining table needs a matching `trino` entry (L5). Extra entries with no code reference → stale (L6). +5. **Compare to `depends_on`:** When a read table is another Python model's output in the same topic, that model's `name` must appear in `depends_on` (L7). ## Conventions reference - **File layout:** `dags/python_models///.python.json` + `.python.py` - **Model ID:** `python______` -- **Source of truth:** `cells` mode → `.python.json`; `hand-written` mode → `.python.py` (metadata in JSON) -- **`_trino_io.py` DML helpers:** `execute_trino`, `overwrite_partition` / `overwrite` (keyword `insert_sql=` or `source_query=` + optional `columns=`), `append`, `merge`, `delete`, `update` -- **Script entry:** `if __name__ == "__main__":` with `build_context_from_env()` from `_config.py` — Airflow and KPO run the file as `__main__` -- **Scaffold convention:** `def run_etl(context)` orchestrates extract/transform/cleanup; recommended but not enforced at discovery -- **Partition contract:** downstream dbt models expect `portal_partition_daily` column in output -- **Table properties for lineage:** `python_model_name`, `python_model_type`, `python_model_namespace`, `python_model_table`, `python_model_description`, `python_model_upstream_sources` +- **Source of truth for code:** `.python.py` — edit it directly (DJ never rewrites it from JSON). `.python.json` holds metadata + a non-authoritative `cells` snapshot. Never hand-edit generated `_trino_io.py` / `_config.py` / `_model_tasks.py` +- **`_trino_io` SQL path:** `execute_sql`, `fetch_value`, `overwrite_partition`, `overwrite`, `append`, `merge`, `update` +- **`_trino_io` pandas path:** `read_partition_to_dataframe`, `overwrite_partition_from_dataframe` +- **Do not use:** `execute_trino`, inline `trino.dbapi.connect`, standalone `delete()` +- **Script entry:** `if __name__ == "__main__":` with `build_context_from_env()` from `_config.py` +- **`upstream_sources` trino format:** `catalog.schema.table` (canonical); `schema.table` legacy +- **`upstream_sources` external format:** `s3::…`, `api::…`, `jenkins::…` +- **Partition contract:** `portal_partition_daily` in partition-based models +- **Table properties:** `python_model_name`, `python_model_type`, `python_model_namespace`, `python_model_table`, `python_model_description`, `python_model_upstream_sources` ## Detailed checklist reference -For extended explanations, code examples, and edge cases for each check, see [references/review-checklist.md](references/review-checklist.md). +For extended explanations, code examples, and edge cases, see [references/review-checklist.md](references/review-checklist.md). + +## Evaluation scenarios + +For prompt → expected-finding cases (which check ID each condition should raise), see [evals/evals.md](evals/evals.md). diff --git a/templates/skills/dj-review-python-model/evals/evals.md b/templates/skills/dj-review-python-model/evals/evals.md new file mode 100644 index 0000000..59752cf --- /dev/null +++ b/templates/skills/dj-review-python-model/evals/evals.md @@ -0,0 +1,55 @@ +# Evals — dj-review-python-model + +Prompt → expected-finding cases for the read-only Python model review. Each case describes an input model condition and the finding the review must raise (with its check ID). Use these to sanity-check the skill after edits. + +The review treats `.python.py` as the **code source of truth** (hand-edits are expected, not drift). `.python.json` holds metadata + a non-authoritative `cells` snapshot. Only `_trino_io.py` / `_config.py` / `_model_tasks.py` are generated. + +--- + +## E1 — Missing upstream_source for a table read + +**Input:** model reads `glue_development.analytics.user_sessions` in SQL but `upstream_sources` omits it. + +**Expected:** **L5 ISSUE** — completeness gap; the table read in code has no matching `trino` entry. Report lists tables-read vs declared and the missing FQN. + +## E2 — Inline Trino connection + +**Input:** `.python.py` uses `from trino.dbapi import connect` / raw HTTP to Trino. + +**Expected:** **F5 ISSUE** — inline Trino code; recommend `_trino_io` helpers. HIGH priority. + +## E3 — Stale upstream_sources entry + +**Input:** `upstream_sources` lists `old_events_table` that no code path reads. + +**Expected:** **L6 WARNING** — stale/phantom entry; misleading graph edge. + +## E4 — Sibling read without depends_on + +**Input:** model reads another Python model's output table but `depends_on` does not list that model. + +**Expected:** **L7 ISSUE** — race condition; `depends_on` must list the sibling for task ordering (in addition to the `upstream_sources` lineage edge). + +## E5 — Import-unsafe module + +**Input:** `.python.py` has a bare `run_etl(context)` at module top level. + +**Expected:** **F8 ISSUE** — not import-safe; would execute on import under Airflow worker / KPO. Recommend the `__main__` block as the sole entry. + +## E6 — Pure pandas model is not penalized + +**Input:** coherent pandas end-to-end model (`extract → transform → load`). + +**Expected:** **P1 PASS** — do **not** flag solely because pandas is used. Only flag genuine inconsistencies with the chosen style. + +## E7 — Legacy 2-part trino FQN + +**Input:** `upstream_sources` uses `opus_python_source.tenant_events` (2-part). + +**Expected:** **L3 WARNING (LOW)** — legacy format still parsed; recommend migrating to `catalog.schema.table`. + +## E8 — Read-only guarantee + +**Input:** any model. + +**Expected:** the skill produces a structured report and changes **no** files; it does not execute SQL against Trino. diff --git a/templates/skills/dj-review-python-model/references/review-checklist.md b/templates/skills/dj-review-python-model/references/review-checklist.md index f88e626..af5daf1 100644 --- a/templates/skills/dj-review-python-model/references/review-checklist.md +++ b/templates/skills/dj-review-python-model/references/review-checklist.md @@ -12,115 +12,130 @@ Extended explanations, code examples, and edge cases for each check in the DJ Py ```json { - "name": "backstage_catalogs", + "name": "tenant_events", "group": "etl", - "topic": "api_data" + "topic": "tenant_status" } ``` -**Fail example:** Missing `topic` field — model cannot be discovered by Airflow or synced by the framework. +**Fail:** Missing `topic` — model cannot be discovered by Airflow or synced by the framework. ### F2: Naming convention **Pattern:** `^[a-z][a-z0-9_]*$` -**Pass:** `backstage_catalogs`, `user_embeddings`, `daily_snapshot` -**Fail:** `BackstageCatalogs` (PascalCase), `backstage-catalogs` (hyphens), `2024_snapshot` (starts with digit) +**Pass:** `tenant_events`, `user_embeddings`, `daily_snapshot` +**Fail:** `TenantEvents` (PascalCase), `tenant-events` (hyphens), `2024_snapshot` (starts with digit) -### F3: Cells array +### F3: Code source of truth -**Pass condition:** `cells` array exists and has at least one code cell. +**Pass condition:** `.python.json` is valid with the required metadata, and a companion `.python.py` exists with authored code. -**Why it matters:** The framework regenerates `.python.py` from `cells`. If `cells` is empty/missing and `.python.py` exists, the sync service skips overwriting (safety valve) — but this means the JSON is no longer the source of truth. +`.python.py` is the source of truth for code. The `cells` array in `.python.json` is an optional, non-authoritative snapshot of the last notebook save — it may be empty when the model is authored by editing `.python.py` directly, which is fine and not a failure. ### F4: `__main__` script entry -**Pass condition:** `.python.py` contains `if __name__ == "__main__":` and uses `build_context_from_env()` from `_config.py` (scaffold default: `run_etl(build_context_from_env())`). +**Pass condition:** `.python.py` contains: -**Why it matters:** Airflow worker and KPO pods execute the script as `__main__` via `runpy.run_path` or `python ` — not import + callable lookup. +```python +if __name__ == "__main__": + from python_models._config import build_context_from_env + run_etl(build_context_from_env()) +``` -**Common failure:** Missing `__main__` block after migrating from the old import-based execution model. +**Why it matters:** Airflow worker and KPO pods execute the script as `__main__` via `python ` — not import + callable lookup. `build_context_from_env()` reads `PYMODEL_DS`, `PYMODEL_DS_NODASH`, `PYMODEL_DATES` and configures stdout logging for KPO. ### F5: `_trino_io` usage **Pass condition:** All Trino operations use imports from `python_models._trino_io`. -**Fail patterns to flag:** +**Fail patterns:** + ```python -# INLINE CONNECTION — violates framework +# INLINE CONNECTION from trino.dbapi import connect -conn = connect(host='trino-host', port=8080, ...) +conn = connect(host="trino-host", port=8080, ...) -# RAW HTTP — violates framework -import requests -response = requests.post('http://trino:8080/v1/statement', ...) +# DEPRECATED API NAME +from python_models._trino_io import execute_trino -# SQLALCHEMY — violates framework -from sqlalchemy import create_engine -engine = create_engine('trino://...') +# RAW HTTP +response = requests.post("http://trino:8080/v1/statement", ...) ``` -**Correct pattern:** +**Correct patterns:** + ```python -from python_models._trino_io import execute_trino, overwrite_partition +# SQL path +from python_models._trino_io import execute_sql, fetch_value, overwrite_partition + +# Pandas path +from python_models._trino_io import read_partition_to_dataframe, overwrite_partition_from_dataframe ``` ### F6: OUTPUT_CONFIG -**Pass condition:** Code contains `OUTPUT_CONFIG = PythonModelConfig(...)` with import from `python_models._config`. +**Pass condition:** ```python from python_models._config import PythonModelConfig OUTPUT_CONFIG = PythonModelConfig( - model_name="backstage_catalogs", + model_name="tenant_events", model_type="python", - namespace="python", - description="Fetches catalog data from Backstage API", + namespace="opus_python_source", + table_name="tenant_events", + description="Daily tenant events from API", ) ``` -**Fail:** Custom dataclass, plain dict, or missing `OUTPUT_CONFIG` variable entirely. +**Fail:** Custom dataclass, plain dict, or missing `OUTPUT_CONFIG`. -### F7: Sync drift detection +### F7: Code source of truth (no JSON overwrite) -**How to check:** Compare the cells in `.python.json` against the content of `.python.py`. The framework generates `.py` by concatenating code cells (excluding the runner cell) separated by `\n\n\n`. +`.python.py` is the source of truth for code; DJ never rewrites it from `.python.json`. Hand-edited code is **expected and preserved** — it is not lost on sync. (The `cells` array in JSON is only a snapshot written when the notebook is saved.) -**Warning condition:** `.python.py` contains code not derivable from JSON cells — indicates hand-edits that will be lost on next sync. +**Pass condition:** `.python.py` is present and contains the authored ETL code. -### F8: Runner cell position +### F8: Import safety + +**Pass condition:** `.python.py` has no side-effectful top-level statements — e.g. no bare `run_etl(context)` at module scope (which would execute on import). Execution goes through the `__main__` block. -**Pass condition:** The last code cell in the `cells` array is the notebook runner: ```python +# FAIL — runs on import run_etl(context) + +# PASS — script entry only +if __name__ == "__main__": + from python_models._config import build_context_from_env + run_etl(build_context_from_env()) ``` -**Why it matters:** This cell is stripped when generating `.python.py` (it's notebook-only for interactive execution). If it's not last, other cells after it would also be stripped. +**Why it matters:** Airflow workers import the module and KPO runs `python ` — a top-level runner call would fire at import time. Only the derived `# Python Model: …` header is notebook-only; all other code in the notebook view is real module code. -### F9: Standard function structure +### F9: Orchestrator structure -**Pass condition:** Code follows the ETL function contract: -```python -def extract(context): - """Fetch data from external source.""" - ... +**Pass — pandas path:** -def transform_and_load(context): - """Transform and write to Iceberg.""" - ... +```python +def extract(context): ... +def transform(df, context): ... +def load(df, context): ... +def run_etl(context): ... +``` -def cleanup(context): - """Drop staging tables.""" - ... +**Pass — SQL path:** -def run_etl(context): - """Orchestrate the ETL.""" - extract(context) - transform_and_load(context) - cleanup(context) +```python +def extract(context): ... # optional +def transform_and_load(context): ... +def cleanup(context): ... +def run_etl(context): ... ``` -**Warning:** Missing `cleanup()` (staging tables may leak). Missing `extract()` is acceptable if data comes directly from existing Trino tables. +**Pass — mixed:** Documented combination (e.g. pandas extract + SQL transform_and_load). + +**Warning:** SQL path with staging but no `cleanup()`. Missing `extract()` is acceptable when data is already in Trino. --- @@ -128,120 +143,139 @@ def run_etl(context): ### L1: Table properties emission -**Pass condition:** `PythonModelConfig` class (from `_config.py`) has a `table_properties` property that returns all required keys: +DJ writes Iceberg properties after each successful MWAA or KPO run from `.python.json` metadata: -```python -@property -def table_properties(self) -> dict[str, str]: - return { - "python_model_name": self.model_name, - "python_model_type": self.model_type, - "python_model_namespace": self.namespace, - "python_model_table": self.table_name or self.model_name, - "python_model_description": self.description, - } -``` +- `python_model_name` ← `name` +- `python_model_namespace` ← `namespace` / `output.schema` +- `python_model_table` ← `table_name` / `output.table` +- `python_model_description` ← `description` +- `python_model_upstream_sources` ← `upstream_sources` -**How to check:** Verify `OUTPUT_CONFIG` fields populate all five keys. Missing `description` is a warning (lineage node renders without description). +Authors do not need `set_extra_properties` calls in model code for lineage. -### L2: `upstream_sources` is set in `.python.json` +### L2: `upstream_sources` when model reads upstream data -**Pass condition:** The `.python.json` file includes an `upstream_sources` array (or single object) with `{ "type": "trino"|"external", "value": "..." }` entries. DJ writes these to the output Iceberg table as `python_model_upstream_sources` after each successful run — authors do not need manual `ALTER TABLE` code in `.python.py`. +**Pass condition:** `.python.json` includes `upstream_sources` when the model reads Trino tables or external systems. -**Example (preferred):** ```json "upstream_sources": [ - { "type": "trino", "value": "opus_python_source.raw_events" }, - { "type": "external", "value": "backstage_api" } + { "type": "trino", "value": "glue_development.opus_python_source.tenant_events" }, + { "type": "external", "value": "api::backstage_tenant_events" } ] ``` -**Fail:** `upstream_sources` is missing or empty — lineage will show the Python model node but no upstream edges beyond the immediate source. - -### L3: Format validation +**Fail:** Model reads upstream data but `upstream_sources` is missing or empty — lineage shows the Python node with no upstream edges. -**Pass condition:** Each `trino` entry's `value` is `schema.table` format. Each `external` entry's `value` is the intended identifier (API name, S3 path, etc.). +**N/A:** Model with no upstream reads (pure external generate with no declared source) — rare; confirm with author. -**Pass (trino):** `{ "type": "trino", "value": "opus_python_source.raw_events" }` -**Fail (trino):** `{ "type": "trino", "value": "iceberg_catalog.opus_python_source.raw_events" }` (includes catalog — lineage uses schema.table only), `{ "type": "trino", "value": "raw_events" }` (missing schema) +### L3: Trino format validation -### L4: Completeness — no missing upstream edges +**Pass (canonical):** -**Procedure:** +```json +{ "type": "trino", "value": "glue_development.opus_python_source.tenant_events" } +``` -1. Extract all table references from SQL in code cells: - ```sql - -- These patterns indicate source tables: - FROM catalog.schema.table - JOIN catalog.schema.table - FROM schema.table - JOIN schema.table - ``` +**Warning (legacy 2-part):** -2. Normalize to `schema.table` (strip catalog prefix). +```json +{ "type": "trino", "value": "opus_python_source.tenant_events" } +``` -3. Exclude: - - The model's own output table - - Staging tables (`stg_tmp_*`, `tmp_*`) - - Inline CTEs (WITH clause names) +Still parsed with catalog from parent context — recommend migrating to 3-part FQN. -4. Every remaining table must appear in `upstream_sources` (as a `trino` entry). +**Fail:** -**Example finding:** +```json +{ "type": "trino", "value": "tenant_events" } ``` -Tables in SQL: .raw_events, analytics.user_sessions, reference.countries -Tables in upstream_sources: .raw_events, analytics.user_sessions -Missing from upstream_sources: reference.countries + +Missing schema — cannot resolve table location. + +### L4: External format validation + +**Pass:** + +```json +{ "type": "external", "value": "s3::data-lake/tenant_migrations/{ds}/migrations.csv" } +{ "type": "external", "value": "api::backstage_tenant_events" } +{ "type": "external", "value": "jenkins::nightly-export-job" } ``` -### L5: Accuracy — no stale entries +**Warning:** Free-form string with no prefix convention — acceptable if stable and documented. + +### L5: Completeness — no missing upstream edges -**Procedure:** Every entry in `upstream_sources` must correspond to a source actually read by the model. +**Procedure:** + +1. Extract Trino tables from: + - SQL strings: `FROM catalog.schema.table`, `JOIN schema.table` + - `read_partition_to_dataframe("catalog.schema.table", …)` + - F-strings building FQNs from `get_catalog_name()` +2. Normalize to `catalog.schema.table`. +3. Exclude: own output table, `stg_tmp_*` tables, CTE aliases. +4. Every remaining table must match a `trino` entry in `upstream_sources`. **Example finding:** + ``` -Tables in upstream_sources: .raw_events, .old_events -Tables in SQL: .raw_events -Stale in upstream_sources: .old_events +Tables read in code: glue_development.opus_python_source.tenant_events, glue_development.analytics.user_sessions +upstream_sources trino entries: glue_development.opus_python_source.tenant_events +Missing: glue_development.analytics.user_sessions ``` -**Severity:** Warning (stale entries don't break lineage but create misleading edges in the graph). +### L6: Accuracy — no stale entries + +Every `upstream_sources` entry must correspond to a source actually read. -### L6: DAG assignment +**Example finding:** -**Pass:** `"dags": ["daily_etl_dag"]` — model is scheduled. -**Warning:** `"dags": []` — model is a utility module. Lineage still works via table properties, but there's no Airflow task node in the DAG graph. +``` +upstream_sources: tenant_events, old_events_table +Tables read in code: tenant_events only +Stale: old_events_table +``` + +**Severity:** Warning — misleading graph edges. ### L7: `depends_on` alignment -**Pass condition:** If model A reads from the output table of model B, then model A's `depends_on` should include model B's name. +**Pass condition:** When model A reads the output table of Python model B, model A's `depends_on` includes B's `name`. -**Example:** ```json -// Model: user_features.python.json { - "depends_on": ["user_embeddings"], // user_embeddings writes to python.user_embeddings - ... + "name": "tenant_migration_status", + "depends_on": ["tenant_events", "tenant_migrations"], + "upstream_sources": [ + { "type": "trino", "value": "glue_development.opus_python_source.tenant_events" }, + { "type": "trino", "value": "glue_development.opus_python_source.tenant_migrations" } + ] } ``` -**Why it matters:** `depends_on` wires Airflow task dependencies. Without it, model A may execute before model B finishes writing, reading stale/empty data. +**Why both matter:** + +- `depends_on` — Airflow waits for upstream tasks +- `upstream_sources` — Model Lineage draws edges -### L8: Model ID consistency +DJ may auto-infer trino `upstream_sources` from `depends_on` — inferred entries are valid. Missing `depends_on` when reading a sibling Python output is an **issue** (race condition). -**Pass condition:** File path `dags/python_models/etl/api_data/backstage_catalogs.python.json` → model ID `python__etl__api_data__backstage_catalogs` → matches `python_model_name` in table properties. +### L8: DAG assignment -### L9: Companion `.python.py` exists +**Pass:** `"dags": ["daily_etl_dag"]` +**Warning:** `"dags": []` — utility module; lineage via table properties only, no Airflow task node. -**Pass condition:** For file `.python.json`, the file `.python.py` exists in the same directory. +### L9: Model ID consistency -**Fail:** Airflow cannot discover or execute the model without the `.py` file. +File path `dags/python_models/etl/tenant_status/tenant_events.python.json` → `python__etl__tenant_status__tenant_events` → matches `python_model_name` in Iceberg properties. -### L10: Downstream source reference match +### L10: Companion `.python.py` exists -**Pass condition:** If any `.source.json` in the project references this Python model's output table, the table name must match `OUTPUT_CONFIG.table_name` (or `model_name` fallback). +**Fail:** Missing `.python.py` — Airflow cannot execute the model. -**How to check:** Search project for `.source.json` files containing the model's output table name in their `tables` array. +### L11: Downstream source reference match + +If `.source.json` files reference this model's output, table name must match `OUTPUT_CONFIG.table_name` (or `model_name` fallback). --- @@ -249,181 +283,141 @@ Stale in upstream_sources: .old_events ### D1: Standard catalog/schema -**Pass condition:** Output table uses the project-standard catalog and schema for Python model outputs. - -**Default convention:** `..` +**Default:** `glue_development.opus_python_source.` -**When to flag:** Custom schema is acceptable but should be documented in `description` or `variables`. Flag if namespace doesn't match any known project convention. +Custom schema is acceptable when documented. Flag unknown namespaces. ### D2: Partition column emission -**Pass condition:** Output INSERT SQL includes `portal_partition_daily` column: +**Pass (SQL):** ```sql -INSERT INTO {table} -SELECT - col_a, - col_b, - '{ds}' AS portal_partition_daily -FROM ... +SELECT col_a, '{ds}' AS portal_partition_daily FROM ... ``` -**Why it matters:** Downstream dbt models and the DJ framework expect this column for incremental processing. Without it, `overwrite_partition` DML will fail. +**Pass (pandas):** -**Exception:** Models with `output.write_mode: "overwrite"` (full table refresh) may omit partition columns — flag as suggestion, not issue. +```python +df["portal_partition_daily"] = context["ds"] +``` + +**Exception:** `output.write_mode: "overwrite"` full refresh — suggestion only. ### D3: Column naming -**Pass condition:** All output column names in the INSERT SQL follow `snake_case`: -- `user_id` (pass) -- `created_at` (pass) -- `userId` (fail — camelCase) -- `Created At` (fail — spaces) +All output columns in `snake_case`. Flag `userId`, `Created At`. ### D4: No `SELECT *` -**Pass condition:** Production INSERT statements enumerate columns explicitly. - ```sql --- FAIL: -INSERT INTO {table} SELECT * FROM staging +-- FAIL +INSERT INTO target SELECT * FROM staging --- PASS: -INSERT INTO {table} -SELECT user_id, email, created_at, portal_partition_daily -FROM staging +-- PASS +INSERT INTO target SELECT user_id, email, portal_partition_daily FROM staging ``` -**Why it matters:** `SELECT *` makes the output schema dependent on upstream schema changes. A column added upstream silently changes the downstream contract. - ### D5: Table name matches model name -**Pass condition:** `OUTPUT_CONFIG.table_name` (or `model_name` fallback) matches what downstream consumers reference. - -Convention: `python______` → table `` in namespace `python`. +Convention: table `` in output schema matches model `name` unless `table_name` override is intentional and documented. ### D6: Deterministic output -**Warning conditions:** -- Non-deterministic functions without ORDER BY (e.g., `ROW_NUMBER()` without deterministic tie-breaking) -- UNION without explicit column ordering -- JSON field extraction where key order varies +**Warning:** `ROW_NUMBER()` without deterministic tie-break, `NOW()` as data column, non-deterministic UDFs. --- ## 4. Performance & Best Practices (P) -### P1: SQL-first adherence +### P1: Orchestrator consistency (not SQL-first) -**Flag these as violations:** +**Accept — pandas end-to-end:** ```python -# VIOLATION: filtering in pandas instead of SQL WHERE -df = df[df['status'] == 'active'] +df = extract(context) +df = transform(df, context) +load(df, context) # overwrite_partition_from_dataframe +``` -# VIOLATION: aggregation in pandas instead of SQL GROUP BY -result = df.groupby('category').agg({'amount': 'sum'}) +**Accept — SQL end-to-end:** -# VIOLATION: join in pandas instead of SQL JOIN -merged = pd.merge(df1, df2, on='user_id') +```python +extract(context) +transform_and_load(context) # overwrite_partition +cleanup(context) +``` -# VIOLATION: dedup in pandas instead of SQL ROW_NUMBER -df = df.drop_duplicates(subset=['user_id']) +**Accept — mixed** when documented. -# VIOLATION: type casting in pandas instead of SQL CAST -df['amount'] = df['amount'].astype(float) -``` +**Flag as warning — inconsistency only:** -**Acceptable pandas usage:** ```python -# OK: parsing API JSON response into rows -df = pd.json_normalize(api_response['data']) - -# OK: chunking for batch INSERT staging -for chunk in np.array_split(df, num_batches): - stage_to_trino(chunk) +# Author chose SQL path but filters in pandas after staging unnecessarily +df = df[df["status"] == "active"] # could be SQL WHERE ``` +Do **not** fail solely because pandas is used. Fail only when the chosen style is contradicted without justification. + ### P2: Staging table cleanup -**Pass condition:** `cleanup()` function drops all staging tables: +**SQL path pass:** ```python def cleanup(context): - execute_trino("DROP TABLE IF EXISTS iceberg_catalog.stg_tmp.backstage_raw") + execute_sql(f"DROP TABLE IF EXISTS {staging_fqn}") ``` -**Fail:** No `cleanup()` function, or staging tables created in `extract()` not referenced in `cleanup()`. - -**Impact:** Orphaned staging tables accumulate storage costs and namespace pollution. +**Fail:** Staging tables created but never dropped. ### P3: Partition predicate pushdown -**Pass condition:** Queries against large source tables include partition column in WHERE: - ```sql --- PASS: partition filter enables pushdown -SELECT * FROM source_table -WHERE portal_partition_daily = '{ds}' +-- PASS +SELECT * FROM source WHERE portal_partition_daily = '{ds}' --- FAIL: full table scan -SELECT * FROM source_table +-- FAIL on large partitioned tables +SELECT * FROM source WHERE name = 'foo' ``` -### P4: Batch staging +Also applies to `read_partition_to_dataframe` — should pass partition column and value. + +### P4: Batch staging / writes -**Pass condition:** For large datasets, staging uses batched inserts: +**Pass:** ```python -BATCH_SIZE = 10000 -for i in range(0, len(df), BATCH_SIZE): - batch = df.iloc[i:i+BATCH_SIZE] - values = format_values(batch) - execute_trino(f"INSERT INTO staging_table VALUES {values}") +overwrite_partition_from_dataframe(..., batch_size=1000) ``` -**Warning trigger:** Single INSERT with > 100,000 rows or no batching visible for API pagination results. +**Warning:** Single unbounded INSERT of millions of rows. ### P5: Write mode appropriateness -| Data pattern | Correct write mode | Incorrect | -|-------------|-------------------|-----------| -| Daily snapshot, rerun-safe | `overwrite_partition` | `append` (duplicates on rerun) | -| Event stream, never reprocess | `append` | `overwrite_partition` (loses history) | -| Full refresh, small table | `overwrite(..., "true", insert_sql=...)` | `overwrite_partition` (unnecessary complexity) | -| Upsert by key | `merge` | `overwrite(..., "true", ...)` (loses unmatched rows) | +| Pattern | Correct | Incorrect | +| ------- | ------- | --------- | +| Daily idempotent | `overwrite_partition` / `overwrite_partition_from_dataframe` | `append` | +| Event stream | `append` | `overwrite_partition` | +| Full refresh | `overwrite(…, "true", …)` | unnecessary partition logic | ### P6: Partition filters on large tables -**Flag:** Any FROM/JOIN against a known large table (Iceberg table with partitions) without a WHERE clause filtering on partition columns. +Flag FROM/JOIN against partitioned Iceberg tables without partition predicates. -### P7: DataFrame restriction +### P7: Correct API names -**Pass condition:** pandas/polars usage is limited to: -- API response parsing (`pd.json_normalize`, `pd.DataFrame(response_data)`) -- CSV/file reading (`pd.read_csv`) -- Batch chunking for staging -- Data that cannot exist in Trino yet (external ingestion) +| Wrong | Correct | +| ----- | ------- | +| `execute_trino(sql)` | `execute_sql(sql)` | +| `execute_trino(sql, return_result=True)` | `fetch_value(sql)` | +| `delete(table, where)` | `overwrite_partition(…)` or `overwrite(…)` | -**Fail:** Using DataFrame operations on data already in Trino (pull → transform → push back). +### P8: Type safety -### P8: Explicit CAST - -**Pass condition:** Type conversions in SQL INSERT use explicit CAST: - -```sql -SELECT - CAST(raw_amount AS DECIMAL(18,2)) AS amount, - CAST(event_date AS DATE) AS event_date -FROM staging -``` - -**Fail:** Relying on implicit coercion or pandas `.astype()` for data that lands in Trino. +**SQL path:** `CAST(raw_amount AS DECIMAL(18,2))` +**Pandas path:** appropriate dtypes before `overwrite_partition_from_dataframe` ### P9: Error handling -**Pass condition:** `run_etl()` ensures cleanup runs even on failure: - ```python def run_etl(context): try: @@ -433,34 +427,34 @@ def run_etl(context): cleanup(context) ``` -**Fail:** Linear call without try/finally — staging tables leak on transform failure. - ### P10: Idempotency -**Pass condition:** For `overwrite_partition` models, re-running with the same `ds` produces identical output (DELETE + INSERT pattern). - -**Warning conditions:** -- `append` mode used with partition-based data (duplicates on rerun) -- Non-deterministic expressions in SELECT without dedup logic -- Timestamp-based columns that change on each execution (e.g., `NOW()` as `loaded_at`) +For partition-based models, rerun with same `ds` replaces the partition — no duplicate rows. --- ## Edge Cases -### Models without SQL (pure Python-to-S3) +### Models without SQL (pure pandas) -If `output_type: "s3"`, several checks become N/A: -- D2 (partition column) — S3 models may use Hive-style partitioning differently -- P3/P6 (partition pushdown) — no Trino source queries -- L2 (upstream_sources) — may not be discoverable via Iceberg properties +Checks D4, P3, P6 may be N/A. `upstream_sources` with `external` entries still required when reading external systems. Trino reads via `read_partition_to_dataframe` still need `trino` entries. -Mark these as `N/A (S3 output model)` in the report. +### Models without Trino (S3 output) + +When `output_type: "s3"`: + +- D2, P3, P6 — N/A +- L2 — external `upstream_sources` still recommended +- F5 — `_trino_io` may be unused — N/A with reason ### Utility modules (empty `dags`) -Models with `"dags": []` are importable Python utilities, not scheduled ETL. Checks L6, L7, L9 become N/A. Framework compliance checks still apply. +L8, L7 (scheduling) become N/A. Framework compliance still applies. + +### Auto-inferred upstream_sources + +When `depends_on` is set and trino entries were auto-inferred at create time, L5 may pass even if the author did not hand-write `upstream_sources`. Verify inferred FQNs match actual output tables of upstream models. -### Models reading from other Python model outputs +### Topic pipelines -If this model reads from another Python model's output table (not an external source), `depends_on` is critical (L7) and the upstream table should appear in `python_model_upstream_sources` (L4). +Review each model in the topic. Cross-check that integration models list all sibling `depends_on` entries and that `upstream_sources` trino FQNs use `catalog.schema.table`. diff --git a/web/src/context/app.tsx b/web/src/context/app.tsx index 18d797e..2c8e20f 100644 --- a/web/src/context/app.tsx +++ b/web/src/context/app.tsx @@ -1861,7 +1861,7 @@ SELECT * FROM final`, case 'framework-get-available-dags': { return resolve( apiResponse({ - dags: ['source_etl', 'cloud_dir_dag', 'data_pipeline'], + dags: ['cloud_dir_dag', 'data_pipeline'], }), ); } diff --git a/web/src/forms/FieldSelectSingle.tsx b/web/src/forms/FieldSelectSingle.tsx index bd4325b..41cce30 100644 --- a/web/src/forms/FieldSelectSingle.tsx +++ b/web/src/forms/FieldSelectSingle.tsx @@ -11,6 +11,8 @@ export type FieldSelectSingleProps = ControllerRenderProps & { inputClassName?: string; labelClass?: string; helpIcon?: React.ReactNode; // Help icon for Assist Me + disabled?: boolean; + placeholder?: string; /** * Virtualize the options list. Enable for large option sets (hundreds+), * e.g. Trino tables, so the dropdown doesn't render every row into the DOM diff --git a/web/src/pages/ModelLineage/LineageGraph.tsx b/web/src/pages/ModelLineage/LineageGraph.tsx index 09c6dc8..e5f730d 100644 --- a/web/src/pages/ModelLineage/LineageGraph.tsx +++ b/web/src/pages/ModelLineage/LineageGraph.tsx @@ -350,6 +350,7 @@ export default function LineageGraph({ name: currentNode.name, type: currentNode.type, description: currentNode.description, + tags: currentNode.tags, path: currentNode.path, pathSystem: currentNode.pathSystem, isCurrent: true, @@ -387,6 +388,7 @@ export default function LineageGraph({ name: node.name, type: node.type, description: node.description, + tags: node.tags, path: node.path, pathSystem: node.pathSystem, isCurrent: false, @@ -450,6 +452,7 @@ export default function LineageGraph({ name: node.name, type: node.type, description: node.description, + tags: node.tags, path: node.path, pathSystem: node.pathSystem, isCurrent: false, @@ -512,6 +515,7 @@ export default function LineageGraph({ name: node.name, type: node.type, description: node.description, + tags: node.tags, path: node.path, pathSystem: node.pathSystem, isCurrent: false, diff --git a/web/src/pages/ModelLineage/ModelNode.tsx b/web/src/pages/ModelLineage/ModelNode.tsx index aa01c81..8afaf96 100644 --- a/web/src/pages/ModelLineage/ModelNode.tsx +++ b/web/src/pages/ModelLineage/ModelNode.tsx @@ -1,4 +1,5 @@ import { + ArrowDownTrayIcon, CircleStackIcon, CodeBracketIcon, CogIcon, @@ -28,6 +29,39 @@ const TYPE_LABELS = { python: 'PY', }; +type PythonLineageKind = 'model' | 'trino-ref' | 'external'; + +function getPythonLineageKind(tags?: string[]): PythonLineageKind | null { + if (!tags?.includes('python')) { + return null; + } + if (tags.includes('external')) { + return 'external'; + } + if (tags.includes('trino-ref')) { + return 'trino-ref'; + } + return 'model'; +} + +const PYTHON_KIND_ICONS: Record = { + model: CodeBracketIcon, + 'trino-ref': CircleStackIcon, + external: ArrowDownTrayIcon, +}; + +const PYTHON_KIND_LABELS: Record = { + model: 'PY', + 'trino-ref': 'TRINO', + external: 'EXT', +}; + +const PYTHON_KIND_BADGE_STYLES: Record = { + model: 'bg-surface border-neutral', + 'trino-ref': 'bg-cyan-600/15 border-cyan-600/40', + external: 'bg-amber-600/15 border-amber-600/40', +}; + const MATERIALIZATION_LABELS: Record = { ephemeral: 'ephemeral', incremental: 'incremental', @@ -67,14 +101,21 @@ export default function ModelNode({ data }: { data: ModelNodeData }) { onExpandUpstream, onExpandDownstream, onViewColumns, + tags, } = data; + const pythonKind = type === 'python' ? getPythonLineageKind(tags) : null; + // Determine if we need to compile before running // Compile if: model is outdated (source changed) OR no compiled file exists const needsCompilation = isOutdated === true || hasCompiledFile === false; - const TypeIcon = TYPE_ICONS[type]; - const typeLabel = TYPE_LABELS[type]; + const TypeIcon = + pythonKind != null ? PYTHON_KIND_ICONS[pythonKind] : TYPE_ICONS[type]; + const typeLabel = + pythonKind != null ? PYTHON_KIND_LABELS[pythonKind] : TYPE_LABELS[type]; + const pythonBadgeStyle = + pythonKind != null ? PYTHON_KIND_BADGE_STYLES[pythonKind] : ''; // Build title with name and description const nodeTitle = description ? `${name}\n\n${description}` : name; @@ -156,7 +197,12 @@ export default function ModelNode({ data }: { data: ModelNodeData }) { !isStale && hasDownstream && !isDownstreamExpanded && onExpandDownstream; const isClickable = !isStale && type !== 'python'; - const displayName = type === 'python' && description ? description : name; + const displayName = + pythonKind === 'trino-ref' && description + ? description + : type === 'python' && description && pythonKind === 'model' + ? description + : name; return (
{/* Type icon and label */} -
+
{typeLabel} diff --git a/web/src/pages/PythonModelCreate.tsx b/web/src/pages/PythonModelCreate.tsx index fcd3207..f590782 100644 --- a/web/src/pages/PythonModelCreate.tsx +++ b/web/src/pages/PythonModelCreate.tsx @@ -71,6 +71,7 @@ export function PythonModelCreate() { const [isProjectsLoading, setIsProjectsLoading] = useState(true); const [isSubmitting, setIsSubmitting] = useState(false); const [availableDags, setAvailableDags] = useState([]); + const [isDagsLoading, setIsDagsLoading] = useState(false); const [kpoSizes, setKpoSizes] = useState>({}); const projectName = watch('projectName'); @@ -112,10 +113,11 @@ export function PythonModelCreate() { ); const disableSubmit = useMemo(() => { - return ( - !projectName || !name || !group || !topic || !dags || dags.length === 0 - ); - }, [projectName, name, group, topic, dags]); + if (isDagsLoading || !projectName || !group || !topic || !dags?.length) { + return true; + } + return !name; + }, [isDagsLoading, projectName, name, group, topic, dags]); const kpoSizeOptions = useMemo( () => @@ -223,8 +225,11 @@ export function PythonModelCreate() { const fetchAvailableDags = async () => { if (!projectName) { setAvailableDags([]); + setIsDagsLoading(false); return; } + setAvailableDags([]); + setIsDagsLoading(true); try { const response = await api.post({ type: 'framework-get-available-dags', @@ -234,6 +239,8 @@ export function PythonModelCreate() { } catch (err) { console.error('Error fetching available DAGs:', err); setAvailableDags([]); + } finally { + setIsDagsLoading(false); } }; void fetchAvailableDags(); @@ -441,6 +448,10 @@ export function PythonModelCreate() { error={errors.dags as any} label="Select DAG" options={availableDags.map((d) => ({ label: d, value: d }))} + disabled={isDagsLoading || !projectName} + placeholder={ + isDagsLoading ? 'Loading DAGs...' : 'Select a DAG' + } tooltipText="The DAG where this model will run. The model will execute as part of the run_python_models task in the selected DAG." /> )} From 8754b8eba5ab3017b3569ade6a7b6ef545bcc37f Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Tue, 22 Sep 2026 11:10:56 +0530 Subject: [PATCH 07/11] Update python model skills --- CHANGELOG.md | 1 + docs/AGENT_SKILLS.md | 13 ++++++++ templates/_agents-dj/_AGENTS.md | 2 +- .../skills/dj-create-python-model/_SKILL.md | 20 +++++++++---- .../dj-create-python-model/evals/evals.md | 2 +- .../references/etl-patterns.md | 10 +++---- .../references/worked-example.md | 30 +++++++++---------- .../dj-migrate-notebook-to-pymodel/_SKILL.md | 7 +++++ .../skills/dj-review-python-model/_SKILL.md | 9 +++++- .../dj-review-python-model/evals/evals.md | 2 +- .../references/review-checklist.md | 20 ++++++------- 11 files changed, 76 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e00176..cbd3a87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ ### Python models +- **Project defaults for python-model skills.** Commit `.agents/project/skills//project-defaults.md` in your dbt project so `dj-create-python-model`, `dj-review-python-model`, and `dj-migrate-notebook-to-pymodel` pick up your catalog, schema, authoring mode, and migration rules in Step 0 — without forking DJ skill files. Example templates live in the DJ repo under `templates/_agents-project/skills/` (not deployed). DJ never overwrites `.agents/project/`. - **Declare upstream lineage in `.python.json`.** Add `"upstream_sources": [{ "type": "trino"|"external", "value": "..." }]` on any Python model; DJ auto-suggests trino entries from `depends_on` on JSON sync when omitted. After each successful Airflow or KPO run, DJ writes `python_model_name` and the structured `python_model_upstream_sources` property to the output Iceberg table — no manual `ALTER TABLE` in your `.python.py`. - **Model Lineage walks the full Python chain.** When a downstream dbt model sources a Python output table, Data Explorer recursively reads each upstream table's Iceberg properties until it reaches roots (external sources or tables with no further upstream), so a chain like D → C → B, A renders in one view. - **Cleaner lineage for Python models.** Python models now show their real upstream sources in the lineage graph, without the extra placeholder node that used to appear. diff --git a/docs/AGENT_SKILLS.md b/docs/AGENT_SKILLS.md index a7effce..6c649e9 100644 --- a/docs/AGENT_SKILLS.md +++ b/docs/AGENT_SKILLS.md @@ -48,6 +48,19 @@ DJ never overwrites `.agents/project/`. Starters live in the DJ repo at `templat - **You stay in control of DJ commands.** Skills can't run VS Code commands themselves; they'll ask you to run things like **`DJ: Sync to SQL and YML`** or **`DJ: Refresh Projects`** at the right moment. - **Some skills are read-only.** `dj-review-python-model`, `dj-govern-model`, `dj-trino-analyzer`, and `dj-verify-pymodel-parity` produce reports and change nothing. +## Project defaults (python-model skills) + +Downstream dbt projects can commit **project-owned defaults** that DJ python-model skills read in Step 0 — without forking the deployed `dj-*` skill files (which `DJ: Refresh Projects` overwrites). + +```text +.agents/project/skills/ + dj-create-python-model/project-defaults.md + dj-review-python-model/project-defaults.md + dj-migrate-notebook-to-pymodel/project-defaults.md +``` + +Optional sibling files (e.g. `trino-dialect-cheatsheet.md`) live in the same folder. Example templates ship in the DJ extension repo under `templates/_agents-project/skills/` for copy-paste into your project. + ## The skills DJ provides 20 skills, grouped below by what they help you do. diff --git a/templates/_agents-dj/_AGENTS.md b/templates/_agents-dj/_AGENTS.md index 0b962ac..4edb553 100644 --- a/templates/_agents-dj/_AGENTS.md +++ b/templates/_agents-dj/_AGENTS.md @@ -27,7 +27,7 @@ Follow the **always-relevant** rules below — project structure, naming, struct | Run a read-only Trino query to inspect data or schema — CLI resolution, connection env, invocation | [reference/running-trino.md](reference/running-trino.md) | | Run the Lightdash CLI (start-preview / download / upload) — connection env, settings, guardrails | [reference/running-lightdash.md](reference/running-lightdash.md) | | Commit, branch, or stage DJ work in git — what to commit, what to ignore, commit hygiene | [reference/git-workflow.md](reference/git-workflow.md) | -| Customize python-model agent skills (catalog, schema, migration rules) | `.agents/project/skills//project-defaults.md` (project-owned; DJ never overwrites) | +| Customize python-model agent skills for this project (catalog, schema, authoring mode, migration rules) | `.agents/project/skills//project-defaults.md` (project-owned; DJ never overwrites) | | Invoke DJ operations via the CLI bridge (`.dj/bin/dj`) — discovery, authoring, compile, query | `.agents/skills/dj-cli/SKILL.md` (skill/CLI routing: `dj-cli-registry`) | --- diff --git a/templates/skills/dj-create-python-model/_SKILL.md b/templates/skills/dj-create-python-model/_SKILL.md index ec61c2a..773f480 100644 --- a/templates/skills/dj-create-python-model/_SKILL.md +++ b/templates/skills/dj-create-python-model/_SKILL.md @@ -50,6 +50,14 @@ Use this skill when the user mentions: python model, ETL, data ingestion, API fe Shared topic helpers live at `dags/python_models///_helpers.py` (one per topic). Import from sibling models as needed. +## Step 0 — Resolve project defaults (before proposing values) + +1. Read `.agents/project/skills/dj-create-python-model/project-defaults.md` if it exists. +2. Read any sibling files in that folder linked from the defaults. +3. If missing, discover from the repo: existing models under `dags/python_models/`, python-source `.source.json` files, `dbt_project.yml` vars, `dj.*` VS Code settings. +4. Use resolved defaults for output catalog/schema, DAG, authoring mode, and groups — state them before scaffolding. +5. Never assume org-specific catalog, schema, or DAG names without confirmation. + ## Interactive gathering workflow Ask the following in order. Batch related questions to minimize round-trips. **Skip the Q&A entirely** when the user supplies a complete spec (name, group, topic, DAG, transform style, output, dependencies). @@ -102,8 +110,8 @@ Present defaults and confirm overrides: | Field | Default | Override? | | --------------------- | ---------------------------- | ----------------------- | -| `output.database` | `glue_development` | Yes | -| `output.schema` | `opus_python_source` | Yes | +| `output.database` | From Step 0 (``) | Yes | +| `output.schema` | From Step 0 (``) | Yes | | `output_type` | `iceberg` | Yes (alternative: `s3`) | | `output.write_mode` | `overwrite_partitions` | Yes | | `output.partition_by` | `["portal_partition_daily"]` | Yes | @@ -139,7 +147,7 @@ Declare in `.python.json` only — DJ writes properties after each run; **do not ```json "upstream_sources": [ - { "type": "trino", "value": "glue_development.opus_python_source.raw_events" }, + { "type": "trino", "value": "..raw_events" }, { "type": "external", "value": "api::backstage_catalog" }, { "type": "external", "value": "s3::my-bucket/path/data.csv" }, { "type": "external", "value": "jenkins::my-job-name" } @@ -148,7 +156,7 @@ Declare in `.python.json` only — DJ writes properties after each run; **do not | `type` | `value` format | Notes | | ------ | -------------- | ----- | -| **trino** | **`catalog.schema.table`** (canonical) | e.g. `glue_development.opus_python_source.model_a` | +| **trino** | **`catalog.schema.table`** (canonical) | e.g. `..model_a` | | **trino** | `schema.table` (legacy) | Still parsed; prefer 3-part FQN for new models | | **external** | Prefix convention | `s3::…`, `api::…`, `jenkins::…`, or other stable identifier | @@ -161,7 +169,7 @@ New models default to **`compute: "kpo"`** (Kubernetes pod). Set **`compute: "ai | `compute` | `kpo` (default) or `airflow` | | `kpo_size` | Pod preset when `compute` is `kpo`: `small`, `medium`, `large`, `xlarge` | -**KPO script path:** The `mwaa-opus-kpo-runner` image syncs `s3:///dags/python_models/` into the pod and runs `python //.python.py`. Set Airflow Variable **`python_models_s3_bucket`**; optional **`python_model_s3_region`** (default `us-west-2`). Configure **`airflow_runner_cfg`** (image, tag, namespace, cluster). +**KPO script path:** The configured KPO runner image (from project defaults or `airflow_runner_cfg`) syncs `s3:///dags/python_models/` into the pod and runs `python //.python.py`. Set Airflow Variable **`python_models_s3_bucket`**; optional **`python_model_s3_region`**. Configure **`airflow_runner_cfg`** (image, tag, namespace, cluster). **`PYMODEL_*` environment:** Airflow sets `PYMODEL_DS`, `PYMODEL_DS_NODASH`, and optionally `PYMODEL_DATES` (JSON array) before execution. `build_context_from_env()` in `_config.py` reads them and configures `logging.basicConfig` to stdout (required for KPO log visibility). @@ -349,7 +357,7 @@ Chains like downstream dbt → model C → model B + model A render in one graph Downstream dbt SQL models reference the output table using DJ source notation: ``` -glue_development__opus_python_source. +__. ``` (Double underscore separates catalog from schema.) diff --git a/templates/skills/dj-create-python-model/evals/evals.md b/templates/skills/dj-create-python-model/evals/evals.md index 8879b45..a91e68d 100644 --- a/templates/skills/dj-create-python-model/evals/evals.md +++ b/templates/skills/dj-create-python-model/evals/evals.md @@ -59,7 +59,7 @@ Reference the correct file model throughout: ## E5 — Destructive-write confirmation -**Prompt:** "Point this model at `glue_production.opus_python_source`." +**Prompt:** "Point this model at `.`." **Expected:** - The skill states the resolved `output.database` / `output.schema` / `write_mode` and asks for explicit confirmation before writing to a non-default (production) target, because `overwrite_partitions` is destructive. diff --git a/templates/skills/dj-create-python-model/references/etl-patterns.md b/templates/skills/dj-create-python-model/references/etl-patterns.md index a8d66ed..dae04cf 100644 --- a/templates/skills/dj-create-python-model/references/etl-patterns.md +++ b/templates/skills/dj-create-python-model/references/etl-patterns.md @@ -2,7 +2,7 @@ Stage-level patterns for both orchestrators. All Trino I/O goes through **`python_models._trino_io`** — no inline `trino.dbapi.connect`, no `execute_trino`, no standalone `delete()`. -Default output target: `glue_development.opus_python_source` (override via `output` in `.python.json` and `get_catalog_name()` / `get_namespace_name()` at runtime). +Default output target: `.` (resolve from Step 0 / `output` in `.python.json` and `get_catalog_name()` at runtime). Build `table_fqn` for every helper call: @@ -179,7 +179,7 @@ from python_models._trino_io import read_partition_to_dataframe def extract(context: dict) -> pd.DataFrame: """Read an upstream Iceberg partition declared in upstream_sources.""" ds = context["ds"] - upstream_fqn = f"{get_catalog_name()}.opus_python_source.upstream_model_a" + upstream_fqn = f"{get_catalog_name()}..upstream_model_a" return read_partition_to_dataframe( upstream_fqn, "portal_partition_daily", @@ -192,7 +192,7 @@ Declare in `.python.json`: ```json "upstream_sources": [ - { "type": "trino", "value": "glue_development.opus_python_source.upstream_model_a" } + { "type": "trino", "value": "..upstream_model_a" } ] ``` @@ -370,8 +370,8 @@ def transform_and_load(context: dict) -> None: ### Join with existing Trino tables ```python - source_fqn = f"{catalog}.opus_python_source.raw_events" - lookup_fqn = f"{catalog}.opus_python_source.region_lookup" + source_fqn = f"{catalog}..raw_events" + lookup_fqn = f"{catalog}..region_lookup" source_query = f""" SELECT diff --git a/templates/skills/dj-create-python-model/references/worked-example.md b/templates/skills/dj-create-python-model/references/worked-example.md index f449643..410f259 100644 --- a/templates/skills/dj-create-python-model/references/worked-example.md +++ b/templates/skills/dj-create-python-model/references/worked-example.md @@ -23,7 +23,7 @@ DJ auto-infers trino `upstream_sources` on model 3 from `depends_on` when you sa "topic": "tenant_status", "description": "Daily tenant events from Backstage API", "model_type": "python", - "dags": ["opus_python_source"], + "dags": [""], "compute": "kpo", "kpo_size": "small", "depends_on": [], @@ -31,8 +31,8 @@ DJ auto-infers trino `upstream_sources` on model 3 from `depends_on` when you sa { "type": "external", "value": "api::backstage_tenant_events" } ], "output": { - "database": "glue_development", - "schema": "opus_python_source", + "database": "", + "schema": "python_output_schema", "table": "tenant_events", "partition_by": ["portal_partition_daily"], "write_mode": "overwrite_partitions" @@ -50,7 +50,7 @@ from python_models._config import PythonModelConfig, get_catalog_name, build_con OUTPUT_CONFIG = PythonModelConfig( model_name="tenant_events", - namespace="opus_python_source", + namespace="python_output_schema", table_name="tenant_events", model_type="python", description="Daily tenant events from Backstage API", @@ -96,8 +96,8 @@ if __name__ == "__main__": { "type": "external", "value": "s3::data-lake/tenant_migrations/{ds}/migrations.csv" } ], "output": { - "database": "glue_development", - "schema": "opus_python_source", + "database": "", + "schema": "python_output_schema", "table": "tenant_migrations", "partition_by": ["portal_partition_daily"], "write_mode": "overwrite_partitions" @@ -118,12 +118,12 @@ Pandas extract reads S3; `load()` uses `overwrite_partition_from_dataframe` the "description": "Combines tenant events and migrations into a daily status table", "depends_on": ["tenant_events", "tenant_migrations"], "upstream_sources": [ - { "type": "trino", "value": "glue_development.opus_python_source.tenant_events" }, - { "type": "trino", "value": "glue_development.opus_python_source.tenant_migrations" } + { "type": "trino", "value": ".python_output_schema.tenant_events" }, + { "type": "trino", "value": ".python_output_schema.tenant_migrations" } ], "output": { - "database": "glue_development", - "schema": "opus_python_source", + "database": "", + "schema": "python_output_schema", "table": "tenant_migration_status", "partition_by": ["portal_partition_daily"], "write_mode": "overwrite_partitions" @@ -144,14 +144,14 @@ from python_models._config import PythonModelConfig, get_catalog_name, build_con OUTPUT_CONFIG = PythonModelConfig( model_name="tenant_migration_status", - namespace="opus_python_source", + namespace="python_output_schema", table_name="tenant_migration_status", model_type="python", description="Combines tenant events and migrations into a daily status table", ) -EVENTS_FQN = f"{get_catalog_name()}.opus_python_source.tenant_events" -MIGRATIONS_FQN = f"{get_catalog_name()}.opus_python_source.tenant_migrations" +EVENTS_FQN = f"{get_catalog_name()}.python_output_schema.tenant_events" +MIGRATIONS_FQN = f"{get_catalog_name()}.python_output_schema.tenant_migrations" def extract(context: dict) -> dict[str, pd.DataFrame]: ds = context["ds"] @@ -233,7 +233,7 @@ After successful KPO runs, each output table carries: ```json { "python_model_name": "tenant_migration_status", - "python_model_upstream_sources": "[{\"type\":\"trino\",\"value\":\"glue_development.opus_python_source.tenant_events\"},{\"type\":\"trino\",\"value\":\"glue_development.opus_python_source.tenant_migrations\"}]" + "python_model_upstream_sources": "[{\"type\":\"trino\",\"value\":\".python_output_schema.tenant_events\"},{\"type\":\"trino\",\"value\":\".python_output_schema.tenant_migrations\"}]" } ``` @@ -251,7 +251,7 @@ Model Lineage renders: Register each output table, then reference in `.model.json`: ``` -glue_development__opus_python_source.tenant_migration_status +__python_output_schema.tenant_migration_status ``` Use **`dj-create-source`** after the first successful run. diff --git a/templates/skills/dj-migrate-notebook-to-pymodel/_SKILL.md b/templates/skills/dj-migrate-notebook-to-pymodel/_SKILL.md index 7b5dfc8..83e2ba6 100644 --- a/templates/skills/dj-migrate-notebook-to-pymodel/_SKILL.md +++ b/templates/skills/dj-migrate-notebook-to-pymodel/_SKILL.md @@ -48,6 +48,7 @@ Use when the user mentions: migrate/convert/port a notebook, convert an existing 5. State resolved defaults in the migration plan before proposing identity fields. 6. Never assume org-specific catalog, schema, or DAG names without confirmation. +<<<<<<< HEAD ## DJ CLI (preferred when DJ is running) When `.dj/bin/dj system.ping` succeeds, use `trino.columns` (and `trino.tables` when needed) to validate tables named in the migration plan against the live warehouse. For post-migration data parity, delegate execution to **`dj-verify-pymodel-parity`** (`query.execute`). Invocation patterns → `dj-cli`. @@ -56,6 +57,12 @@ When `.dj/bin/dj system.ping` succeeds, use `trino.columns` (and `trino.tables` Before Phase B, inspect existing models under `dags/python_models/` and follow **that** layout: +======= +## Authoring shape — match the workspace + +Before Phase B, inspect existing models under `dags/python_models/` and follow **that** layout: + +>>>>>>> ff72b7f (Update python model skills) | Workspace pattern | What Phase B produces | | ----------------- | --------------------- | | **Hand-written `.python.py`** (metadata-only `.python.json` + `def run_etl`) | Write `.python.json` **and** `.python.py`; mirror sibling models | diff --git a/templates/skills/dj-review-python-model/_SKILL.md b/templates/skills/dj-review-python-model/_SKILL.md index ba09551..98f68b7 100644 --- a/templates/skills/dj-review-python-model/_SKILL.md +++ b/templates/skills/dj-review-python-model/_SKILL.md @@ -29,6 +29,13 @@ Use this skill when the user mentions: review python model, audit python model, - Verifying the model's output *data* against a legacy/reference table (this skill audits code, not data) → `dj-verify-pymodel-parity` - Migrating a legacy notebook into a python model before it can be reviewed → `dj-migrate-notebook-to-pymodel` +## Step 0 — Resolve project defaults (before checklist) + +1. Read `.agents/project/skills/dj-review-python-model/project-defaults.md` if it exists. +2. Read `.agents/project/skills/dj-create-python-model/project-defaults.md` for shared catalog/schema/authoring-mode defaults. +3. If missing, discover from sibling models under `dags/python_models/` and project source registrations. +4. Apply project `authoring_mode` when interpreting F3/F7/F9 (hand-written models may have empty `cells` in JSON). + ## Workflow - [ ] **1. Resolve scope.** Default to the open `.python.json` in the editor if any. Otherwise ask which Python model to review. Accept a file path or model name (`python______`). @@ -89,7 +96,7 @@ DJ auto-infers trino `upstream_sources` from `depends_on` at create/sync — inf | Check | What to validate | | ----- | ---------------- | -| D1 | Output table uses standard catalog/schema convention (`glue_development.opus_python_source` or project equivalent) | +| D1 | Output table uses project catalog/schema convention from Step 0 (`.`) | | D2 | Partition column `portal_partition_daily` emitted in output (SQL or pandas) when using partition writes | | D3 | Column names follow `snake_case` | | D4 | No `SELECT *` in production SQL INSERT statements | diff --git a/templates/skills/dj-review-python-model/evals/evals.md b/templates/skills/dj-review-python-model/evals/evals.md index 59752cf..0a46f8e 100644 --- a/templates/skills/dj-review-python-model/evals/evals.md +++ b/templates/skills/dj-review-python-model/evals/evals.md @@ -44,7 +44,7 @@ The review treats `.python.py` as the **code source of truth** (hand-edits are e ## E7 — Legacy 2-part trino FQN -**Input:** `upstream_sources` uses `opus_python_source.tenant_events` (2-part). +**Input:** `upstream_sources` uses `.tenant_events` (2-part). **Expected:** **L3 WARNING (LOW)** — legacy format still parsed; recommend migrating to `catalog.schema.table`. diff --git a/templates/skills/dj-review-python-model/references/review-checklist.md b/templates/skills/dj-review-python-model/references/review-checklist.md index af5daf1..d253f26 100644 --- a/templates/skills/dj-review-python-model/references/review-checklist.md +++ b/templates/skills/dj-review-python-model/references/review-checklist.md @@ -83,7 +83,7 @@ from python_models._config import PythonModelConfig OUTPUT_CONFIG = PythonModelConfig( model_name="tenant_events", model_type="python", - namespace="opus_python_source", + namespace="python_output_schema", table_name="tenant_events", description="Daily tenant events from API", ) @@ -159,7 +159,7 @@ Authors do not need `set_extra_properties` calls in model code for lineage. ```json "upstream_sources": [ - { "type": "trino", "value": "glue_development.opus_python_source.tenant_events" }, + { "type": "trino", "value": ".python_output_schema.tenant_events" }, { "type": "external", "value": "api::backstage_tenant_events" } ] ``` @@ -173,13 +173,13 @@ Authors do not need `set_extra_properties` calls in model code for lineage. **Pass (canonical):** ```json -{ "type": "trino", "value": "glue_development.opus_python_source.tenant_events" } +{ "type": "trino", "value": ".python_output_schema.tenant_events" } ``` **Warning (legacy 2-part):** ```json -{ "type": "trino", "value": "opus_python_source.tenant_events" } +{ "type": "trino", "value": "python_output_schema.tenant_events" } ``` Still parsed with catalog from parent context — recommend migrating to 3-part FQN. @@ -219,9 +219,9 @@ Missing schema — cannot resolve table location. **Example finding:** ``` -Tables read in code: glue_development.opus_python_source.tenant_events, glue_development.analytics.user_sessions -upstream_sources trino entries: glue_development.opus_python_source.tenant_events -Missing: glue_development.analytics.user_sessions +Tables read in code: .python_output_schema.tenant_events, .analytics.user_sessions +upstream_sources trino entries: .python_output_schema.tenant_events +Missing: .analytics.user_sessions ``` ### L6: Accuracy — no stale entries @@ -247,8 +247,8 @@ Stale: old_events_table "name": "tenant_migration_status", "depends_on": ["tenant_events", "tenant_migrations"], "upstream_sources": [ - { "type": "trino", "value": "glue_development.opus_python_source.tenant_events" }, - { "type": "trino", "value": "glue_development.opus_python_source.tenant_migrations" } + { "type": "trino", "value": ".python_output_schema.tenant_events" }, + { "type": "trino", "value": ".python_output_schema.tenant_migrations" } ] } ``` @@ -283,7 +283,7 @@ If `.source.json` files reference this model's output, table name must match `OU ### D1: Standard catalog/schema -**Default:** `glue_development.opus_python_source.` +**Default:** `.python_output_schema.` Custom schema is acceptable when documented. Flag unknown namespaces. From e9faae7a6c7e93a96c46a056dd2fc26a84785672 Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Tue, 27 Oct 2026 14:39:21 +0530 Subject: [PATCH 08/11] chore(airflow): set xlarge KPO preset memory to 12Gi Align kpo_factory fallbacks and kpo_sizes.yml with the xlarge pod memory target. --- airflow/v2_10/kpo_factory.py | 2 +- airflow/v2_10/kpo_sizes.yml | 2 +- airflow/v2_7/kpo_factory.py | 2 +- airflow/v2_7/kpo_sizes.yml | 2 +- src/services/framework/utils/__tests__/kpo-sizes.test.ts | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/airflow/v2_10/kpo_factory.py b/airflow/v2_10/kpo_factory.py index 59157a9..2c06ace 100644 --- a/airflow/v2_10/kpo_factory.py +++ b/airflow/v2_10/kpo_factory.py @@ -37,7 +37,7 @@ "small": ("1", "2Gi"), "medium": ("2", "4Gi"), "large": ("2", "8Gi"), - "xlarge": ("2", "16Gi"), + "xlarge": ("2", "12Gi"), } _SIZES_YAML_PATH = Path(__file__).resolve().parent / "kpo_sizes.yml" diff --git a/airflow/v2_10/kpo_sizes.yml b/airflow/v2_10/kpo_sizes.yml index 868dd7f..3288a18 100644 --- a/airflow/v2_10/kpo_sizes.yml +++ b/airflow/v2_10/kpo_sizes.yml @@ -15,5 +15,5 @@ sizes: guidance: "Heavier transforms — wider schemas, multi-step pandas work, ~2–8 GB." xlarge: cpu: "2" - memory: "16Gi" + memory: "12Gi" guidance: "Large workloads — full-table scans, wide fin-style exports, ~8–16 GB." diff --git a/airflow/v2_7/kpo_factory.py b/airflow/v2_7/kpo_factory.py index 59157a9..2c06ace 100644 --- a/airflow/v2_7/kpo_factory.py +++ b/airflow/v2_7/kpo_factory.py @@ -37,7 +37,7 @@ "small": ("1", "2Gi"), "medium": ("2", "4Gi"), "large": ("2", "8Gi"), - "xlarge": ("2", "16Gi"), + "xlarge": ("2", "12Gi"), } _SIZES_YAML_PATH = Path(__file__).resolve().parent / "kpo_sizes.yml" diff --git a/airflow/v2_7/kpo_sizes.yml b/airflow/v2_7/kpo_sizes.yml index 868dd7f..3288a18 100644 --- a/airflow/v2_7/kpo_sizes.yml +++ b/airflow/v2_7/kpo_sizes.yml @@ -15,5 +15,5 @@ sizes: guidance: "Heavier transforms — wider schemas, multi-step pandas work, ~2–8 GB." xlarge: cpu: "2" - memory: "16Gi" + memory: "12Gi" guidance: "Large workloads — full-table scans, wide fin-style exports, ~8–16 GB." diff --git a/src/services/framework/utils/__tests__/kpo-sizes.test.ts b/src/services/framework/utils/__tests__/kpo-sizes.test.ts index 1419d50..a0b793d 100644 --- a/src/services/framework/utils/__tests__/kpo-sizes.test.ts +++ b/src/services/framework/utils/__tests__/kpo-sizes.test.ts @@ -18,6 +18,6 @@ describe('readKpoSizesTemplate', () => { }); expect(sizes.large?.cpu).toBe('2'); expect(sizes.large?.memory).toBe('8Gi'); - expect(sizes.xlarge?.memory).toBe('16Gi'); + expect(sizes.xlarge?.memory).toBe('12Gi'); }); }); From 7a62e4fe9cd9ffe7ad83e6b6fa8761fb444cb9d6 Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Tue, 27 Oct 2026 14:48:48 +0530 Subject: [PATCH 09/11] fix(extension): default python models to Airflow worker Omit compute in new .python.json unless KPO is chosen; document under 2.5.0 and bump extension version. --- CHANGELOG.md | 6 +++++ airflow/v2_10/etl_helper.py | 7 +++--- airflow/v2_7/etl_helper.py | 7 +++--- package-lock.json | 4 ++-- package.json | 2 +- schemas/python-model.schema.json | 3 +-- .../__tests__/python-model-roundtrip.test.ts | 13 +++++++++++ src/services/framework/index.ts | 14 +++++------ .../schema/types/python-model.schema.d.ts | 2 +- .../skills/dj-create-python-model/_SKILL.md | 10 ++++---- .../dj-create-python-model/evals/evals.md | 2 +- web/src/pages/PythonModelCreate.tsx | 23 ++++++++++--------- 12 files changed, 56 insertions(+), 37 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbd3a87..ef5a6b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Change Log +## 2.5.0 + +### Python models + +- **MWAA worker is the default compute.** New `.python.json` files omit `compute` (Create Python Model defaults to Airflow worker). Set `"compute": "kpo"` and optional `kpo_size` to run on Kubernetes; pod presets remain in `dags/_ext_/kpo_sizes.yml`. + ## 2.4.0 ### Agent / CLI diff --git a/airflow/v2_10/etl_helper.py b/airflow/v2_10/etl_helper.py index bd00e94..9e28b8b 100644 --- a/airflow/v2_10/etl_helper.py +++ b/airflow/v2_10/etl_helper.py @@ -568,9 +568,10 @@ def discover_models(dag_id: str | None = None) -> list[dict]: "depends_on": config.get("depends_on", []), "task_group": config.get("task_group"), "topic": config.get("topic", ""), - "compute": config.get("compute", "kpo"), - "kpo_size": config.get("kpo_size", "small"), + "compute": config.get("compute", "airflow"), } + if config.get("compute") == "kpo": + model["kpo_size"] = config.get("kpo_size", "small") if config.get("output"): model["table_name"] = config["output"].get("table", "") @@ -825,7 +826,7 @@ def resolve_dates_in(dates_in: str | list | dict, ds: str) -> list[str]: def _model_uses_kpo(model: dict) -> bool: - return model.get("compute", "kpo") != "airflow" + return model.get("compute") == "kpo" def _current_dag(): diff --git a/airflow/v2_7/etl_helper.py b/airflow/v2_7/etl_helper.py index bd00e94..9e28b8b 100644 --- a/airflow/v2_7/etl_helper.py +++ b/airflow/v2_7/etl_helper.py @@ -568,9 +568,10 @@ def discover_models(dag_id: str | None = None) -> list[dict]: "depends_on": config.get("depends_on", []), "task_group": config.get("task_group"), "topic": config.get("topic", ""), - "compute": config.get("compute", "kpo"), - "kpo_size": config.get("kpo_size", "small"), + "compute": config.get("compute", "airflow"), } + if config.get("compute") == "kpo": + model["kpo_size"] = config.get("kpo_size", "small") if config.get("output"): model["table_name"] = config["output"].get("table", "") @@ -825,7 +826,7 @@ def resolve_dates_in(dates_in: str | list | dict, ds: str) -> list[str]: def _model_uses_kpo(model: dict) -> bool: - return model.get("compute", "kpo") != "airflow" + return model.get("compute") == "kpo" def _current_dag(): diff --git a/package-lock.json b/package-lock.json index 86ff31b..608572e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "dj", - "version": "2.4.0", + "version": "2.5.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "dj", - "version": "2.4.0", + "version": "2.5.0", "license": "Apache-2.0", "workspaces": [ "web" diff --git a/package.json b/package.json index 9a62e93..532126d 100644 --- a/package.json +++ b/package.json @@ -10,7 +10,7 @@ "type": "git", "url": "https://github.com/Workday/dj.git" }, - "version": "2.4.0", + "version": "2.5.0", "workspaces": [ "web" ], diff --git a/schemas/python-model.schema.json b/schemas/python-model.schema.json index 39250ee..9da4946 100644 --- a/schemas/python-model.schema.json +++ b/schemas/python-model.schema.json @@ -44,8 +44,7 @@ "compute": { "type": "string", "enum": ["kpo", "airflow"], - "default": "kpo", - "description": "Execution backend: kpo (Kubernetes pod, default) or airflow (MWAA worker)" + "description": "Execution backend: omit for MWAA worker (airflow); set to kpo for Kubernetes pod" }, "kpo_size": { "type": "string", diff --git a/src/services/framework/__tests__/python-model-roundtrip.test.ts b/src/services/framework/__tests__/python-model-roundtrip.test.ts index a19b5ee..dea8d60 100644 --- a/src/services/framework/__tests__/python-model-roundtrip.test.ts +++ b/src/services/framework/__tests__/python-model-roundtrip.test.ts @@ -211,6 +211,19 @@ describe('.py -> notebook -> .py stability', () => { expect(CONFIG.kpo_size).toBe('medium'); }); + test('scaffold generation works when compute is omitted (Airflow default)', () => { + const config: PythonModelConfig = { + name: 'worker_model', + group: 'analytics', + topic: 'revenue', + model_type: 'python', + dags: ['full_source_etl'], + }; + const py = generatePythonModelScaffoldPy(config); + expect(py).toContain('worker_model'); + expect(config.compute).toBeUndefined(); + }); + test('scaffold includes __main__ entry with build_context_from_env', () => { const py = generatePythonModelScaffoldPy(CONFIG); expect(py).toContain('if __name__ == "__main__":'); diff --git a/src/services/framework/index.ts b/src/services/framework/index.ts index e20a4f3..62f7f01 100644 --- a/src/services/framework/index.ts +++ b/src/services/framework/index.ts @@ -440,10 +440,12 @@ export class Framework implements ApiEnabledService<'framework'> { const baseTags = tags && tags.length > 0 ? tags : ['python-model', group]; - const compute = reqCompute ?? 'kpo'; const kpoFields = - compute === 'kpo' || reqCompute === undefined - ? { kpo_size: reqKpoSize ?? 'small' } + reqCompute === 'kpo' + ? { + compute: 'kpo' as const, + kpo_size: reqKpoSize ?? 'small', + } : {}; const pythonModelConfig: PythonModelConfig = { @@ -452,7 +454,6 @@ export class Framework implements ApiEnabledService<'framework'> { topic, model_type, dags, - compute, ...kpoFields, ...(description && { description }), ...(enable_notebook && { enable_notebook }), @@ -473,10 +474,7 @@ export class Framework implements ApiEnabledService<'framework'> { `Python model '${pythonModelConfig.name}' already exists in ${group}/${topic}`, ); } catch (err) { - if ( - err instanceof Error && - err.message.includes('already exists') - ) { + if (err instanceof Error && err.message.includes('already exists')) { throw err; } } diff --git a/src/shared/schema/types/python-model.schema.d.ts b/src/shared/schema/types/python-model.schema.d.ts index 32e984d..8ebd84c 100644 --- a/src/shared/schema/types/python-model.schema.d.ts +++ b/src/shared/schema/types/python-model.schema.d.ts @@ -47,7 +47,7 @@ export interface SchemaPythonModel { */ depends_on?: string[]; /** - * Execution backend: kpo (Kubernetes pod, default) or airflow (MWAA worker) + * Execution backend: omit for MWAA worker (airflow); set to kpo for Kubernetes pod */ compute?: 'kpo' | 'airflow'; /** diff --git a/templates/skills/dj-create-python-model/_SKILL.md b/templates/skills/dj-create-python-model/_SKILL.md index 773f480..b46489a 100644 --- a/templates/skills/dj-create-python-model/_SKILL.md +++ b/templates/skills/dj-create-python-model/_SKILL.md @@ -116,8 +116,8 @@ Present defaults and confirm overrides: | `output.write_mode` | `overwrite_partitions` | Yes | | `output.partition_by` | `["portal_partition_daily"]` | Yes | | `namespace` | `python` | Yes | -| `compute` | `kpo` | Yes (`airflow` for MWAA worker) | -| `kpo_size` | `small` | Yes (when `compute` is `kpo`) | +| `compute` | omitted (MWAA worker) | Yes — set `kpo` for Kubernetes pod | +| `kpo_size` | `small` when `compute` is `kpo` | Yes (only with `compute: kpo`) | **Confirm the write target before proceeding.** `overwrite_partitions` is destructive for affected partitions. State resolved `output.database` / `output.schema` / `write_mode` and get explicit confirmation when not using defaults. Never point a new model at production without confirmation. @@ -162,11 +162,11 @@ Declare in `.python.json` only — DJ writes properties after each run; **do not ## KPO and execution requirements -New models default to **`compute: "kpo"`** (Kubernetes pod). Set **`compute: "airflow"`** when the KPO image pipeline is not ready. +New models default to the **MWAA worker** — omit `compute` in `.python.json`. Set **`compute: "kpo"`** when the KPO image pipeline is ready. | Field | Purpose | | ----- | ------- | -| `compute` | `kpo` (default) or `airflow` | +| `compute` | Omit for Airflow worker, or `kpo` for Kubernetes pod | | `kpo_size` | Pod preset when `compute` is `kpo`: `small`, `medium`, `large`, `xlarge` | **KPO script path:** The configured KPO runner image (from project defaults or `airflow_runner_cfg`) syncs `s3:///dags/python_models/` into the pod and runs `python //.python.py`. Set Airflow Variable **`python_models_s3_bucket`**; optional **`python_model_s3_region`**. Configure **`airflow_runner_cfg`** (image, tag, namespace, cluster). @@ -200,7 +200,7 @@ Model ID: `python______` - [ ] Step 3: Determine DAG assignment - [ ] Step 4: Confirm transform strategy (pandas, SQL, or mixed) - [ ] Step 5: Identify data sources -- [ ] Step 6: Confirm output config and compute (`kpo` default) +- [ ] Step 6: Confirm output config and compute (Airflow default; opt into `kpo` when needed) - [ ] Step 7: Set `depends_on`, `upstream_sources`, and optional fields - [ ] Step 8: Read `.dj/schemas/python-model.schema.json` - [ ] Step 9: Write `.python.json` at `dags/python_models///.python.json` diff --git a/templates/skills/dj-create-python-model/evals/evals.md b/templates/skills/dj-create-python-model/evals/evals.md index a91e68d..255f01a 100644 --- a/templates/skills/dj-create-python-model/evals/evals.md +++ b/templates/skills/dj-create-python-model/evals/evals.md @@ -51,7 +51,7 @@ Reference the correct file model throughout: **Prompt:** "Create a model that runs on a large KPO pod." **Expected:** -- `compute: "kpo"` (default) with `kpo_size: "large"`. +- When user asks for KPO: `compute: "kpo"` with `kpo_size: "large"`. When unspecified, omit `compute` (Airflow worker). - Mentions the required Airflow Variable `python_models_s3_bucket` and `airflow_runner_cfg`. - One Airflow task named after the model `name`. diff --git a/web/src/pages/PythonModelCreate.tsx b/web/src/pages/PythonModelCreate.tsx index f590782..9fed5cc 100644 --- a/web/src/pages/PythonModelCreate.tsx +++ b/web/src/pages/PythonModelCreate.tsx @@ -152,9 +152,14 @@ export function PythonModelCreate() { setIsSubmitting(true); setShowErrorDialog(false); setErrorMessage(''); + const { compute: formCompute, kpo_size, ...rest } = values; + const request: Values = + formCompute === 'kpo' + ? { ...rest, compute: 'kpo', kpo_size: kpo_size ?? 'small' } + : rest; const resp = await api.post({ type: 'framework-python-model-create', - request: values, + request, }); setSuccess(resp); void onClose(); @@ -192,8 +197,7 @@ export function PythonModelCreate() { setValue('projectName', _projects[0].name); } setValue('model_type', 'python'); - setValue('compute', 'kpo'); - setValue('kpo_size', 'small'); + setValue('compute', 'airflow'); setValue('description', ''); setValue('namespace', ''); setValue('table_name', ''); @@ -256,8 +260,7 @@ export function PythonModelCreate() { namespace: '', table_name: '', model_type: 'python', - compute: 'kpo', - kpo_size: 'small', + compute: 'airflow', dags: [], enable_notebook: true, tags: undefined, @@ -449,9 +452,7 @@ export function PythonModelCreate() { label="Select DAG" options={availableDags.map((d) => ({ label: d, value: d }))} disabled={isDagsLoading || !projectName} - placeholder={ - isDagsLoading ? 'Loading DAGs...' : 'Select a DAG' - } + placeholder={isDagsLoading ? 'Loading DAGs...' : 'Select a DAG'} tooltipText="The DAG where this model will run. The model will execute as part of the run_python_models task in the selected DAG." /> )} @@ -463,13 +464,13 @@ export function PythonModelCreate() { render={({ field }) => ( )} /> From 140067c0a765445c6bda512e3cbc97c736fa95cd Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Tue, 27 Oct 2026 15:36:56 +0530 Subject: [PATCH 10/11] Update dj_etl config keys --- CHANGELOG.md | 2 + airflow/v2_10/etl_helper.py | 6 +- airflow/v2_7/etl_helper.py | 6 +- docs/SETTINGS.md | 11 ++++ package.json | 8 +++ src/services/config.ts | 2 + .../__tests__/python-model-roundtrip.test.ts | 4 +- .../utils/__tests__/kpo-env-context.test.ts | 15 +++++ .../__tests__/python-model-config.test.ts | 33 ++++++++++ .../framework/utils/python-model-upstream.ts | 13 +++- .../framework/utils/python-model-utils.ts | 8 ++- src/services/sync/PythonModelSyncService.ts | 6 +- src/services/types/config.ts | 2 + .../python-model-catalog-defaults.test.ts | 65 +++++++++++++++++++ .../__tests__/python-model-lineage.test.ts | 26 ++++---- .../python-upstream-inference.test.ts | 16 ++--- .../__tests__/python-upstream-sources.test.ts | 26 ++++---- .../pymodel/python-model-catalog-defaults.ts | 52 +++++++++++++++ .../pymodel/python-upstream-inference.ts | 20 +++--- .../skills/dj-create-python-model/_SKILL.md | 6 +- .../skills/dj-review-python-model/_SKILL.md | 9 +-- .../references/review-checklist.md | 22 +++++-- 22 files changed, 287 insertions(+), 71 deletions(-) create mode 100644 src/shared/pymodel/__tests__/python-model-catalog-defaults.test.ts create mode 100644 src/shared/pymodel/python-model-catalog-defaults.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index c46a62a..ed586ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ ### Python models - **MWAA worker is the default compute.** New `.python.json` files omit `compute` (Create Python Model defaults to Airflow worker). Set `"compute": "kpo"` and optional `kpo_size` to run on Kubernetes; pod presets remain in `dags/_ext_/kpo_sizes.yml`. +- **Airflow-aligned defaults for `depends_on` upstream inference.** When DJ auto-fills trino `upstream_sources` from `depends_on`, catalog and schema resolve from `dj.pythonModel.catalogName` / `dj.pythonModel.databaseName` (mirrors Airflow Variables `dj_python_model_catalog_name` and `dj_python_model_database_name`), KPO env vars, then `glue` and `python_source` — not hardcoded org catalogs. +- **KPO and python models use the same Trino Variables as source ETL.** Regenerated `python_models/_trino_io.py` and KPO pod env read `dj_etl_trino_host`, `dj_etl_trino_user`, and `dj_etl_trino_password` instead of legacy `opus-etl-*` names. ### Model Wizard diff --git a/airflow/v2_10/etl_helper.py b/airflow/v2_10/etl_helper.py index 9e28b8b..132f4c1 100644 --- a/airflow/v2_10/etl_helper.py +++ b/airflow/v2_10/etl_helper.py @@ -909,9 +909,9 @@ def build_kpo_env_vars_from_context(ctx: dict) -> dict[str, str]: ), "python_models_s3_bucket": _airflow_var("python_models_s3_bucket"), "python_model_s3_region": _airflow_var("python_model_s3_region", "us-west-2"), - "OPUS-ETL-CLUSTER-HOST": _airflow_var("opus-etl-cluster-host"), - "OPUS-ETL-USER-NAME": _airflow_var("opus-etl-user-name"), - "OPUS-ETL-CLUSTER-PASSWORD": _airflow_var("opus-etl-cluster-password"), + "DJ_ETL_TRINO_HOST": _airflow_var("dj_etl_trino_host"), + "DJ_ETL_TRINO_USER": _airflow_var("dj_etl_trino_user"), + "DJ_ETL_TRINO_PASSWORD": _airflow_var("dj_etl_trino_password"), } if "dates" in ctx: env["PYMODEL_DATES"] = json.dumps(ctx["dates"]) diff --git a/airflow/v2_7/etl_helper.py b/airflow/v2_7/etl_helper.py index 9e28b8b..132f4c1 100644 --- a/airflow/v2_7/etl_helper.py +++ b/airflow/v2_7/etl_helper.py @@ -909,9 +909,9 @@ def build_kpo_env_vars_from_context(ctx: dict) -> dict[str, str]: ), "python_models_s3_bucket": _airflow_var("python_models_s3_bucket"), "python_model_s3_region": _airflow_var("python_model_s3_region", "us-west-2"), - "OPUS-ETL-CLUSTER-HOST": _airflow_var("opus-etl-cluster-host"), - "OPUS-ETL-USER-NAME": _airflow_var("opus-etl-user-name"), - "OPUS-ETL-CLUSTER-PASSWORD": _airflow_var("opus-etl-cluster-password"), + "DJ_ETL_TRINO_HOST": _airflow_var("dj_etl_trino_host"), + "DJ_ETL_TRINO_USER": _airflow_var("dj_etl_trino_user"), + "DJ_ETL_TRINO_PASSWORD": _airflow_var("dj_etl_trino_password"), } if "dates" in ctx: env["PYMODEL_DATES"] = json.dumps(ctx["dates"]) diff --git a/docs/SETTINGS.md b/docs/SETTINGS.md index 50cc297..72bb1fd 100644 --- a/docs/SETTINGS.md +++ b/docs/SETTINGS.md @@ -47,6 +47,17 @@ Complete guide to configuring the DJ (Data JSON) Framework VS Code extension. - Extension activates venv when running dbt and Python tools - Validated: checks for `bin/activate` file +**`dj.pythonModel.catalogName`** — Default Glue/Trino catalog when DJ infers `upstream_sources` from `depends_on` on `.python.json` sync. Mirrors Airflow Variable `dj_python_model_catalog_name` (runtime default when unset: `glue`). + +**`dj.pythonModel.databaseName`** — Default schema/namespace for the same inference. Mirrors Airflow Variable `dj_python_model_database_name`. + +```json +{ + "dj.pythonModel.catalogName": "glue", + "dj.pythonModel.databaseName": "python_source" +} +``` + **`dj.trinoPath`** - Path to Trino CLI executable (default: `"trino-cli"`) ```json diff --git a/package.json b/package.json index 821f66e..4ac112a 100644 --- a/package.json +++ b/package.json @@ -383,6 +383,14 @@ ], "description": "List of group categories available when creating Python models. Requires 'DJ: Refresh Projects' command or workspace reload to take effect." }, + "dj.pythonModel.catalogName": { + "type": "string", + "description": "Default Trino/Glue catalog for inferring python model upstream_sources from depends_on. Mirrors the Airflow Variable dj_python_model_catalog_name (same default as KPO: glue when unset)." + }, + "dj.pythonModel.databaseName": { + "type": "string", + "description": "Default schema/namespace for inferring python model upstream_sources from depends_on. Mirrors the Airflow Variable dj_python_model_database_name." + }, "dj.trinoPath": { "type": "string", "default": "trino-cli", diff --git a/src/services/config.ts b/src/services/config.ts index 3d02531..57147f1 100644 --- a/src/services/config.ts +++ b/src/services/config.ts @@ -64,6 +64,8 @@ export function getDjConfig(): CoderConfig { 'analytics', 'others', ]), + pythonModelCatalogName: config.get('pythonModel.catalogName', undefined), + pythonModelDatabaseName: config.get('pythonModel.databaseName', undefined), trinoPath: config.get('trinoPath', undefined), lightdashProjectPath: config.get('lightdashProjectPath', undefined), lightdashProfilesPath: config.get('lightdashProfilesPath', undefined), diff --git a/src/services/framework/__tests__/python-model-roundtrip.test.ts b/src/services/framework/__tests__/python-model-roundtrip.test.ts index dea8d60..f940a55 100644 --- a/src/services/framework/__tests__/python-model-roundtrip.test.ts +++ b/src/services/framework/__tests__/python-model-roundtrip.test.ts @@ -23,7 +23,7 @@ const CONFIG: PythonModelConfig = { topic: 'revenue', model_type: 'python', dags: ['full_source_etl'], - namespace: 'opus_python_source', + namespace: 'python_source', table_name: 'my_model', compute: 'kpo', kpo_size: 'medium', @@ -183,7 +183,7 @@ describe('.py -> notebook -> .py stability', () => { '', 'def transform_and_load(context):', ' ds = context["ds"]', - ' table = f"{get_catalog_name()}.opus_python_source.my_model"', + ' table = f"{get_catalog_name()}.python_source.my_model"', ' overwrite_partition(table, "portal_partition_daily", ds, insert_sql="")', '', '', diff --git a/src/services/framework/utils/__tests__/kpo-env-context.test.ts b/src/services/framework/utils/__tests__/kpo-env-context.test.ts index e0c5d90..2cf7032 100644 --- a/src/services/framework/utils/__tests__/kpo-env-context.test.ts +++ b/src/services/framework/utils/__tests__/kpo-env-context.test.ts @@ -4,7 +4,22 @@ */ import { describe, expect, test } from '@jest/globals'; +/** Trino keys injected by etl_helper.build_kpo_env_vars_from_context (must match _trino_io._var). */ +const KPO_TRINO_ENV_KEYS = [ + 'DJ_ETL_TRINO_HOST', + 'DJ_ETL_TRINO_USER', + 'DJ_ETL_TRINO_PASSWORD', +] as const; + describe('KPO env context contract', () => { + test('documents Trino env keys aligned with dj_etl_trino_* Airflow Variables', () => { + expect(KPO_TRINO_ENV_KEYS).toEqual([ + 'DJ_ETL_TRINO_HOST', + 'DJ_ETL_TRINO_USER', + 'DJ_ETL_TRINO_PASSWORD', + ]); + }); + test('resolved context maps to expected PYMODEL keys', () => { const ctx = { ds: '2025-08-25', diff --git a/src/services/framework/utils/__tests__/python-model-config.test.ts b/src/services/framework/utils/__tests__/python-model-config.test.ts index 9fc12a0..120c54e 100644 --- a/src/services/framework/utils/__tests__/python-model-config.test.ts +++ b/src/services/framework/utils/__tests__/python-model-config.test.ts @@ -29,4 +29,37 @@ describe('generateTrinoIoPy', () => { expect(py).toContain('def overwrite_partition_from_dataframe'); expect(py).toContain('Pandas I/O (python-path models only)'); }); + + test('resolves Trino connection from dj_etl_trino_* variables', () => { + const py = generateTrinoIoPy(); + expect(py).toContain('_var("dj_etl_trino_host")'); + expect(py).toContain('_var("dj_etl_trino_user")'); + expect(py).toContain('_var("dj_etl_trino_password")'); + expect(py).not.toContain('opus-etl-cluster-host'); + }); + + test('exports stable public _trino_io API surface', () => { + const py = generateTrinoIoPy(); + const publicDefs = [ + ...py.matchAll(/^def ([a-z][a-z0-9_]*)\(/gm), + ].map((m) => m[1]); + + expect(publicDefs).toEqual([ + 'get_trino_connection', + 'execute_sql', + 'fetch_value', + 'append', + 'overwrite', + 'overwrite_partition', + 'merge', + 'create_schema', + 'create_hive_csv_external_table', + 'create_iceberg_table', + 'set_extra_properties', + 'flush_hive_metadata_cache', + 'update', + 'read_partition_to_dataframe', + 'overwrite_partition_from_dataframe', + ]); + }); }); diff --git a/src/services/framework/utils/python-model-upstream.ts b/src/services/framework/utils/python-model-upstream.ts index 594a01b..447e402 100644 --- a/src/services/framework/utils/python-model-upstream.ts +++ b/src/services/framework/utils/python-model-upstream.ts @@ -1,7 +1,9 @@ +import { getDjConfig } from '@services/config'; import type { PythonModelConfig } from '@shared/framework/types'; import { DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS, enrichPythonModelUpstreamSources, + resolvePythonModelCatalogDefaults, type PythonModelCatalogDefaults, } from '@shared/pymodel/python-upstream-inference'; import * as fs from 'fs'; @@ -9,6 +11,15 @@ import * as path from 'path'; export { DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS }; +/** Catalog/namespace defaults for upstream inference (workspace settings, env, fallbacks). */ +export function getPythonModelCatalogDefaultsForWorkspace(): PythonModelCatalogDefaults { + const dj = getDjConfig(); + return resolvePythonModelCatalogDefaults({ + catalog: dj.pythonModelCatalogName, + namespace: dj.pythonModelDatabaseName, + }); +} + /** Load all python model configs in a topic directory keyed by model name. */ export function loadTopicPythonModelConfigs( topicDir: string, @@ -41,7 +52,7 @@ export function loadTopicPythonModelConfigs( export function applyUpstreamInference( config: PythonModelConfig, topicDir: string, - defaults: PythonModelCatalogDefaults = DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS, + defaults: PythonModelCatalogDefaults = getPythonModelCatalogDefaultsForWorkspace(), ): PythonModelConfig { const siblings = loadTopicPythonModelConfigs(topicDir); siblings.set(config.name, config); diff --git a/src/services/framework/utils/python-model-utils.ts b/src/services/framework/utils/python-model-utils.ts index 9ade40d..53d42d7 100644 --- a/src/services/framework/utils/python-model-utils.ts +++ b/src/services/framework/utils/python-model-utils.ts @@ -21,6 +21,7 @@ from __future__ import annotations from dataclasses import dataclass, field +# Must match PYTHON_MODEL_*_AIRFLOW_VAR in src/shared/pymodel/python-model-catalog-defaults.ts GLUE_CATALOG_VARIABLE = "dj_python_model_catalog_name" NAMESPACE_VARIABLE = "dj_python_model_database_name" @@ -194,12 +195,13 @@ def _var(name: str, default: str = "") -> str: def _resolve_trino_config() -> dict[str, Any]: - host = _var("opus-etl-cluster-host") + # Variable names match airflow variables.py and KPO env in etl_helper.build_kpo_env_vars_from_context + host = _var("dj_etl_trino_host") return { "host": host, "port": 443, - "user": _var("opus-etl-user-name"), - "password": _var("opus-etl-cluster-password"), + "user": _var("dj_etl_trino_user"), + "password": _var("dj_etl_trino_password"), } diff --git a/src/services/sync/PythonModelSyncService.ts b/src/services/sync/PythonModelSyncService.ts index ed9bdf2..b8ae883 100644 --- a/src/services/sync/PythonModelSyncService.ts +++ b/src/services/sync/PythonModelSyncService.ts @@ -30,7 +30,10 @@ import { generatePythonModelScaffoldPy, stripNotebookOnlyCells, } from '@services/framework/utils'; -import { loadTopicPythonModelConfigs } from '@services/framework/utils/python-model-upstream'; +import { + getPythonModelCatalogDefaultsForWorkspace, + loadTopicPythonModelConfigs, +} from '@services/framework/utils/python-model-upstream'; import { workspaceHasPythonModels } from '@services/framework/utils/source-etl-python-workspace'; import type { PythonModelCell, @@ -529,6 +532,7 @@ export class PythonModelSyncService implements vscode.Disposable { const enriched = enrichPythonModelUpstreamSources( fileSet.config, siblings, + getPythonModelCatalogDefaultsForWorkspace(), ); const before = JSON.stringify(fileSet.config.upstream_sources ?? null); const after = JSON.stringify(enriched.upstream_sources ?? null); diff --git a/src/services/types/config.ts b/src/services/types/config.ts index fe49633..f94a86a 100644 --- a/src/services/types/config.ts +++ b/src/services/types/config.ts @@ -23,6 +23,8 @@ export type CoderConfig = { dbtMacroPath?: string; pythonVenvPath?: string; pythonModelGroups: string[]; + pythonModelCatalogName?: string; + pythonModelDatabaseName?: string; trinoPath?: string; lightdashProjectPath?: string; lightdashProfilesPath?: string; diff --git a/src/shared/pymodel/__tests__/python-model-catalog-defaults.test.ts b/src/shared/pymodel/__tests__/python-model-catalog-defaults.test.ts new file mode 100644 index 0000000..e4240a6 --- /dev/null +++ b/src/shared/pymodel/__tests__/python-model-catalog-defaults.test.ts @@ -0,0 +1,65 @@ +import { afterEach, describe, expect, test } from '@jest/globals'; + +import { + PYTHON_MODEL_CATALOG_ENV, + PYTHON_MODEL_CATALOG_FALLBACK, + PYTHON_MODEL_DATABASE_ENV, + PYTHON_MODEL_DATABASE_FALLBACK, + resolvePythonModelCatalogDefaults, +} from '@shared/pymodel/python-model-catalog-defaults'; + +describe('resolvePythonModelCatalogDefaults', () => { + const savedCatalog = process.env[PYTHON_MODEL_CATALOG_ENV]; + const savedDatabase = process.env[PYTHON_MODEL_DATABASE_ENV]; + + afterEach(() => { + if (savedCatalog === undefined) { + delete process.env[PYTHON_MODEL_CATALOG_ENV]; + } else { + process.env[PYTHON_MODEL_CATALOG_ENV] = savedCatalog; + } + if (savedDatabase === undefined) { + delete process.env[PYTHON_MODEL_DATABASE_ENV]; + } else { + process.env[PYTHON_MODEL_DATABASE_ENV] = savedDatabase; + } + }); + + test('uses fallbacks when env and overrides are absent', () => { + delete process.env[PYTHON_MODEL_CATALOG_ENV]; + delete process.env[PYTHON_MODEL_DATABASE_ENV]; + expect(resolvePythonModelCatalogDefaults()).toEqual({ + catalog: PYTHON_MODEL_CATALOG_FALLBACK, + namespace: PYTHON_MODEL_DATABASE_FALLBACK, + }); + }); + + test('prefers explicit overrides over env', () => { + process.env[PYTHON_MODEL_CATALOG_ENV] = 'from_env'; + process.env[PYTHON_MODEL_DATABASE_ENV] = 'from_env_db'; + expect( + resolvePythonModelCatalogDefaults({ + catalog: 'custom_cat', + namespace: 'custom_ns', + }), + ).toEqual({ catalog: 'custom_cat', namespace: 'custom_ns' }); + }); + + test('reads KPO env vars when overrides omitted', () => { + process.env[PYTHON_MODEL_CATALOG_ENV] = 'prod_glue'; + process.env[PYTHON_MODEL_DATABASE_ENV] = 'my_schema'; + expect(resolvePythonModelCatalogDefaults()).toEqual({ + catalog: 'prod_glue', + namespace: 'my_schema', + }); + }); + + test('ignores empty or None env values', () => { + process.env[PYTHON_MODEL_CATALOG_ENV] = 'None'; + process.env[PYTHON_MODEL_DATABASE_ENV] = ' '; + expect(resolvePythonModelCatalogDefaults()).toEqual({ + catalog: PYTHON_MODEL_CATALOG_FALLBACK, + namespace: PYTHON_MODEL_DATABASE_FALLBACK, + }); + }); +}); diff --git a/src/shared/pymodel/__tests__/python-model-lineage.test.ts b/src/shared/pymodel/__tests__/python-model-lineage.test.ts index a80caf9..f98540e 100644 --- a/src/shared/pymodel/__tests__/python-model-lineage.test.ts +++ b/src/shared/pymodel/__tests__/python-model-lineage.test.ts @@ -16,11 +16,11 @@ describe('parsePythonModelProperties', () => { parsePythonModelProperties({ python_model_name: 'model_d', python_model_upstream_sources: - '[{"type":"trino","value":"opus_python_source.model_c"}]', + '[{"type":"trino","value":"python_source.model_c"}]', }), ).toEqual({ pythonModelName: 'model_d', - upstreamSources: [{ type: 'trino', value: 'opus_python_source.model_c' }], + upstreamSources: [{ type: 'trino', value: 'python_source.model_c' }], namespace: undefined, table: undefined, description: undefined, @@ -52,23 +52,23 @@ describe('buildPythonUpstreamGraph', () => { name: 'model_d', type: 'source', path: '', - schema: 'opus_python_source', + schema: 'python_source', database: 'glue', }; test('walks a multi-hop Python model chain D -> C -> B,A', async () => { const propsByTable: Record> = { - 'glue.opus_python_source.model_d': { + 'glue.python_source.model_d': { python_model_name: 'model_d', python_model_upstream_sources: - '[{"type":"trino","value":"opus_python_source.model_c"}]', + '[{"type":"trino","value":"python_source.model_c"}]', }, - 'glue.opus_python_source.model_c': { + 'glue.python_source.model_c': { python_model_name: 'model_c', python_model_upstream_sources: - '[{"type":"trino","value":"opus_python_source.model_b"},{"type":"external","value":"backstage_api"}]', + '[{"type":"trino","value":"python_source.model_b"},{"type":"external","value":"backstage_api"}]', }, - 'glue.opus_python_source.model_b': { + 'glue.python_source.model_b': { python_model_name: 'model_b', }, }; @@ -213,17 +213,17 @@ describe('buildPythonUpstreamGraph', () => { name: 'out', type: 'source', path: '', - schema: 'opus_python_source', - database: 'glue_development', + schema: 'python_source', + database: 'glue', }; const propsByTable: Record> = { - 'glue_development.opus_python_source.out': { + 'glue.python_source.out': { python_model_name: 'out_model', python_model_upstream_sources: - '[{"type":"trino","value":"glue_development.opus_python_source.upstream"}]', + '[{"type":"trino","value":"glue.python_source.upstream"}]', }, - 'glue_development.opus_python_source.upstream': { + 'glue.python_source.upstream': { python_model_name: 'upstream_model', }, }; diff --git a/src/shared/pymodel/__tests__/python-upstream-inference.test.ts b/src/shared/pymodel/__tests__/python-upstream-inference.test.ts index eb75c2c..39c6348 100644 --- a/src/shared/pymodel/__tests__/python-upstream-inference.test.ts +++ b/src/shared/pymodel/__tests__/python-upstream-inference.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from '@jest/globals'; +import { DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS } from '@shared/pymodel/python-model-catalog-defaults'; import { inferTrinoUpstreamFromDependsOn, mergeInferredUpstreamSources, @@ -7,10 +8,7 @@ import { } from '@shared/pymodel/python-upstream-inference'; import type { PythonModelConfig } from '@shared/framework/types'; -const defaults = { - catalog: 'glue_development', - namespace: 'opus_python_source', -}; +const defaults = DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS; describe('resolvePythonModelOutputFqn', () => { test('uses output block when present', () => { @@ -32,7 +30,7 @@ describe('resolvePythonModelOutputFqn', () => { test('falls back to defaults and model name', () => { expect( resolvePythonModelOutputFqn({ name: 'customer_tenant_dc_tracker' }, defaults), - ).toBe('glue_development.opus_python_source.customer_tenant_dc_tracker'); + ).toBe('glue.python_source.customer_tenant_dc_tracker'); }); }); @@ -73,12 +71,12 @@ describe('inferTrinoUpstreamFromDependsOn', () => { { type: 'trino', value: - 'glue_development.opus_python_source.customer_tenant_dc_tracker', + 'glue.python_source.customer_tenant_dc_tracker', }, { type: 'trino', value: - 'glue_development.opus_python_source.customer_tenant_migration_enrichment', + 'glue.python_source.customer_tenant_migration_enrichment', }, ]); }); @@ -94,7 +92,7 @@ describe('inferTrinoUpstreamFromDependsOn', () => { { type: 'trino', value: - 'glue_development.opus_python_source.customer_tenant_dc_tracker', + 'glue.python_source.customer_tenant_dc_tracker', }, ], }; @@ -105,7 +103,7 @@ describe('inferTrinoUpstreamFromDependsOn', () => { { type: 'trino', value: - 'glue_development.opus_python_source.customer_tenant_migration_enrichment', + 'glue.python_source.customer_tenant_migration_enrichment', }, ]); }); diff --git a/src/shared/pymodel/__tests__/python-upstream-sources.test.ts b/src/shared/pymodel/__tests__/python-upstream-sources.test.ts index 3a17d83..a0a4eb5 100644 --- a/src/shared/pymodel/__tests__/python-upstream-sources.test.ts +++ b/src/shared/pymodel/__tests__/python-upstream-sources.test.ts @@ -42,11 +42,11 @@ describe('serializeUpstreamSourcesForProperty', () => { test('serializes as compact JSON array', () => { expect( serializeUpstreamSourcesForProperty([ - { type: 'trino', value: 'opus_python_source.raw_events' }, + { type: 'trino', value: 'python_source.raw_events' }, { type: 'external', value: 'backstage_api' }, ]), ).toBe( - '[{"type":"trino","value":"opus_python_source.raw_events"},{"type":"external","value":"backstage_api"}]', + '[{"type":"trino","value":"python_source.raw_events"},{"type":"external","value":"backstage_api"}]', ); }); }); @@ -66,10 +66,10 @@ describe('parseUpstreamSourcesProperty', () => { test('falls back to legacy comma-separated trino entries', () => { expect( parseUpstreamSourcesProperty( - 'opus_python_source.raw_events, analytics.user_sessions', + 'python_source.raw_events, analytics.user_sessions', ), ).toEqual([ - { type: 'trino', value: 'opus_python_source.raw_events' }, + { type: 'trino', value: 'python_source.raw_events' }, { type: 'trino', value: 'analytics.user_sessions' }, ]); }); @@ -84,11 +84,11 @@ describe('parseUpstreamSourcesProperty', () => { // property can surface as a real array rather than a string. expect( parseUpstreamSourcesProperty([ - { type: 'trino', value: 'glue_development.opus_python_sources.wpc_slavepool_details' }, + { type: 'trino', value: 'glue.python_sources.wpc_slavepool_details' }, { type: 'external', value: 'jenkins::dcim_inventory_csv' }, ]), ).toEqual([ - { type: 'trino', value: 'glue_development.opus_python_sources.wpc_slavepool_details' }, + { type: 'trino', value: 'glue.python_sources.wpc_slavepool_details' }, { type: 'external', value: 'jenkins::dcim_inventory_csv' }, ]); }); @@ -97,10 +97,10 @@ describe('parseUpstreamSourcesProperty', () => { expect( parseUpstreamSourcesProperty({ type: 'trino', - value: 'glue_development.opus_python_sources.wpc_slavepool_details', + value: 'glue.python_sources.wpc_slavepool_details', }), ).toEqual([ - { type: 'trino', value: 'glue_development.opus_python_sources.wpc_slavepool_details' }, + { type: 'trino', value: 'glue.python_sources.wpc_slavepool_details' }, ]); }); @@ -114,21 +114,21 @@ describe('parseTrinoUpstreamValue', () => { test('parses 3-part catalog.schema.table', () => { expect( parseTrinoUpstreamValue( - 'glue_development.opus_python_source.model_c', + 'glue.python_source.model_c', 'glue', ), ).toEqual({ - catalog: 'glue_development', - schema: 'opus_python_source', + catalog: 'glue', + schema: 'python_source', table: 'model_c', }); }); test('parses 2-part schema.table with fallback catalog', () => { - expect(parseTrinoUpstreamValue('opus_python_source.model_c', 'glue')).toEqual( + expect(parseTrinoUpstreamValue('python_source.model_c', 'glue')).toEqual( { catalog: 'glue', - schema: 'opus_python_source', + schema: 'python_source', table: 'model_c', }, ); diff --git a/src/shared/pymodel/python-model-catalog-defaults.ts b/src/shared/pymodel/python-model-catalog-defaults.ts new file mode 100644 index 0000000..3a26c6c --- /dev/null +++ b/src/shared/pymodel/python-model-catalog-defaults.ts @@ -0,0 +1,52 @@ +export interface PythonModelCatalogDefaults { + catalog: string; + namespace: string; +} + +/** Airflow Variable for Trino/Glue catalog (see generated `python_models/_config.py`). */ +export const PYTHON_MODEL_CATALOG_AIRFLOW_VAR = 'dj_python_model_catalog_name'; + +/** Airflow Variable for output schema/namespace. */ +export const PYTHON_MODEL_DATABASE_AIRFLOW_VAR = 'dj_python_model_database_name'; + +/** KPO pod env key set from {@link PYTHON_MODEL_CATALOG_AIRFLOW_VAR} in etl_helper. */ +export const PYTHON_MODEL_CATALOG_ENV = 'DJ_PYTHON_MODEL_CATALOG_NAME'; + +/** Optional pod env mirror for {@link PYTHON_MODEL_DATABASE_AIRFLOW_VAR}. */ +export const PYTHON_MODEL_DATABASE_ENV = 'DJ_PYTHON_MODEL_DATABASE_NAME'; + +/** Matches `build_kpo_env_vars_from_context` default in etl_helper.py. */ +export const PYTHON_MODEL_CATALOG_FALLBACK = 'glue'; + +export const PYTHON_MODEL_DATABASE_FALLBACK = 'python_source'; + +function nonEmpty(value: string | undefined): string | undefined { + const trimmed = value?.trim(); + return trimmed && trimmed !== 'None' ? trimmed : undefined; +} + +/** + * Resolve catalog/namespace defaults for upstream FQN inference. + * Precedence: explicit overrides → KPO env vars → fallbacks. + */ +export function resolvePythonModelCatalogDefaults( + overrides?: Partial, +): PythonModelCatalogDefaults { + const catalog = + nonEmpty(overrides?.catalog) ?? + nonEmpty(process.env[PYTHON_MODEL_CATALOG_ENV]) ?? + PYTHON_MODEL_CATALOG_FALLBACK; + + const namespace = + nonEmpty(overrides?.namespace) ?? + nonEmpty(process.env[PYTHON_MODEL_DATABASE_ENV]) ?? + PYTHON_MODEL_DATABASE_FALLBACK; + + return { catalog, namespace }; +} + +/** Static fallbacks only (stable for tests); prefer {@link resolvePythonModelCatalogDefaults}. */ +export const DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS: PythonModelCatalogDefaults = { + catalog: PYTHON_MODEL_CATALOG_FALLBACK, + namespace: PYTHON_MODEL_DATABASE_FALLBACK, +}; diff --git a/src/shared/pymodel/python-upstream-inference.ts b/src/shared/pymodel/python-upstream-inference.ts index d7b0596..09fdb85 100644 --- a/src/shared/pymodel/python-upstream-inference.ts +++ b/src/shared/pymodel/python-upstream-inference.ts @@ -3,12 +3,17 @@ import type { PythonModelUpstreamSource, } from '@shared/framework/types'; +import { + resolvePythonModelCatalogDefaults, + type PythonModelCatalogDefaults, +} from './python-model-catalog-defaults'; import { normalizeUpstreamSources } from './python-upstream-sources'; -export interface PythonModelCatalogDefaults { - catalog: string; - namespace: string; -} +export type { PythonModelCatalogDefaults } from './python-model-catalog-defaults'; +export { + DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS, + resolvePythonModelCatalogDefaults, +} from './python-model-catalog-defaults'; /** Resolve catalog.schema.table for a python model output table. */ export function resolvePythonModelOutputFqn( @@ -79,16 +84,11 @@ export function inferTrinoUpstreamFromDependsOn( return inferred; } -export const DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS: PythonModelCatalogDefaults = { - catalog: 'glue_development', - namespace: 'opus_python_source', -}; - /** Apply depends_on inference and return a config with merged upstream_sources. */ export function enrichPythonModelUpstreamSources( config: PythonModelConfig, siblingConfigs: Map, - defaults: PythonModelCatalogDefaults = DEFAULT_PYTHON_MODEL_CATALOG_DEFAULTS, + defaults: PythonModelCatalogDefaults = resolvePythonModelCatalogDefaults(), ): PythonModelConfig { const inferred = inferTrinoUpstreamFromDependsOn( config, diff --git a/templates/skills/dj-create-python-model/_SKILL.md b/templates/skills/dj-create-python-model/_SKILL.md index b46489a..6b36bbf 100644 --- a/templates/skills/dj-create-python-model/_SKILL.md +++ b/templates/skills/dj-create-python-model/_SKILL.md @@ -173,7 +173,9 @@ New models default to the **MWAA worker** — omit `compute` in `.python.json`. **`PYMODEL_*` environment:** Airflow sets `PYMODEL_DS`, `PYMODEL_DS_NODASH`, and optionally `PYMODEL_DATES` (JSON array) before execution. `build_context_from_env()` in `_config.py` reads them and configures `logging.basicConfig` to stdout (required for KPO log visibility). -**`__main__` block (required):** Scaffolded `.python.py` must end with: +**Trino connection (`_trino_io.py`):** Uses the same Airflow Variables as source ETL — `dj_etl_trino_host`, `dj_etl_trino_user`, `dj_etl_trino_password`. KPO runs may also set env `DJ_ETL_TRINO_HOST`, `DJ_ETL_TRINO_USER`, `DJ_ETL_TRINO_PASSWORD` (from `etl_helper`). Local VS Code upstream inference can mirror catalog/schema via `dj.pythonModel.catalogName` / `dj.pythonModel.databaseName`. + +**`__main__` block (default scaffold):** New `.python.py` files from DJ should end with: ```python if __name__ == "__main__": @@ -181,7 +183,7 @@ if __name__ == "__main__": run_etl(build_context_from_env()) ``` -MWAA worker and KPO pod both execute the file as `__main__`. `run_etl` is scaffold convention, not a framework-enforced name. +MWAA worker and KPO `python ` execution use pattern above. Projects with a custom `entrypoint.py` that imports models and calls `run_etl(context)` may omit `__main__` if every model defines `run_etl` and has no bare top-level runner call. `run_etl` is scaffold convention, not a framework-enforced name. **Airflow UI:** Each KPO model is **one task** named after the model `name` field. diff --git a/templates/skills/dj-review-python-model/_SKILL.md b/templates/skills/dj-review-python-model/_SKILL.md index 98f68b7..30940f3 100644 --- a/templates/skills/dj-review-python-model/_SKILL.md +++ b/templates/skills/dj-review-python-model/_SKILL.md @@ -54,11 +54,11 @@ Use this skill when the user mentions: review python model, audit python model, | F1 | `.python.json` has required fields: `name`, `group`, `topic` | | F2 | Name/group/topic match pattern `^[a-z][a-z0-9_]*$` | | F3 | `.python.json` is valid with required metadata; `cells` is an optional non-authoritative snapshot (may be empty when the model is authored via `.python.py` directly) | -| F4 | `.python.py` contains `if __name__ == "__main__":` calling `run_etl(build_context_from_env())` | +| F4 | Script entry: **`__main__`** + `run_etl(build_context_from_env())` (MWAA worker / KPO `python` path) **or** **`def run_etl(context)`** only when a documented import/entrypoint orchestrator calls it (no bare top-level `run_etl(...)`) | | F5 | Uses `_trino_io` helpers (`from python_models._trino_io import …`) — no inline Trino connection code (`trino.dbapi.connect`, `create_engine`, raw HTTP to Trino) | | F6 | `OUTPUT_CONFIG` uses `PythonModelConfig` from `python_models._config` | | F7 | `.python.py` is present and treated as the code source of truth — hand-edited code is expected and preserved (DJ never overwrites it from JSON) | -| F8 | `.python.py` is import-safe — no side-effectful top-level statement such as a bare `run_etl(context)`; execution runs through the `__main__` block | +| F8 | `.python.py` is import-safe — no side-effectful top-level statement such as a bare `run_etl(context)` at module scope | | F9 | ETL follows a coherent orchestrator: pandas (`extract` / `transform` / `load`) **or** SQL (`extract` / `transform_and_load` / `cleanup`) **or** documented mixed pattern | ### 2. Lineage Readiness (L) @@ -178,7 +178,7 @@ DJ auto-infers trino `upstream_sources` from `depends_on` at create/sync — inf | Priority | Criteria | | -------- | -------- | -| CRITICAL | Lineage broken (L2/L5/L7 failures), missing `__main__`, no partition column | +| CRITICAL | Lineage broken (L2/L5/L7 failures), missing script entry (F4), no partition column | | HIGH | Inline Trino code, missing cleanup, stale `upstream_sources`, missing `depends_on` for Python upstream | | MEDIUM | Orchestrator inconsistency, missing batch staging, no error handling | | LOW | Legacy 2-part trino FQN, naming style, documentation gaps | @@ -208,7 +208,8 @@ DJ auto-infers trino `upstream_sources` from `depends_on` at create/sync — inf - **`_trino_io` SQL path:** `execute_sql`, `fetch_value`, `overwrite_partition`, `overwrite`, `append`, `merge`, `update` - **`_trino_io` pandas path:** `read_partition_to_dataframe`, `overwrite_partition_from_dataframe` - **Do not use:** `execute_trino`, inline `trino.dbapi.connect`, standalone `delete()` -- **Script entry:** `if __name__ == "__main__":` with `build_context_from_env()` from `_config.py` +- **Trino connection (generated `_trino_io.py`):** Airflow Variables `dj_etl_trino_host`, `dj_etl_trino_user`, `dj_etl_trino_password` (same as source ETL). KPO may prefetch `DJ_ETL_TRINO_HOST`, `DJ_ETL_TRINO_USER`, `DJ_ETL_TRINO_PASSWORD` into the pod env. +- **Script entry:** `if __name__ == "__main__":` with `build_context_from_env()` from `_config.py` for MWAA worker / KPO `python ` — or `def run_etl(context)` when a project entrypoint imports the module and invokes it - **`upstream_sources` trino format:** `catalog.schema.table` (canonical); `schema.table` legacy - **`upstream_sources` external format:** `s3::…`, `api::…`, `jenkins::…` - **Partition contract:** `portal_partition_daily` in partition-based models diff --git a/templates/skills/dj-review-python-model/references/review-checklist.md b/templates/skills/dj-review-python-model/references/review-checklist.md index d253f26..70df683 100644 --- a/templates/skills/dj-review-python-model/references/review-checklist.md +++ b/templates/skills/dj-review-python-model/references/review-checklist.md @@ -33,9 +33,11 @@ Extended explanations, code examples, and edge cases for each check in the DJ Py `.python.py` is the source of truth for code. The `cells` array in `.python.json` is an optional, non-authoritative snapshot of the last notebook save — it may be empty when the model is authored by editing `.python.py` directly, which is fine and not a failure. -### F4: `__main__` script entry +### F4: Script entry (`__main__` or importable `run_etl`) -**Pass condition:** `.python.py` contains: +**Pass condition (either pattern):** + +**A — DJ default (MWAA worker / KPO `python .python.py`):** `.python.py` contains: ```python if __name__ == "__main__": @@ -43,7 +45,9 @@ if __name__ == "__main__": run_etl(build_context_from_env()) ``` -**Why it matters:** Airflow worker and KPO pods execute the script as `__main__` via `python ` — not import + callable lookup. `build_context_from_env()` reads `PYMODEL_DS`, `PYMODEL_DS_NODASH`, `PYMODEL_DATES` and configures stdout logging for KPO. +**B — Custom KPO / batch entrypoint (`dags/python_models/entrypoint.py` or similar):** module defines `def run_etl(context: dict) -> None:` (or equivalent orchestrator) and has **no** bare top-level `run_etl(...)` call. The orchestrator loads the file with `importlib` and calls `run_etl(context)` with a built context dict. + +**Why it matters:** The MWAA worker path uses `runpy.run_path(..., run_name="__main__")`, so pattern A is required there. Pattern B is valid when deployment docs state that execution is import-based, not `python` as `__main__`. `build_context_from_env()` reads `PYMODEL_DS`, `PYMODEL_DS_NODASH`, `PYMODEL_DATES` for pattern A and for KPO env parity. ### F5: `_trino_io` usage @@ -99,19 +103,23 @@ OUTPUT_CONFIG = PythonModelConfig( ### F8: Import safety -**Pass condition:** `.python.py` has no side-effectful top-level statements — e.g. no bare `run_etl(context)` at module scope (which would execute on import). Execution goes through the `__main__` block. +**Pass condition:** `.python.py` has no side-effectful top-level statements — e.g. no bare `run_etl(context)` at module scope (which would execute on import). ```python -# FAIL — runs on import +# FAIL — runs on import (breaks entrypoint import path too) run_etl(context) -# PASS — script entry only +# PASS — MWAA / KPO script path (F4 pattern A) if __name__ == "__main__": from python_models._config import build_context_from_env run_etl(build_context_from_env()) + +# PASS — entrypoint path (F4 pattern B): define only, orchestrator calls run_etl +def run_etl(context: dict) -> None: + ... ``` -**Why it matters:** Airflow workers import the module and KPO runs `python ` — a top-level runner call would fire at import time. Only the derived `# Python Model: …` header is notebook-only; all other code in the notebook view is real module code. +**Why it matters:** Custom entrypoints import the module before calling `run_etl`. A top-level runner call fires at import time. MWAA worker `__main__` execution still requires pattern A in F4 when that path is used. ### F9: Orchestrator structure From d84e93a4f7ac9280aac2b868d108d0780b6a6672 Mon Sep 17 00:00:00 2001 From: Ajayjayendran Arulraj Date: Tue, 27 Oct 2026 15:39:20 +0530 Subject: [PATCH 11/11] Update FieldSelectSingle.tsx --- web/src/forms/FieldSelectSingle.tsx | 1 - 1 file changed, 1 deletion(-) diff --git a/web/src/forms/FieldSelectSingle.tsx b/web/src/forms/FieldSelectSingle.tsx index 3d8fe72..68212b3 100644 --- a/web/src/forms/FieldSelectSingle.tsx +++ b/web/src/forms/FieldSelectSingle.tsx @@ -13,7 +13,6 @@ export type FieldSelectSingleProps = ControllerRenderProps & { labelClass?: string; helpIcon?: React.ReactNode; // Help icon for Assist Me disabled?: boolean; - placeholder?: string; /** * Virtualize the options list. Enable for large option sets (hundreds+), * e.g. Trino tables, so the dropdown doesn't render every row into the DOM