diff --git a/infra/airflow/scripts/generate_cassandra_telemetry_data.py b/infra/airflow/scripts/generate_cassandra_telemetry_data.py index e61ddf3..bd7ae4e 100644 --- a/infra/airflow/scripts/generate_cassandra_telemetry_data.py +++ b/infra/airflow/scripts/generate_cassandra_telemetry_data.py @@ -51,37 +51,18 @@ def main(): print(f"Generating {TARGET_ROWS} device metrics...") print(f"Batch size: {BATCH_SIZE}") + insert_stmt = session.prepare( + f"INSERT INTO {KEYSPACE}.{TABLE_NAME} (device_id, metric_ts, metric_type, metric_value, payload) " + f"VALUES (?, ?, ?, ?, ?)" + ) + total_generated = 0 - batch = [] - for i in range(TARGET_ROWS): - batch.append(generate_fake_device_metric()) - - if len(batch) >= BATCH_SIZE: - session.execute( - f""" - INSERT INTO {KEYSPACE}.{TABLE_NAME} (device_id, metric_ts, metric_type, metric_value, payload) - VALUES (%s, %s, %s, %s, %s) - """, - batch - ) - total_generated += len(batch) - batch = [] - - if total_generated % 100000 == 0: - print(f"Generated {total_generated} rows...") - - # Insert remaining rows - if batch: - session.execute( - f""" - INSERT INTO {KEYSPACE}.{TABLE_NAME} (device_id, metric_ts, metric_type, metric_value, payload) - VALUES (%s, %s, %s, %s, %s) - """, - batch - ) - total_generated += len(batch) - + session.execute(insert_stmt, generate_fake_device_metric()) + total_generated += 1 + if total_generated % 1000 == 0: + print(f"Generated {total_generated} rows...") + session.shutdown() cluster.shutdown() diff --git a/infra/airflow/scripts/generate_neo4j_graph_data.py b/infra/airflow/scripts/generate_neo4j_graph_data.py index ef6f7b4..0faafeb 100644 --- a/infra/airflow/scripts/generate_neo4j_graph_data.py +++ b/infra/airflow/scripts/generate_neo4j_graph_data.py @@ -25,6 +25,24 @@ PRODUCT_CATEGORIES = ["Electronics", "Clothing", "Food", "Furniture", "Toys", "B SUPPLIER_REGIONS = ["EU", "APAC", "LATAM", "NA", "EMEA"] RELATIONSHIP_TYPES = ["SUPPLIES", "RELATED_TO", "COMPATIBLE_WITH", "PART_OF"] +def _flush_rels(session, rels): + """Create relationships without APOC by grouping on the (fixed) rel type.""" + from collections import defaultdict + groups = defaultdict(list) + for r in rels: + groups[r["rel_type"]].append(r) + for rtype, items in groups.items(): + session.run( + f""" + UNWIND $batch as row + MATCH (p:Product {{product_id: row.product_id}}) + MATCH (s:Supplier {{supplier_id: row.supplier_id}}) + MERGE (p)-[:`{rtype}`]->(s) + """, + batch=items, + ) + + def generate_fake_product(): """Generate a single fake product node""" product_id = str(uuid.uuid4()) @@ -190,16 +208,7 @@ def main(): }) if len(batch) >= BATCH_SIZE: - session.run( - """ - UNWIND $batch as row - MATCH (p:Product {product_id: row.product_id}) - MATCH (s:Supplier {supplier_id: row.supplier_id}) - CALL apoc.create.relationship(p, row.rel_type, {}, s) YIELD rel - RETURN rel - """, - batch=batch - ) + _flush_rels(session, batch) total_relationships += len(batch) batch = [] @@ -207,16 +216,7 @@ def main(): print(f"Created {total_relationships} relationships...") if batch: - session.run( - """ - UNWIND $batch as row - MATCH (p:Product {product_id: row.product_id}) - MATCH (s:Supplier {supplier_id: row.supplier_id}) - CALL apoc.create.relationship(p, row.rel_type, {}, s) YIELD rel - RETURN rel - """, - batch=batch - ) + _flush_rels(session, batch) total_relationships += len(batch) print(f"Created {total_relationships} relationships.")