Lakehouse/config/airflow/dags/scripts/generate_neo4j_graph_data.py

#!/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
TARGET_NODES = 500000  # Approximately 1GB of data with relationships
BATCH_SIZE = 1000

# 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 = 10000
        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()