infra: mirror light per-source Airflow DAGs + generators (GEN_ROWS), fix script path

2026-06-26 01:07:23 +00:00
parent 9400f9410c
commit 62b416d210
6 changed files with 690 additions and 0 deletions
@@ -0,0 +1,65 @@
+"""Per-database light data generation DAGs.
+
+Each DAG runs one generator script and accepts a `rows` value via the
+dag_run conf (passed by the Command Center), exported as GEN_ROWS.
+Triggerable independently so every database gets its own button.
+"""
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+from datetime import datetime
+import os
+import subprocess
+
+SCRIPTS_DIR = "/opt/airflow/dags/scripts"
+
+SOURCES = {
+    "postgres": "generate_postgres_sales_data.py",
+    "mysql": "generate_mysql_employee_data.py",
+    "mongodb": "generate_mongodb_events_data.py",
+    "cassandra": "generate_cassandra_telemetry_data.py",
+    "neo4j": "generate_neo4j_graph_data.py",
+}
+
+DEFAULT_ROWS = {"neo4j": "2000"}
+
+default_args = {"owner": "airflow", "retries": 0}
+
+
+def make_runner(script: str, default_rows: str):
+    def _run(**context):
+        dag_run = context.get("dag_run")
+        conf = (dag_run.conf if dag_run else {}) or {}
+        rows = str(conf.get("rows") or default_rows)
+        env = dict(os.environ)
+        env["GEN_ROWS"] = rows
+        print(f"Running {script} with GEN_ROWS={rows}")
+        result = subprocess.run(
+            ["python3", os.path.join(SCRIPTS_DIR, script)],
+            capture_output=True, text=True, env=env,
+        )
+        if result.stdout:
+            print(result.stdout[-4000:])
+        if result.stderr:
+            print("STDERR:", result.stderr[-4000:])
+        if result.returncode != 0:
+            raise Exception(f"{script} failed with return code {result.returncode}")
+    return _run
+
+
+for _src, _script in SOURCES.items():
+    _dag_id = f"gen_{_src}"
+    _dag = DAG(
+        dag_id=_dag_id,
+        default_args=default_args,
+        description=f"Generate light data into {_src} (GEN_ROWS via conf.rows)",
+        schedule=None,
+        start_date=datetime(2025, 1, 1),
+        catchup=False,
+        tags=["data", "generation", _src],
+    )
+    PythonOperator(
+        task_id=f"generate_{_src}",
+        python_callable=make_runner(_script, DEFAULT_ROWS.get(_src, "5000")),
+        dag=_dag,
+    )
+    globals()[_dag_id] = _dag