From 9cbdccb40354a13d9557aef28f3e5dc4c02b8801 Mon Sep 17 00:00:00 2001 From: mo Date: Fri, 26 Jun 2026 08:41:44 +0000 Subject: [PATCH] infra: fix mongodb uuid encoding + combined DAG honors rows conf --- infra/airflow/generate_data_dag.py | 75 ++++++++++++++++ infra/airflow/generate_mongodb_events_data.py | 90 +++++++++++++++++++ .../scripts/generate_mongodb_events_data.py | 2 +- 3 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 infra/airflow/generate_data_dag.py create mode 100644 infra/airflow/generate_mongodb_events_data.py diff --git a/infra/airflow/generate_data_dag.py b/infra/airflow/generate_data_dag.py new file mode 100644 index 0000000..e239754 --- /dev/null +++ b/infra/airflow/generate_data_dag.py @@ -0,0 +1,75 @@ +"""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 diff --git a/infra/airflow/generate_mongodb_events_data.py b/infra/airflow/generate_mongodb_events_data.py new file mode 100644 index 0000000..adc5d76 --- /dev/null +++ b/infra/airflow/generate_mongodb_events_data.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +""" +Script to generate fake event data for MongoDB +Generates approximately 1GB of data +""" + +import pymongo +import random +from datetime import datetime, timedelta +import uuid +import sys + +# Database connection details +DB_HOST = "10.0.21.51" +DB_PORT = "27017" +DB_NAME = "supplychain" +COLLECTION_NAME = "events" + +# Data generation settings +import os +TARGET_DOCUMENTS = int(os.getenv("GEN_ROWS", "5000")) # light, configurable via GEN_ROWS +BATCH_SIZE = min(5000, max(500, TARGET_DOCUMENTS)) + +# Sample data +EVENT_TYPES = ["INSERT", "UPDATE", "DELETE", "CREATE", "MODIFY"] +REGIONS = ["EU", "APAC", "LATAM", "NA", "EMEA"] +SOURCES = ["ERP", "WMS", "CRM", "SCM", "TMS"] + +def generate_fake_event(): + """Generate a single fake event""" + event_id = str(uuid.uuid4()) + event_type = random.choice(EVENT_TYPES) + region = random.choice(REGIONS) + source = random.choice(SOURCES) + + # Random timestamp within the last year + days_ago = random.randint(0, 365) + ts = datetime.now() - timedelta(days=days_ago, hours=random.randint(0, 23), + minutes=random.randint(0, 59)) + + amount = random.uniform(100.0, 50000.0) + + # Generate a long payload field (like the existing data) + payload = "X" * 500 + + return { + "event_id": event_id, + "type": event_type, + "region": region, + "source": source, + "amount": amount, + "ts": ts, + "payload": payload + } + +def main(): + print(f"Connecting to MongoDB at {DB_HOST}:{DB_PORT}...") + + client = pymongo.MongoClient(f"mongodb://{DB_HOST}:{DB_PORT}/") + db = client[DB_NAME] + collection = db[COLLECTION_NAME] + + print(f"Generating {TARGET_DOCUMENTS} events...") + print(f"Batch size: {BATCH_SIZE}") + + total_generated = 0 + batch = [] + + for i in range(TARGET_DOCUMENTS): + batch.append(generate_fake_event()) + + if len(batch) >= BATCH_SIZE: + collection.insert_many(batch) + total_generated += len(batch) + batch = [] + + if total_generated % 100000 == 0: + print(f"Generated {total_generated} documents...") + + # Insert remaining documents + if batch: + collection.insert_many(batch) + total_generated += len(batch) + + client.close() + + print(f"Completed! Generated {total_generated} events.") + +if __name__ == "__main__": + main() diff --git a/infra/airflow/scripts/generate_mongodb_events_data.py b/infra/airflow/scripts/generate_mongodb_events_data.py index 2be9fa8..adc5d76 100644 --- a/infra/airflow/scripts/generate_mongodb_events_data.py +++ b/infra/airflow/scripts/generate_mongodb_events_data.py @@ -28,7 +28,7 @@ SOURCES = ["ERP", "WMS", "CRM", "SCM", "TMS"] def generate_fake_event(): """Generate a single fake event""" - event_id = uuid.uuid4() + event_id = str(uuid.uuid4()) event_type = random.choice(EVENT_TYPES) region = random.choice(REGIONS) source = random.choice(SOURCES)