From bbeb7ef4b6784e57288ec83ace75e2f63c0d28b0 Mon Sep 17 00:00:00 2001 From: dextmorgn Date: Thu, 6 Nov 2025 15:31:16 +0100 Subject: [PATCH] refactor(core): follow best practive principles --- .../src/flowsint_core/core/graph_db.py | 203 ++++++++++--- .../flowsint_core/core/graph_repository.py | 267 ++++++++++++++++++ .../flowsint_core/core/graph_serializer.py | 169 +++++++++++ .../src/flowsint_core/core/graph_service.py | 217 ++++++++++++++ .../src/flowsint_core/core/transform_base.py | 181 ++++++------ 5 files changed, 900 insertions(+), 137 deletions(-) create mode 100644 flowsint-core/src/flowsint_core/core/graph_repository.py create mode 100644 flowsint-core/src/flowsint_core/core/graph_serializer.py create mode 100644 flowsint-core/src/flowsint_core/core/graph_service.py diff --git a/flowsint-core/src/flowsint_core/core/graph_db.py b/flowsint-core/src/flowsint_core/core/graph_db.py index 7b28362e..0d7e200c 100644 --- a/flowsint-core/src/flowsint_core/core/graph_db.py +++ b/flowsint-core/src/flowsint_core/core/graph_db.py @@ -1,56 +1,177 @@ import os -from neo4j import GraphDatabase +from threading import Lock +from typing import Optional, Dict, Any, List +from neo4j import GraphDatabase, Driver, Session from dotenv import load_dotenv load_dotenv() class Neo4jConnection: - def __init__(self, uri: str, user: str, password: str): - self.__uri = uri - self.__user = user - self.__password = password - self.__driver = GraphDatabase.driver( - self.__uri, auth=(self.__user, self.__password) - ) + """ + Singleton Neo4j connection manager with proper connection pooling. - def query(self, query: str, parameters: dict = None): - with self.__driver.session() as session: - result = session.run(query, parameters) + This class implements the Singleton pattern to ensure only one Neo4j driver + instance exists throughout the application lifecycle, providing efficient + connection pooling and resource management. + """ + + _instance: Optional['Neo4jConnection'] = None + _lock: Lock = Lock() + _driver: Optional[Driver] = None + + def __new__(cls, uri: str = None, user: str = None, password: str = None): + """ + Create or return the singleton instance. + + Thread-safe singleton implementation using double-checked locking. + """ + if cls._instance is None: + with cls._lock: + if cls._instance is None: + instance = super().__new__(cls) + cls._instance = instance + return cls._instance + + def __init__(self, uri: str = None, user: str = None, password: str = None): + """ + Initialize the Neo4j connection (only once). + + Args: + uri: Neo4j connection URI + user: Neo4j username + password: Neo4j password + """ + # Only initialize once + if self._driver is None: + self._uri = uri or os.getenv("NEO4J_URI_BOLT") + self._user = user or os.getenv("NEO4J_USERNAME") + self._password = password or os.getenv("NEO4J_PASSWORD") + + if not all([self._uri, self._user, self._password]): + raise ValueError("Neo4j connection credentials are required") + + self._driver = GraphDatabase.driver( + self._uri, + auth=(self._user, self._password), + max_connection_pool_size=50, + connection_acquisition_timeout=60.0 + ) + + @classmethod + def get_instance(cls) -> 'Neo4jConnection': + """ + Get the singleton instance. + + Returns: + The singleton Neo4jConnection instance + """ + if cls._instance is None: + cls._instance = cls() + return cls._instance + + def get_driver(self) -> Driver: + """ + Get the Neo4j driver instance. + + Returns: + Neo4j Driver instance + """ + return self._driver + + def query(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]: + """ + Execute a single query. + + Args: + query: Cypher query string + parameters: Query parameters + + Returns: + List of result records as dictionaries + """ + with self._driver.session() as session: + result = session.run(query, parameters or {}) return result.data() + def execute_write(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]: + """ + Execute a write query within a write transaction. -URI = os.getenv("NEO4J_URI_BOLT") -USERNAME = os.getenv("NEO4J_USERNAME") -PASSWORD = os.getenv("NEO4J_PASSWORD") + Args: + query: Cypher query string + parameters: Query parameters -neo4j_connection = Neo4jConnection(URI, USERNAME, PASSWORD) + Returns: + List of result records as dictionaries + """ + def _execute(tx): + result = tx.run(query, parameters or {}) + return result.data() + + with self._driver.session() as session: + return session.execute_write(_execute) + + def execute_batch(self, queries: List[tuple[str, Dict[str, Any]]]) -> None: + """ + Execute multiple queries in a single transaction. + + Args: + queries: List of (query, parameters) tuples + """ + def _execute_batch(tx): + for query, params in queries: + tx.run(query, params or {}) + + with self._driver.session() as session: + session.execute_write(_execute_batch) + + def close(self) -> None: + """Close the driver connection.""" + if self._driver: + self._driver.close() + self._driver = None + + def verify_connectivity(self) -> bool: + """ + Verify the connection to Neo4j. + + Returns: + True if connection is successful, False otherwise + """ + try: + with self._driver.session() as session: + result = session.run("RETURN 1") + return result.single()[0] == 1 + except Exception: + return False + + @classmethod + def reset_instance(cls) -> None: + """ + Reset the singleton instance (mainly for testing). + + WARNING: This should only be used in test scenarios. + """ + with cls._lock: + if cls._instance and cls._instance._driver: + cls._instance._driver.close() + cls._instance = None + cls._driver = None -# import os -# from neo4j import GraphDatabase -# from dotenv import load_dotenv -# from pydantic_neo4j import PydanticNeo4j -# load_dotenv() +# Create the default singleton instance +try: + URI = os.getenv("NEO4J_URI_BOLT") + USERNAME = os.getenv("NEO4J_USERNAME") + PASSWORD = os.getenv("NEO4J_PASSWORD") -# class Neo4jConnection: -# def __init__(self, uri: str, user: str, password: str): -# self.__uri = uri -# self.__user = user -# self.__password = password -# self.__driver = GraphDatabase.driver(self.__uri, auth=(self.__user, self.__password)) - -# def query(self, query: str, parameters: dict = None): -# with self.__driver.session() as session: -# result = session.run(query, parameters) -# return result.data() - - -# URI = os.getenv("NEO4J_URI_BOLT") -# USERNAME = os.getenv("NEO4J_USERNAME") -# PASSWORD = os.getenv("NEO4J_PASSWORD") - -# neo4j_connection = PydanticNeo4j(URI, USERNAME, PASSWORD) -# match_util = neo4j_connection.match_utilities -# create_util = neo4j_connection.create_utilities -# database_operations = neo4j_connection.database_operations + if all([URI, USERNAME, PASSWORD]): + neo4j_connection = Neo4jConnection(URI, USERNAME, PASSWORD) + else: + # Don't create instance if credentials are missing + neo4j_connection = None +except Exception as e: + import logging + logging.getLogger(__name__).warning(f"Failed to initialize Neo4j connection: {e}") + neo4j_connection = None diff --git a/flowsint-core/src/flowsint_core/core/graph_repository.py b/flowsint-core/src/flowsint_core/core/graph_repository.py new file mode 100644 index 00000000..bbebeb93 --- /dev/null +++ b/flowsint-core/src/flowsint_core/core/graph_repository.py @@ -0,0 +1,267 @@ +""" +Graph database repository for Neo4j operations. + +This module provides a repository pattern implementation for Neo4j, +handling node and relationship operations with batching support. +""" + +from typing import Dict, Any, List, Optional, Tuple +from .graph_db import Neo4jConnection +from .graph_serializer import GraphSerializer + + +class GraphRepository: + """ + Repository for Neo4j graph database operations. + + This class follows the Repository pattern, providing a clean abstraction + over Neo4j operations and handling batching for improved performance. + """ + + def __init__(self, neo4j_connection: Optional[Neo4jConnection] = None): + """ + Initialize the graph repository. + + Args: + neo4j_connection: Optional Neo4j connection instance. + If None, uses the singleton instance. + """ + self._connection = neo4j_connection or Neo4jConnection.get_instance() + self._batch_operations: List[Tuple[str, Dict[str, Any]]] = [] + self._batch_size = 100 + + def create_node( + self, + node_type: str, + key_prop: str, + key_value: str, + sketch_id: str, + **properties: Any + ) -> None: + """ + Create or update a single node in Neo4j. + + Args: + node_type: Node label (e.g., "domain", "ip") + key_prop: Property name used as unique identifier + key_value: Value of the key property + sketch_id: Investigation sketch ID + **properties: Additional node properties + """ + if not self._connection: + return + + # Serialize properties + serialized_props = GraphSerializer.serialize_properties(properties) + + # Add required properties + serialized_props["type"] = node_type.lower() + serialized_props["sketch_id"] = sketch_id + serialized_props["label"] = serialized_props.get("label", key_value) + + # Build SET clauses + set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys()] + params = {key_prop: key_value, **serialized_props} + + # Build and execute query + query = f""" + MERGE (n:{node_type} {{{key_prop}: ${key_prop}}}) + SET {', '.join(set_clauses)} + """ + + self._connection.execute_write(query, params) + + def create_relationship( + self, + from_type: str, + from_key: str, + from_value: str, + to_type: str, + to_key: str, + to_value: str, + rel_type: str, + sketch_id: str, + **properties: Any + ) -> None: + """ + Create a relationship between two nodes. + + Args: + from_type: Source node label + from_key: Source node key property + from_value: Source node key value + to_type: Target node label + to_key: Target node key property + to_value: Target node key value + rel_type: Relationship type + sketch_id: Investigation sketch ID + **properties: Additional relationship properties + """ + if not self._connection: + return + + # Serialize relationship properties + serialized_props = GraphSerializer.serialize_properties(properties) + serialized_props["sketch_id"] = sketch_id + + # Build relationship properties string + if serialized_props: + props_str = ", ".join([f"{k}: ${k}" for k in serialized_props.keys()]) + rel_props = f"{{{props_str}}}" + else: + rel_props = "{sketch_id: $sketch_id}" + + query = f""" + MATCH (from:{from_type} {{{from_key}: $from_value}}) + MATCH (to:{to_type} {{{to_key}: $to_value}}) + MERGE (from)-[:{rel_type} {rel_props}]->(to) + """ + + params = { + "from_value": from_value, + "to_value": to_value, + **serialized_props + } + + self._connection.execute_write(query, params) + + def add_to_batch( + self, + operation_type: str, + **kwargs: Any + ) -> None: + """ + Add an operation to the batch queue. + + Args: + operation_type: Type of operation ("node" or "relationship") + **kwargs: Operation parameters + """ + if operation_type == "node": + query, params = self._build_node_query(**kwargs) + elif operation_type == "relationship": + query, params = self._build_relationship_query(**kwargs) + else: + raise ValueError(f"Unknown operation type: {operation_type}") + + self._batch_operations.append((query, params)) + + # Auto-flush if batch is full + if len(self._batch_operations) >= self._batch_size: + self.flush_batch() + + def _build_node_query( + self, + node_type: str, + key_prop: str, + key_value: str, + sketch_id: str, + **properties: Any + ) -> Tuple[str, Dict[str, Any]]: + """Build a node creation query.""" + serialized_props = GraphSerializer.serialize_properties(properties) + serialized_props["type"] = node_type.lower() + serialized_props["sketch_id"] = sketch_id + serialized_props["label"] = serialized_props.get("label", key_value) + + set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys()] + params = {key_prop: key_value, **serialized_props} + + query = f""" + MERGE (n:{node_type} {{{key_prop}: ${key_prop}}}) + SET {', '.join(set_clauses)} + """ + + return query, params + + def _build_relationship_query( + self, + from_type: str, + from_key: str, + from_value: str, + to_type: str, + to_key: str, + to_value: str, + rel_type: str, + sketch_id: str, + **properties: Any + ) -> Tuple[str, Dict[str, Any]]: + """Build a relationship creation query.""" + serialized_props = GraphSerializer.serialize_properties(properties) + serialized_props["sketch_id"] = sketch_id + + if serialized_props: + props_str = ", ".join([f"{k}: ${k}" for k in serialized_props.keys()]) + rel_props = f"{{{props_str}}}" + else: + rel_props = "{sketch_id: $sketch_id}" + + query = f""" + MATCH (from:{from_type} {{{from_key}: $from_value}}) + MATCH (to:{to_type} {{{to_key}: $to_value}}) + MERGE (from)-[:{rel_type} {rel_props}]->(to) + """ + + params = { + "from_value": from_value, + "to_value": to_value, + **serialized_props + } + + return query, params + + def flush_batch(self) -> None: + """Execute all batched operations in a single transaction.""" + if not self._batch_operations: + return + + if not self._connection: + self._batch_operations.clear() + return + + try: + self._connection.execute_batch(self._batch_operations) + finally: + self._batch_operations.clear() + + def clear_batch(self) -> None: + """Clear the batch without executing.""" + self._batch_operations.clear() + + def set_batch_size(self, size: int) -> None: + """ + Set the batch size for auto-flushing. + + Args: + size: Number of operations to batch before auto-flush + """ + if size < 1: + raise ValueError("Batch size must be at least 1") + self._batch_size = size + + def query(self, cypher: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]: + """ + Execute a custom Cypher query. + + Args: + cypher: Cypher query string + parameters: Query parameters + + Returns: + List of result records + """ + if not self._connection: + return [] + + return self._connection.query(cypher, parameters) + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit - auto-flush batch.""" + if exc_type is None: + self.flush_batch() + else: + self.clear_batch() diff --git a/flowsint-core/src/flowsint_core/core/graph_serializer.py b/flowsint-core/src/flowsint_core/core/graph_serializer.py new file mode 100644 index 00000000..63f658b6 --- /dev/null +++ b/flowsint-core/src/flowsint_core/core/graph_serializer.py @@ -0,0 +1,169 @@ +""" +Graph property serialization utilities. + +This module provides utilities for serializing complex Python objects +into Neo4j-compatible primitive types, following the Single Responsibility Principle. +""" + +from typing import Any, Dict, List + + +class GraphSerializer: + """ + Handles serialization of complex objects to Neo4j-compatible types. + + This class is responsible for converting Pydantic models, nested objects, + and other complex types into primitive types that can be stored in Neo4j. + """ + + @staticmethod + def serialize_properties(properties: Dict[str, Any]) -> Dict[str, Any]: + """ + Convert properties to Neo4j-compatible values. + + Handles: + - Pydantic models (flattened into individual properties) + - Nested objects (flattened with prefixed keys) + - Lists (converted to primitive types) + - Dictionaries (flattened with prefixed keys) + - None values (converted to empty strings) + + Args: + properties: Dictionary of properties to serialize + + Returns: + Dictionary of Neo4j-compatible properties + """ + serialized = {} + + for key, value in properties.items(): + if value is None: + serialized[key] = "" + elif GraphSerializer._is_pydantic_model(value): + # Flatten Pydantic models + flattened = GraphSerializer._flatten_pydantic(key, value) + serialized.update(flattened) + elif isinstance(value, dict): + # Flatten dictionaries + flattened = GraphSerializer._flatten_dict(key, value) + serialized.update(flattened) + elif isinstance(value, list): + # Handle lists + serialized[key] = GraphSerializer._serialize_list(value) + elif isinstance(value, (str, int, float, bool)): + # Keep primitive types as-is + serialized[key] = "" if value == "None" else value + else: + # Convert other complex types to strings + serialized[key] = str(value) if value != "None" else "" + + return serialized + + @staticmethod + def _is_pydantic_model(obj: Any) -> bool: + """Check if an object is a Pydantic model.""" + return ( + hasattr(obj, "__dict__") + and not isinstance(obj, (str, int, float, bool)) + and (hasattr(obj, "model_dump") or hasattr(obj, "dict")) + ) + + @staticmethod + def _flatten_pydantic(prefix: str, model: Any) -> Dict[str, Any]: + """ + Flatten a Pydantic model into individual properties. + + Args: + prefix: Key prefix for nested properties + model: Pydantic model instance + + Returns: + Flattened dictionary + """ + flattened = {} + + # Try Pydantic v2 first, then v1 + if hasattr(model, "model_dump"): + data = model.model_dump() + elif hasattr(model, "dict"): + data = model.dict() + else: + # Fallback to __dict__ + data = model.__dict__ + + for nested_key, nested_value in data.items(): + if nested_value is not None: + new_key = f"{prefix}_{nested_key}" + if isinstance(nested_value, (str, int, float, bool)): + flattened[new_key] = nested_value + else: + # Recursively handle nested complex types + flattened[new_key] = str(nested_value) + + return flattened + + @staticmethod + def _flatten_dict(prefix: str, data: Dict[str, Any]) -> Dict[str, Any]: + """ + Flatten a dictionary into individual properties. + + Args: + prefix: Key prefix for nested properties + data: Dictionary to flatten + + Returns: + Flattened dictionary + """ + flattened = {} + + for dict_key, dict_value in data.items(): + if dict_value is not None: + new_key = f"{prefix}_{dict_key}" + if isinstance(dict_value, (str, int, float, bool)): + flattened[new_key] = dict_value + else: + flattened[new_key] = str(dict_value) + + return flattened + + @staticmethod + def _serialize_list(items: List[Any]) -> List[Any]: + """ + Serialize a list to Neo4j-compatible types. + + Args: + items: List of items to serialize + + Returns: + List of serialized items + """ + serialized = [] + + for item in items: + if isinstance(item, (str, int, float, bool)): + serialized.append(item) + elif GraphSerializer._is_pydantic_model(item): + # Convert complex objects to strings + serialized.append(str(item)) + else: + serialized.append(str(item)) + + return serialized + + @staticmethod + def normalize_property_value(value: Any) -> Any: + """ + Normalize a single property value for Neo4j. + + Args: + value: Value to normalize + + Returns: + Normalized value + """ + if value is None: + return "" + elif isinstance(value, (str, int, float, bool)): + return "" if value == "None" else value + else: + return str(value) diff --git a/flowsint-core/src/flowsint_core/core/graph_service.py b/flowsint-core/src/flowsint_core/core/graph_service.py new file mode 100644 index 00000000..31152a3e --- /dev/null +++ b/flowsint-core/src/flowsint_core/core/graph_service.py @@ -0,0 +1,217 @@ +""" +Graph service for high-level graph operations. + +This module provides a service layer for graph operations, +integrating repository and logging functionality. +""" + +from typing import Dict, Any, Optional, Protocol +from uuid import UUID +from .graph_repository import GraphRepository +from .graph_db import Neo4jConnection + + +class LoggerProtocol(Protocol): + """Protocol for logger implementations.""" + + @staticmethod + def graph_append(sketch_id: str, message: Dict[str, Any]) -> None: + """Log a graph append message.""" + ... + + +class GraphService: + """ + High-level service for graph operations. + + This service provides a clean interface for transform operations, + handling both graph persistence and logging with proper separation of concerns. + """ + + def __init__( + self, + sketch_id: str, + neo4j_connection: Optional[Neo4jConnection] = None, + logger: Optional[LoggerProtocol] = None, + enable_batching: bool = True + ): + """ + Initialize the graph service. + + Args: + sketch_id: Investigation sketch ID + neo4j_connection: Optional Neo4j connection + logger: Optional logger instance + enable_batching: Enable batch operations + """ + self._sketch_id = sketch_id + self._repository = GraphRepository(neo4j_connection) + self._logger = logger + self._enable_batching = enable_batching + + @property + def sketch_id(self) -> str: + """Get the sketch ID.""" + return self._sketch_id + + @property + def repository(self) -> GraphRepository: + """Get the underlying repository.""" + return self._repository + + def create_node( + self, + node_type: str, + key_prop: str, + key_value: str, + **properties: Any + ) -> None: + """ + Create or update a node in the graph. + + Args: + node_type: Node label (e.g., "domain", "ip") + key_prop: Property name used as unique identifier + key_value: Value of the key property + **properties: Additional node properties + """ + if self._enable_batching: + self._repository.add_to_batch( + "node", + node_type=node_type, + key_prop=key_prop, + key_value=key_value, + sketch_id=self._sketch_id, + **properties + ) + else: + self._repository.create_node( + node_type=node_type, + key_prop=key_prop, + key_value=key_value, + sketch_id=self._sketch_id, + **properties + ) + + def create_relationship( + self, + from_type: str, + from_key: str, + from_value: str, + to_type: str, + to_key: str, + to_value: str, + rel_type: str, + **properties: Any + ) -> None: + """ + Create a relationship between two nodes. + + Args: + from_type: Source node label + from_key: Source node key property + from_value: Source node key value + to_type: Target node label + to_key: Target node key property + to_value: Target node key value + rel_type: Relationship type + **properties: Additional relationship properties + """ + if self._enable_batching: + self._repository.add_to_batch( + "relationship", + from_type=from_type, + from_key=from_key, + from_value=from_value, + to_type=to_type, + to_key=to_key, + to_value=to_value, + rel_type=rel_type, + sketch_id=self._sketch_id, + **properties + ) + else: + self._repository.create_relationship( + from_type=from_type, + from_key=from_key, + from_value=from_value, + to_type=to_type, + to_key=to_key, + to_value=to_value, + rel_type=rel_type, + sketch_id=self._sketch_id, + **properties + ) + + def log_graph_message(self, message: str) -> None: + """ + Log a graph operation message. + + Args: + message: Message to log + """ + if self._logger: + self._logger.graph_append(self._sketch_id, {"message": message}) + + def flush(self) -> None: + """Flush any pending batch operations.""" + if self._enable_batching: + self._repository.flush_batch() + + def query(self, cypher: str, parameters: Dict[str, Any] = None) -> list: + """ + Execute a custom Cypher query. + + Args: + cypher: Cypher query string + parameters: Query parameters + + Returns: + List of result records + """ + return self._repository.query(cypher, parameters) + + def set_batch_size(self, size: int) -> None: + """ + Set the batch size for operations. + + Args: + size: Number of operations to batch + """ + self._repository.set_batch_size(size) + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit - auto-flush batch.""" + if exc_type is None: + self.flush() + + +def create_graph_service( + sketch_id: str, + neo4j_connection: Optional[Neo4jConnection] = None, + enable_batching: bool = True +) -> GraphService: + """ + Factory function to create a GraphService instance. + + Args: + sketch_id: Investigation sketch ID + neo4j_connection: Optional Neo4j connection + enable_batching: Enable batch operations + + Returns: + Configured GraphService instance + """ + # Import Logger here to avoid circular imports + from .logger import Logger + + return GraphService( + sketch_id=sketch_id, + neo4j_connection=neo4j_connection, + logger=Logger, + enable_batching=enable_batching + ) diff --git a/flowsint-core/src/flowsint_core/core/transform_base.py b/flowsint-core/src/flowsint_core/core/transform_base.py index f8f7d196..130cf7c8 100644 --- a/flowsint-core/src/flowsint_core/core/transform_base.py +++ b/flowsint-core/src/flowsint_core/core/transform_base.py @@ -5,7 +5,9 @@ from pydantic.config import ConfigDict from .graph_db import Neo4jConnection from .logger import Logger from .vault import VaultProtocol +from .graph_service import GraphService, create_graph_service from ..utils import resolve_type +import warnings class InvalidTransformParams(Exception): @@ -112,14 +114,27 @@ class Transform(ABC): params_schema: Optional[List[Dict[str, Any]]] = None, vault: Optional[VaultProtocol] = None, params: Optional[Dict[str, Any]] = None, + graph_service: Optional[GraphService] = None, ): self.scan_id = scan_id or "default" self.sketch_id = sketch_id or "system" - self.neo4j_conn = neo4j_conn + self.neo4j_conn = neo4j_conn # Kept for backward compatibility self.vault = vault self.params_schema = params_schema or [] self.ParamsModel = build_params_model(self.params_schema) self.params: Dict[str, Any] = params or {} + + # Initialize graph service (new architecture) + if graph_service: + self._graph_service = graph_service + else: + # Create graph service with the provided or singleton connection + self._graph_service = create_graph_service( + sketch_id=self.sketch_id, + neo4j_connection=neo4j_conn, + enable_batching=True + ) + # Params is filled synchronously by the constructor. This params is generally constructed of # vaultSecret references, not the key directly. The idea is that the real key values are resolved after calling # async_init(), right before the execution. @@ -356,6 +371,9 @@ class Transform(ABC): results = await self.scan(preprocessed) processed = self.postprocess(results, preprocessed) + # Flush any pending batch operations + self._graph_service.flush() + if self.name() != "transform_orchestrator": Logger.completed( self.sketch_id, {"message": f"Transform {self.name()} finished."} @@ -374,90 +392,40 @@ class Transform(ABC): def create_node( self, node_type: str, key_prop: str, key_value: str, **properties ) -> None: - """Simple helper to create a single Neo4j node.""" - if not self.neo4j_conn: - return - - # Serialize properties to handle nested Pydantic objects - serialized_properties = self._serialize_properties(properties) - - # Ensure all values are Neo4j-compatible primitive types - final_properties = {} - for key, value in serialized_properties.items(): - if value is None: - final_properties[key] = "" - elif isinstance(value, (str, int, float, bool)): - final_properties[key] = "" if value == "None" else value - else: - # Convert any remaining complex types to strings - final_properties[key] = str(value) - - final_properties["type"] = node_type.lower() - final_properties["sketch_id"] = self.sketch_id - final_properties["label"] = final_properties.get("label", key_value) - - set_clauses = [f"n.{prop} = ${prop}" for prop in final_properties.keys()] - params = {key_prop: key_value, **final_properties} - - query = f""" - MERGE (n:{node_type} {{{key_prop}: ${key_prop}}}) - SET {', '.join(set_clauses)} """ - self.neo4j_conn.query(query, params) + Create a single Neo4j node. + + This method now uses the GraphService for improved performance and + better separation of concerns. + + Args: + node_type: Node label (e.g., "domain", "ip") + key_prop: Property name used as unique identifier + key_value: Value of the key property + **properties: Additional node properties + """ + self._graph_service.create_node( + node_type=node_type, + key_prop=key_prop, + key_value=key_value, + **properties + ) def _serialize_properties(self, properties: dict) -> dict: - """Convert properties to Neo4j-compatible values, handling nested Pydantic objects.""" - serialized = {} + """ + Convert properties to Neo4j-compatible values. - for key, value in properties.items(): + DEPRECATED: This method is kept for backward compatibility. + New code should use GraphSerializer directly. - if hasattr(value, "__dict__") and not isinstance( - value, (str, int, float, bool) - ): - # Handle Pydantic objects and other complex objects - if hasattr(value, "model_dump"): # Pydantic v2 - # Flatten the Pydantic object into individual properties - flattened = value.model_dump() - for nested_key, nested_value in flattened.items(): - if nested_value is not None: - serialized[f"{key}_{nested_key}"] = nested_value - elif hasattr(value, "dict"): # Pydantic v1 - # Flatten the Pydantic object into individual properties - flattened = value.dict() - for nested_key, nested_value in flattened.items(): - if nested_value is not None: - serialized[f"{key}_{nested_key}"] = nested_value - else: - # For other objects, try to convert to dict or string - try: - flattened = value.__dict__ - for nested_key, nested_value in flattened.items(): - if nested_value is not None: - serialized[f"{key}_{nested_key}"] = nested_value - except: - serialized[key] = str(value) - elif isinstance(value, list): - # Handle lists - convert all items to primitive types - serialized_list = [] - for item in value: - if hasattr(item, "__dict__") and not isinstance( - item, (str, int, float, bool) - ): - # Convert complex objects to strings - serialized_list.append(str(item)) - else: - serialized_list.append(item) - serialized[key] = serialized_list - elif isinstance(value, dict): - # Handle dictionaries - flatten them - for dict_key, dict_value in value.items(): - if dict_value is not None: - serialized[f"{key}_{dict_key}"] = dict_value - else: - # Keep primitive types as-is - serialized[key] = value + Args: + properties: Dictionary of properties to serialize - return serialized + Returns: + Dictionary of serialized properties + """ + from .graph_serializer import GraphSerializer + return GraphSerializer.serialize_properties(properties) def create_relationship( self, @@ -469,25 +437,46 @@ class Transform(ABC): to_value: str, rel_type: str, ) -> None: - """Simple helper to create a relationship between two nodes.""" - if not self.neo4j_conn: - return - - query = f""" - MATCH (from:{from_type} {{{from_key}: $from_value}}) - MATCH (to:{to_type} {{{to_key}: $to_value}}) - MERGE (from)-[:{rel_type} {{sketch_id: $sketch_id}}]->(to) """ + Create a relationship between two nodes. - self.neo4j_conn.query( - query, - { - "from_value": from_value, - "to_value": to_value, - "sketch_id": self.sketch_id, - }, + This method now uses the GraphService for improved performance and + better separation of concerns. + + Args: + from_type: Source node label + from_key: Source node key property + from_value: Source node key value + to_type: Target node label + to_key: Target node key property + to_value: Target node key value + rel_type: Relationship type + """ + self._graph_service.create_relationship( + from_type=from_type, + from_key=from_key, + from_value=from_value, + to_type=to_type, + to_key=to_key, + to_value=to_value, + rel_type=rel_type ) def log_graph_message(self, message: str) -> None: - """Simple helper to log a graph message.""" - Logger.graph_append(self.sketch_id, {"message": message}) + """ + Log a graph operation message. + + Args: + message: Message to log + """ + self._graph_service.log_graph_message(message) + + @property + def graph_service(self) -> GraphService: + """ + Get the graph service instance. + + Returns: + GraphService instance for advanced operations + """ + return self._graph_service