infra: fix mongodb uuid encoding + combined DAG honors rows conf
This commit is contained in:
@@ -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
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user