mo
76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
"""Generate fake data for all databases in one run.
|
|
|
|
Runs all five generator scripts in parallel. Honors a `rows` value passed via
|
|
the dag_run conf (from the Command Center), exported to each script as
|
|
GEN_ROWS so the "All sources" button respects the row count.
|
|
"""
|
|
from airflow import DAG
|
|
from airflow.operators.python import PythonOperator
|
|
from datetime import datetime, timedelta
|
|
import os
|
|
import subprocess
|
|
|
|
default_args = {
|
|
'owner': 'airflow',
|
|
'depends_on_past': False,
|
|
'start_date': datetime(2025, 1, 1),
|
|
'email_on_failure': False,
|
|
'email_on_retry': False,
|
|
'retries': 1,
|
|
'retry_delay': timedelta(minutes=2),
|
|
}
|
|
|
|
dag = DAG(
|
|
'generate_data_all_databases',
|
|
default_args=default_args,
|
|
description='Generate fake data for all databases',
|
|
schedule=None,
|
|
catchup=False,
|
|
tags=['data', 'generation', 'fake-data'],
|
|
)
|
|
|
|
SCRIPTS_DIR = '/opt/airflow/dags/scripts'
|
|
|
|
SCRIPTS = {
|
|
'postgres': 'generate_postgres_sales_data.py',
|
|
'mongodb': 'generate_mongodb_events_data.py',
|
|
'mysql': 'generate_mysql_employee_data.py',
|
|
'neo4j': 'generate_neo4j_graph_data.py',
|
|
'cassandra': 'generate_cassandra_telemetry_data.py',
|
|
}
|
|
|
|
DEFAULT_ROWS = {'neo4j': '2000'}
|
|
|
|
|
|
def make_runner(script, default_rows):
|
|
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
|
|
|
|
|
|
tasks = []
|
|
for _src, _script in SCRIPTS.items():
|
|
tasks.append(PythonOperator(
|
|
task_id=f"generate_{_src}_data",
|
|
python_callable=make_runner(_script, DEFAULT_ROWS.get(_src, '5000')),
|
|
dag=dag,
|
|
))
|
|
|
|
# all in parallel
|
|
tasks
|