infra: mirror light per-source Airflow DAGs + generators (GEN_ROWS), fix script path

2026-06-26 01:07:23 +00:00
parent 9400f9410c
commit 62b416d210
6 changed files with 690 additions and 0 deletions
@@ -0,0 +1,228 @@
+#!/usr/bin/env python3
+"""
+Script to generate fake graph data for Neo4j
+Generates approximately 1GB of data with nodes and relationships
+"""
+
+from neo4j import GraphDatabase
+import random
+import uuid
+import sys
+
+# Database connection details
+DB_HOST = "10.0.21.51"
+DB_PORT = "7687"
+DB_USER = "neo4j"
+DB_PASSWORD = "testpwd"
+
+# Data generation settings
+import os
+TARGET_NODES = int(os.getenv("GEN_ROWS", "2000"))  # light, configurable via GEN_ROWS
+BATCH_SIZE = min(1000, max(200, TARGET_NODES))
+
+# Sample data
+PRODUCT_CATEGORIES = ["Electronics", "Clothing", "Food", "Furniture", "Toys", "Books"]
+SUPPLIER_REGIONS = ["EU", "APAC", "LATAM", "NA", "EMEA"]
+RELATIONSHIP_TYPES = ["SUPPLIES", "RELATED_TO", "COMPATIBLE_WITH", "PART_OF"]
+
+def generate_fake_product():
+    """Generate a single fake product node"""
+    product_id = str(uuid.uuid4())
+    name = f"Product-{random.randint(1000, 999999)}"
+    category = random.choice(PRODUCT_CATEGORIES)
+    price = round(random.uniform(10.0, 1000.0), 2)
+    stock = random.randint(0, 1000)
+    
+    # Generate a long description field
+    description = "X" * 200
+    
+    return {
+        "product_id": product_id,
+        "name": name,
+        "category": category,
+        "price": price,
+        "stock": stock,
+        "description": description
+    }
+
+def generate_fake_supplier():
+    """Generate a single fake supplier node"""
+    supplier_id = str(uuid.uuid4())
+    name = f"Supplier-{random.randint(1000, 999999)}"
+    region = random.choice(SUPPLIER_REGIONS)
+    rating = round(random.uniform(1.0, 5.0), 1)
+    
+    # Generate a long address field
+    address = "X" * 150
+    
+    return {
+        "supplier_id": supplier_id,
+        "name": name,
+        "region": region,
+        "rating": rating,
+        "address": address
+    }
+
+def main():
+    print(f"Connecting to Neo4j at {DB_HOST}:{DB_PORT}...")
+    
+    driver = GraphDatabase.driver(f"bolt://{DB_HOST}:{DB_PORT}", 
+                                   auth=(DB_USER, DB_PASSWORD))
+    
+    with driver.session() as session:
+        print(f"Generating {TARGET_NODES} product nodes...")
+        print(f"Batch size: {BATCH_SIZE}")
+        
+        total_products = 0
+        total_suppliers = 0
+        product_ids = []
+        
+        # Generate product nodes
+        batch = []
+        for i in range(TARGET_NODES):
+            product = generate_fake_product()
+            batch.append(product)
+            product_ids.append(product["product_id"])
+            
+            if len(batch) >= BATCH_SIZE:
+                session.run(
+                    """
+                    UNWIND $batch as row
+                    CREATE (p:Product {
+                        product_id: row.product_id,
+                        name: row.name,
+                        category: row.category,
+                        price: row.price,
+                        stock: row.stock,
+                        description: row.description
+                    })
+                    """,
+                    batch=batch
+                )
+                total_products += len(batch)
+                batch = []
+                
+                if total_products % 50000 == 0:
+                    print(f"Generated {total_products} product nodes...")
+        
+        # Insert remaining products
+        if batch:
+            session.run(
+                """
+                UNWIND $batch as row
+                CREATE (p:Product {
+                    product_id: row.product_id,
+                    name: row.name,
+                    category: row.category,
+                    price: row.price,
+                    stock: row.stock,
+                    description: row.description
+                })
+                """,
+                batch=batch
+            )
+            total_products += len(batch)
+        
+        print(f"Generated {total_products} product nodes.")
+        
+        # Generate supplier nodes (fewer than products)
+        print(f"Generating supplier nodes...")
+        target_suppliers = max(20, TARGET_NODES // 10)
+        batch = []
+        supplier_ids = []
+        
+        for i in range(target_suppliers):
+            supplier = generate_fake_supplier()
+            batch.append(supplier)
+            supplier_ids.append(supplier["supplier_id"])
+            
+            if len(batch) >= BATCH_SIZE:
+                session.run(
+                    """
+                    UNWIND $batch as row
+                    CREATE (s:Supplier {
+                        supplier_id: row.supplier_id,
+                        name: row.name,
+                        region: row.region,
+                        rating: row.rating,
+                        address: row.address
+                    })
+                    """,
+                    batch=batch
+                )
+                total_suppliers += len(batch)
+                batch = []
+        
+        if batch:
+            session.run(
+                """
+                UNWIND $batch as row
+                CREATE (s:Supplier {
+                    supplier_id: row.supplier_id,
+                    name: row.name,
+                    region: row.region,
+                    rating: row.rating,
+                    address: row.address
+                })
+                """,
+                batch=batch
+            )
+            total_suppliers += len(batch)
+        
+        print(f"Generated {total_suppliers} supplier nodes.")
+        
+        # Create relationships between products and suppliers
+        print(f"Creating relationships...")
+        batch = []
+        total_relationships = 0
+        
+        for product_id in product_ids:
+            # Each product is supplied by 1-3 random suppliers
+            num_suppliers = random.randint(1, 3)
+            for _ in range(num_suppliers):
+                supplier_id = random.choice(supplier_ids)
+                rel_type = random.choice(RELATIONSHIP_TYPES)
+                
+                batch.append({
+                    "product_id": product_id,
+                    "supplier_id": supplier_id,
+                    "rel_type": rel_type
+                })
+                
+                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
+                    )
+                    total_relationships += len(batch)
+                    batch = []
+                    
+                    if total_relationships % 50000 == 0:
+                        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
+            )
+            total_relationships += len(batch)
+        
+        print(f"Created {total_relationships} relationships.")
+    
+    driver.close()
+    print(f"Completed! Generated {total_products} products, {total_suppliers} suppliers, and {total_relationships} relationships.")
+
+if __name__ == "__main__":
+    main()