Phase 5: Neo4j Graph Database Integration¶

import asyncio
from neo4j import AsyncGraphDatabase
from netbox_client import NetBoxAPI

class GraphPopulationEngine:
    def __init__(self, neo4j_uri: str, netbox_client: NetBoxAPI):
        self.driver = AsyncGraphDatabase.driver(neo4j_uri, auth=('neo4j', 'password'))
        self.netbox = netbox_client

    async def populate_initial_graph(self):
        """
        Parallel extraction and graph population from NetBox
        """
        # Extract all entity types concurrently
        extraction_tasks = [
            self.extract_sites(),
            self.extract_devices(), 
            self.extract_interfaces(),
            self.extract_cables(),
            self.extract_vlans(),
            self.extract_ip_addresses(),
            self.extract_power_infrastructure()
        ]

        extracted_data = await asyncio.gather(*extraction_tasks)

        # Build graph in dependency order
        async with self.driver.session() as session:
            await session.execute_write(self.create_sites, extracted_data[0])
            await session.execute_write(self.create_devices, extracted_data[1])
            await session.execute_write(self.create_interfaces, extracted_data[2])
            await session.execute_write(self.create_relationships, extracted_data)

        print("Initial graph population completed")

    async def create_devices(self, tx, devices_data):
        """
        Batch create device nodes with properties
        """
        cypher = """
        UNWIND $devices as device
        CREATE (d:Device {
            id: device.id,
            name: device.name,
            device_type: device.device_type.model,
            serial: device.serial,
            status: device.status.value,
            site_id: device.site.id,
            rack_id: device.rack.id
        })
        """
        await tx.run(cypher, devices=devices_data)

    async def create_relationships(self, tx, all_data):
        """
        Create relationships between entities in batch operations
        """
        # Site -> Device relationships
        site_device_cypher = """
        MATCH (s:Site {id: $site_id})
        MATCH (d:Device {site_id: $site_id})
        CREATE (s)-[:CONTAINS]->(d)
        """

        # Interface -> Device relationships  
        interface_device_cypher = """
        MATCH (d:Device {id: $device_id})
        MATCH (i:Interface {device_id: $device_id})
        CREATE (d)-[:HAS_INTERFACE]->(i)
        """

        # Cable connections (most complex)
        cable_connection_cypher = """
        MATCH (c:Cable {id: $cable_id})
        MATCH (i1:Interface {id: $termination_a_id})
        MATCH (i2:Interface {id: $termination_b_id})
        CREATE (i1)-[:CONNECTED_VIA]->(c)<-[:CONNECTED_VIA]-(i2)
        """

        # Execute relationship creation in parallel where possible
        relationship_tasks = [
            tx.run(site_device_cypher, site_id=site_id)
            for site_id in extract_site_ids(all_data)
        ]

        await asyncio.gather(*relationship_tasks)

from fastapi import FastAPI, BackgroundTasks
from neo4j import AsyncGraphDatabase
import json

app = FastAPI()

class GraphSyncEngine:
    def __init__(self, neo4j_uri: str):
        self.driver = AsyncGraphDatabase.driver(neo4j_uri)

    async def handle_webhook_event(self, event_type: str, object_type: str, 
                                 data: dict, snapshot: dict):
        """
        Process NetBox webhook events and sync to graph
        """
        handler_map = {
            ('created', 'device'): self.handle_device_created,
            ('updated', 'device'): self.handle_device_updated, 
            ('deleted', 'device'): self.handle_device_deleted,
            ('created', 'cable'): self.handle_cable_created,
            ('deleted', 'cable'): self.handle_cable_deleted,
        }

        handler = handler_map.get((event_type, object_type))
        if handler:
            await handler(data, snapshot)
        else:
            print(f"Unhandled webhook: {event_type} {object_type}")

    async def handle_device_created(self, data: dict, snapshot: dict):
        """
        Add new device to graph with relationships
        """
        async with self.driver.session() as session:
            # Create device node
            await session.execute_write(self.create_device_node, data)

            # Create site relationship if site exists
            if data.get('site'):
                await session.execute_write(self.link_device_to_site, 
                                          data['id'], data['site']['id'])

            # Create rack relationship if rack assigned
            if data.get('rack'):
                await session.execute_write(self.link_device_to_rack,
                                          data['id'], data['rack']['id'])

    async def handle_cable_created(self, data: dict, snapshot: dict):
        """
        Add cable with termination relationships
        """
        async with self.driver.session() as session:
            cypher = """
            CREATE (c:Cable {
                id: $cable_id,
                label: $label,
                type: $cable_type,
                length: $length,
                status: $status
            })

            WITH c
            MATCH (i1:Interface {id: $term_a_id})
            MATCH (i2:Interface {id: $term_b_id})
            CREATE (i1)-[:CONNECTED_VIA]->(c)<-[:CONNECTED_VIA]-(i2)
            """

            await session.execute_write(
                lambda tx: tx.run(cypher, 
                    cable_id=data['id'],
                    label=data.get('label', ''),
                    cable_type=data.get('type', {}).get('name', ''),
                    length=data.get('length'),
                    status=data.get('status', {}).get('value', ''),
                    term_a_id=data['a_terminations'][0]['object_id'],
                    term_b_id=data['b_terminations'][0]['object_id']
                )
            )

