mo
248 lines
8.2 KiB
Python
248 lines
8.2 KiB
Python
#!/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
|
|
import os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
try:
|
|
from cdc_emit import emit_changes
|
|
except Exception:
|
|
def emit_changes(*a, **k):
|
|
return 0
|
|
|
|
NODES_TOPIC = "neo4j_graph.cdc.nodes"
|
|
RELS_TOPIC = "neo4j_graph.cdc.relationships"
|
|
|
|
# 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 _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())
|
|
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:
|
|
# Indexes make relationship MATCH/MERGE fast even with many existing nodes
|
|
session.run("CREATE INDEX product_id_idx IF NOT EXISTS FOR (p:Product) ON (p.product_id)")
|
|
session.run("CREATE INDEX supplier_id_idx IF NOT EXISTS FOR (s:Supplier) ON (s.supplier_id)")
|
|
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)
|
|
emit_changes(NODES_TOPIC, "graph", "Product", 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)
|
|
emit_changes(NODES_TOPIC, "graph", "Product", 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)
|
|
emit_changes(NODES_TOPIC, "graph", "Supplier", 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)
|
|
emit_changes(NODES_TOPIC, "graph", "Supplier", 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:
|
|
_flush_rels(session, batch)
|
|
total_relationships += len(batch)
|
|
emit_changes(RELS_TOPIC, "graph", "RELATIONSHIP", batch)
|
|
batch = []
|
|
|
|
if total_relationships % 50000 == 0:
|
|
print(f"Created {total_relationships} relationships...")
|
|
|
|
if batch:
|
|
_flush_rels(session, batch)
|
|
total_relationships += len(batch)
|
|
emit_changes(RELS_TOPIC, "graph", "RELATIONSHIP", 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()
|