mo
66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
"""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
|