sync_engine = GraphSyncEngine("bolt://localhost:7687")

@app.post("/netbox-webhook")
async def netbox_webhook(webhook_data: dict, background_tasks: BackgroundTasks):
    """
    Receive NetBox webhook and queue graph sync
    """
    background_tasks.add_task(
        sync_engine.handle_webhook_event,
        webhook_data['event'],
        webhook_data['model'],  
        webhook_data['data'],
        webhook_data['snapshots']
    )

    return {"status": "accepted"}

import schedule
import time
from datetime import datetime, timedelta

class GraphConsistencyChecker:
    def __init__(self, netbox_client, neo4j_driver):
        self.netbox = netbox_client
        self.driver = neo4j_driver

    async def daily_consistency_check(self):
        """
        Compare NetBox and Neo4j data for inconsistencies
        """
        print(f"Starting consistency check at {datetime.now()}")

        inconsistencies = []

        # Check device count consistency
        netbox_device_count = len(await self.netbox.dcim.devices.all())

        async with self.driver.session() as session:
            neo4j_result = await session.run("MATCH (d:Device) RETURN count(d) as device_count")
            neo4j_device_count = (await neo4j_result.single())['device_count']

        if netbox_device_count != neo4j_device_count:
            inconsistencies.append({
                'type': 'device_count_mismatch',
                'netbox_count': netbox_device_count,
                'neo4j_count': neo4j_device_count,
                'severity': 'high'
            })

        # Check for orphaned relationships
        async with self.driver.session() as session:
            orphaned_query = """
            MATCH (c:Cable)-[:CONNECTED_VIA]-(i:Interface)
            WHERE NOT EXISTS {
                MATCH (i)<-[:HAS_INTERFACE]-(:Device)
            }
            RETURN count(DISTINCT c) as orphaned_cables
            """
            orphaned_result = await session.run(orphaned_query)
            orphaned_count = (await orphaned_result.single())['orphaned_cables']

        if orphaned_count > 0:
            inconsistencies.append({
                'type': 'orphaned_cable_relationships',
                'count': orphaned_count,
                'severity': 'medium'
            })

        if inconsistencies:
            await self.handle_inconsistencies(inconsistencies)
        else:
            print("Graph consistency check passed")

    async def handle_inconsistencies(self, inconsistencies):
        """
        Handle detected inconsistencies with automated repair
        """
        for issue in inconsistencies:
            if issue['severity'] == 'high':
                print(f"CRITICAL: {issue}")
                await self.trigger_full_resync()
            else:
                print(f"WARNING: {issue}")
                await self.repair_specific_issue(issue)

# Schedule daily consistency checks
schedule.every().day.at("02:00").do(lambda: asyncio.run(
    GraphConsistencyChecker(netbox_client, neo4j_driver).daily_consistency_check()
))

class HybridQueryRouter:
    def __init__(self, neo4j_driver, netbox_client):
        self.graph_db = neo4j_driver
        self.netbox = netbox_client
        self.performance_tracker = QueryPerformanceTracker()

    async def route_query(self, classified_intent: Dict[str, Any]) -> ExecutionStrategy:
        """
        Decide whether to use graph database or direct API calls
        """
        query_characteristics = self.analyze_query_characteristics(classified_intent)

        # Graph database is optimal for:
        if (query_characteristics['relationship_depth'] > 2 or
            query_characteristics['involves_path_finding'] or
            query_characteristics['requires_aggregation'] or
            query_characteristics['involves_multiple_object_types'] > 3):

            # Check graph data freshness
            graph_freshness = await self.check_graph_freshness(classified_intent['target_objects'])

            if graph_freshness['is_acceptable']:
                return ExecutionStrategy(
                    primary='neo4j_graph',
                    fallback='netbox_api',
                    estimated_performance='sub_second'
                )
            else:
                return ExecutionStrategy(
                    primary='hybrid_refresh',
                    steps=['refresh_graph_subset', 'execute_graph_query'],
                    estimated_performance='fast'
                )

        # Direct API is better for:
        elif (query_characteristics['single_object_lookup'] or
              query_characteristics['requires_realtime_data'] or
              query_characteristics['simple_list_query']):

            return ExecutionStrategy(
                primary='netbox_api',
                optimization='use_cache',
                estimated_performance='fast'
            )

        # Hybrid approach for complex edge cases
        else:
            return ExecutionStrategy(
                primary='hybrid_combination',
                steps=['graph_analysis', 'api_enrichment', 'combine_results'],
                estimated_performance='moderate'
            )

    async def execute_graph_query(self, cypher_query: str, parameters: dict) -> dict:
        """
        Execute Cypher query with performance monitoring
        """
        start_time = time.time()

        async with self.graph_db.session() as session:
            result = await session.run(cypher_query, parameters)
            data = await result.data()

        execution_time = time.time() - start_time

        self.performance_tracker.record_query_performance(
            query_type='graph',
            execution_time=execution_time,
            result_count=len(data)
        )

        return {
            'data': data,
            'execution_time': execution_time,
            'source': 'neo4j_graph',
            'cache_eligible': True
        }