diff --git a/Makefile b/Makefile index 721af2ae..c9e26a03 100644 --- a/Makefile +++ b/Makefile @@ -10,14 +10,14 @@ open-browser: @open http://localhost:5173 2>/dev/null || xdg-open http://localhost:5173 2>/dev/null || echo "✅ Flowsint ready at http://localhost:5173" dev: - @echo "🚀 Starting Flowsint in DEVELOPMENT mode..." + @echo "🐙 Starting Flowsint in DEVELOPMENT mode..." $(MAKE) check-env docker compose -f docker-compose.dev.yml up --build -d $(MAKE) open-browser docker compose -f docker-compose.dev.yml logs -f prod: - @echo "🚀 Starting Flowsint in PRODUCTION mode..." + @echo "🐙 Starting Flowsint in PRODUCTION mode..." $(MAKE) check-env docker compose -f docker-compose.prod.yml up --build -d $(MAKE) open-browser @@ -50,7 +50,7 @@ test: cd $(PROJECT_ROOT)/flowsint-enrichers && poetry run pytest install: - @echo "🚀 Installing Flowsint project modules..." + @echo "🐙 Installing Flowsint project modules..." @if ! command -v poetry >/dev/null 2>&1; then \ echo "⚠️ Poetry is not installed. Please install it:"; \ echo "pipx install poetry"; \ @@ -73,7 +73,7 @@ api: cd $(PROJECT_ROOT)/flowsint-api && poetry run uvicorn app.main:app --host 0.0.0.0 --port 5001 --reload frontend: - @echo "🚀 Starting frontend and opening browser..." + @echo "🐙 Starting frontend and opening browser..." @docker compose up -d flowsint-app @bash -c 'until curl -s http://localhost:5173 > /dev/null 2>&1; do sleep 1; done; open http://localhost:5173 2>/dev/null || xdg-open http://localhost:5173 2>/dev/null || echo "✅ Frontend ready at http://localhost:5173"' @@ -84,7 +84,7 @@ celery: cd $(PROJECT_ROOT)/flowsint-core && poetry run celery -A flowsint_core.core.celery worker --loglevel=info --pool=solo run: - @echo "🚀 Starting all services..." + @echo "🐙 Starting all services..." docker compose up -d @echo "⏳ Waiting for frontend to be ready..." @bash -c 'until curl -s http://localhost:5173 > /dev/null 2>&1; do sleep 1; done' diff --git a/flowsint-api/app/api/routes/sketches.py b/flowsint-api/app/api/routes/sketches.py index 63ef53a9..6340f4d3 100644 --- a/flowsint-api/app/api/routes/sketches.py +++ b/flowsint-api/app/api/routes/sketches.py @@ -9,6 +9,7 @@ from fastapi import ( Form, BackgroundTasks, ) +from flowsint_types import TYPE_REGISTRY from pydantic import BaseModel, Field from typing import Literal, List, Optional, Dict, Any from datetime import datetime, timezone @@ -29,12 +30,10 @@ router = APIRouter() class NodeData(BaseModel): label: str = Field(default="Node", description="Label/name of the node") - color: str = Field(default="Node", description="Color of the node") type: str = Field(default="Node", description="Type of the node") - # Add any other specific data fields that might be common across nodes class Config: - extra = "allow" # Accept any additional fields + extra = "allow" class NodeInput(BaseModel): @@ -47,10 +46,6 @@ class NodeInput(BaseModel): extra = "allow" # Accept any additional fields -def dict_to_cypher_props(props: dict, prefix: str = "") -> str: - return ", ".join(f"{key}: ${prefix}{key}" for key in props) - - class NodeDeleteInput(BaseModel): nodeIds: List[str] @@ -224,6 +219,29 @@ async def get_sketch_nodes( return {"nds": nodes, "rls": rels} +def clean_empty_values(data: dict) -> dict: + """Remove empty string values from dict to avoid Pydantic validation errors.""" + cleaned = {} + for key, value in data.items(): + if value == "" or value is None: + continue + if isinstance(value, dict): + cleaned_nested = clean_empty_values(value) + if cleaned_nested: + cleaned[key] = cleaned_nested + elif isinstance(value, list): + cleaned_list = [ + clean_empty_values(item) if isinstance(item, dict) else item + for item in value + if item != "" and item is not None + ] + if cleaned_list: + cleaned[key] = cleaned_list + else: + cleaned[key] = value + return cleaned + + @router.post("/{sketch_id}/nodes/add") @update_sketch_timestamp def add_node( @@ -241,60 +259,41 @@ def add_node( ) node_data = node.data.model_dump() - node_type = node_data["type"] - properties = { - "type": node_type.lower(), - "sketch_id": sketch_id, - "caption": node_data["label"], - "label": node_data["label"], - } + DetectedType = TYPE_REGISTRY.get_lowercase(node_type) + if not DetectedType: + raise HTTPException(status_code=400, detail=f"Unknown type: {node_type}") - if node_data: - flattened_data = flatten(node_data) - properties.update(flattened_data) - - cypher_props = dict_to_cypher_props(properties) - - # Add created_at to parameters - properties_with_timestamp = { - **properties, - "created_at": datetime.now(timezone.utc).isoformat(), - } - - create_query = f""" - MERGE (d:`{node_type}` {{ {cypher_props} }}) - ON CREATE SET d.created_at = $created_at - RETURN d as node, elementId(d) as id - """ + cleaned_data = clean_empty_values(node_data) try: - create_result = neo4j_connection.query(create_query, properties_with_timestamp) + pydantic_obj = DetectedType(**cleaned_data) + except Exception as e: + print(f"Pydantic validation error: {e}") + raise HTTPException( + status_code=400, detail=f"Invalid data for type {node_type}: {str(e)}" + ) + + try: + graph_repo = GraphRepository(neo4j_connection) + element_id = graph_repo.create_node(pydantic_obj, sketch_id=sketch_id) except Exception as e: - print(f"Query execution error: {e}") raise HTTPException(status_code=500, detail=f"Database error: {str(e)}") - if not create_result: - raise HTTPException( - status_code=400, detail="Node creation failed - no result returned" - ) + if not element_id: + raise HTTPException(status_code=400, detail="Node creation failed") - try: - new_node = create_result[0]["node"] - new_node["id"] = create_result[0]["id"] - except (IndexError, KeyError) as e: - print(f"Error extracting node_id: {e}, result: {create_result}") - raise HTTPException( - status_code=500, detail="Failed to extract node data from response" - ) - - new_node["data"] = node_data - new_node["data"]["id"] = new_node["id"] + pydantic_data = pydantic_obj.model_dump(mode="json") + pydantic_data["id"] = element_id + pydantic_data["type"] = node_type return { "status": "node added", - "node": new_node, + "node": { + "id": element_id, + "data": pydantic_data, + }, } @@ -363,39 +362,48 @@ def edit_node( node_data = node_edit.data.model_dump() node_type = node_data.get("type", "Node") - # Prepare properties to update - properties = { - "type": node_type.lower(), - "caption": node_data.get("label", "Node"), - "label": node_data.get("label", "Node"), - } + # Get the Pydantic type from registry + DetectedType = TYPE_REGISTRY.get_lowercase(node_type) + if not DetectedType: + raise HTTPException(status_code=400, detail=f"Unknown type: {node_type}") - # Add any additional data from the flattened node_data - if node_data: - flattened_data = flatten(node_data) - properties.update(flattened_data) + # Clean empty values + cleaned_data = clean_empty_values(node_data) - # Remove sketch_id from properties to avoid conflict (it's passed separately for security) - properties.pop("sketch_id", None) + # Convert to Pydantic object + try: + pydantic_obj = DetectedType(**cleaned_data) + except Exception as e: + print(f"Pydantic validation error: {e}") + raise HTTPException( + status_code=400, detail=f"Invalid data for type {node_type}: {str(e)}" + ) # Update node using GraphRepository try: graph_repo = GraphRepository(neo4j_connection) - updated_node = graph_repo.update_node_by_element_id( - element_id=node_edit.nodeId, sketch_id=sketch_id, **properties + updated_element_id = graph_repo.update_node( + element_id=node_edit.nodeId, + node_obj=pydantic_obj, + sketch_id=sketch_id, ) except Exception as e: print(f"Node update error: {e}") raise HTTPException(status_code=500, detail="Failed to update node") - if not updated_node: + if not updated_element_id: raise HTTPException(status_code=404, detail="Node not found or not accessible") - updated_node["data"] = node_data + # Return updated node with its data + pydantic_data = pydantic_obj.model_dump(mode="json") + pydantic_data["id"] = updated_element_id return { "status": "node updated", - "node": updated_node, + "node": { + "id": updated_element_id, + "data": pydantic_data, + }, } @@ -520,7 +528,6 @@ def merge_nodes( db: Session = Depends(get_db), current_user: Profile = Depends(get_current_user), ): - # 1. Verify the sketch exists sketch = db.query(Sketch).filter(Sketch.id == sketch_id).first() if not sketch: raise HTTPException(status_code=404, detail="Sketch not found") @@ -531,125 +538,35 @@ def merge_nodes( if not oldNodes or len(oldNodes) == 0: raise HTTPException(status_code=400, detail="oldNodes cannot be empty") - oldNodeIds = oldNodes - - # 2. Prepare the merged node data node_data = newNode.data.model_dump() if newNode.data else {} node_type = node_data.get("type", "Node") - # Build properties for the new merged node properties = { "type": node_type.lower(), - "sketch_id": sketch_id, "label": node_data.get("label", "Merged Node"), - "caption": node_data.get("label", "Merged Node"), } - # Add all other data from the node flattened_data = flatten(node_data) properties.update(flattened_data) - # 3. Check if the newNode.id is one of the old nodes (reusing existing node) - # or if we need to create a brand new node - is_reusing_node = newNode.id in oldNodeIds - - if is_reusing_node: - # Update the existing node that we're keeping - set_clause = ", ".join(f"n.{key} = ${key}" for key in properties.keys()) - create_query = f""" - MATCH (n) - WHERE elementId(n) = $nodeId AND n.sketch_id = $sketch_id - SET {set_clause} - RETURN elementId(n) as newElementId - """ - params = {"nodeId": newNode.id, "sketch_id": sketch_id, **properties} - else: - # Create a completely new node with created_at timestamp - properties["created_at"] = datetime.now(timezone.utc).isoformat() - - create_query = f""" - CREATE (n:`{node_type}`) - SET n = $properties - RETURN elementId(n) as newElementId - """ - params = {"properties": properties} - try: - result = neo4j_connection.query(create_query, params) - if not result: - raise HTTPException( - status_code=500, detail="Failed to create/update merged node" - ) - new_node_element_id = result[0]["newElementId"] - except Exception as e: - print(f"Error creating/updating merged node: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to create merged node: {str(e)}" - ) - - # 4. Copy all relationships from old nodes to the new node - # This handles both incoming and outgoing relationships while preserving types and properties - copy_relationships_query = """ - MATCH (new) WHERE elementId(new) = $newElementId - - UNWIND $oldNodeIds AS oldNodeId - MATCH (old) WHERE elementId(old) = oldNodeId AND old.sketch_id = $sketch_id - - // Copy incoming relationships - get all unique combinations - WITH new, collect(old) as oldNodes - UNWIND oldNodes as old - MATCH (src)-[r]->(old) - WHERE elementId(src) NOT IN $oldNodeIds AND elementId(src) <> $newElementId - WITH new, src, type(r) as relType, properties(r) as relProps, r - MERGE (src)-[newRel:RELATED_TO {sketch_id: $sketch_id}]->(new) - SET newRel = relProps - - WITH new, $oldNodeIds as oldNodeIds - UNWIND oldNodeIds AS oldNodeId - MATCH (old) WHERE elementId(old) = oldNodeId AND old.sketch_id = $sketch_id - - // Copy outgoing relationships - MATCH (old)-[r]->(dst) - WHERE elementId(dst) NOT IN oldNodeIds AND elementId(dst) <> $newElementId - WITH new, dst, type(r) as relType, properties(r) as relProps - MERGE (new)-[newRel:RELATED_TO {sketch_id: $sketch_id}]->(dst) - SET newRel = relProps - """ - - try: - neo4j_connection.query( - copy_relationships_query, - { - "newElementId": new_node_element_id, - "oldNodeIds": oldNodeIds, - "sketch_id": sketch_id, - }, + graph_repo = GraphRepository(neo4j_connection) + new_node_element_id = graph_repo.merge_nodes( + old_node_ids=oldNodes, + new_node_data=properties, + new_node_id=newNode.id, + sketch_id=sketch_id, ) except Exception as e: - print(f"Error copying relationships: {e}") - # Don't fail if relationship copying has issues, continue to deletion + print(f"Node merge error: {e}") + raise HTTPException(status_code=500, detail=f"Failed to merge nodes: {str(e)}") - # 5. Delete the old nodes (except if we're reusing one) - nodes_to_delete = [nid for nid in oldNodeIds if nid != new_node_element_id] - - if nodes_to_delete: - delete_query = """ - UNWIND $nodeIds AS nodeId - MATCH (old) - WHERE elementId(old) = nodeId AND old.sketch_id = $sketch_id - DETACH DELETE old - """ - try: - neo4j_connection.query( - delete_query, {"nodeIds": nodes_to_delete, "sketch_id": sketch_id} - ) - except Exception as e: - print(f"Error deleting old nodes: {e}") - raise HTTPException(status_code=500, detail="Failed to delete old nodes") + if not new_node_element_id: + raise HTTPException(status_code=500, detail="Failed to merge nodes") return { "status": "nodes merged", - "count": len(oldNodeIds), + "count": len(oldNodes), "new_node_id": new_node_element_id, } @@ -661,7 +578,6 @@ def get_related_nodes( db: Session = Depends(get_db), current_user: Profile = Depends(get_current_user), ): - # First verify the sketch exists and belongs to the user sketch = db.query(Sketch).filter(Sketch.id == sketch_id).first() if not sketch: raise HTTPException(status_code=404, detail="Sketch not found") @@ -669,134 +585,17 @@ def get_related_nodes( current_user.id, sketch.investigation_id, actions=["read"], db=db ) - # Query to get all direct relationships and connected nodes - # First, let's get the center node - center_query = """ - MATCH (n) - WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id - RETURN elementId(n) as id, labels(n) as labels, properties(n) as data - """ - try: - center_result = neo4j_connection.query( - center_query, {"sketch_id": sketch_id, "node_id": node_id} - ) - except Exception as e: - print(f"Center node query error: {e}") - raise HTTPException(status_code=500, detail="Failed to retrieve center node") - - if not center_result: - raise HTTPException(status_code=404, detail="Node not found") - - # Now get all relationships and connected nodes - relationships_query = """ - MATCH (n) - WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id - OPTIONAL MATCH (n)-[r]->(other) - WHERE other.sketch_id = $sketch_id - OPTIONAL MATCH (other)-[r2]->(n) - WHERE other.sketch_id = $sketch_id - RETURN - elementId(r) as rel_id, - type(r) as rel_type, - properties(r) as rel_data, - elementId(other) as other_node_id, - labels(other) as other_node_labels, - properties(other) as other_node_data, - 'outgoing' as direction - UNION - MATCH (n) - WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id - OPTIONAL MATCH (other)-[r]->(n) - WHERE other.sketch_id = $sketch_id - RETURN - elementId(r) as rel_id, - type(r) as rel_type, - properties(r) as rel_data, - elementId(other) as other_node_id, - labels(other) as other_node_labels, - properties(other) as other_node_data, - 'incoming' as direction - """ - - try: - result = neo4j_connection.query( - relationships_query, {"sketch_id": sketch_id, "node_id": node_id} - ) + graph_repo = GraphRepository(neo4j_connection) + result = graph_repo.get_related_nodes(node_id=node_id, sketch_id=sketch_id) except Exception as e: print(f"Related nodes query error: {e}") raise HTTPException(status_code=500, detail="Failed to retrieve related nodes") - # Extract center node info - center_record = center_result[0] - center_node = { - "id": center_record["id"], - "labels": center_record["labels"], - "data": center_record["data"], - "label": center_record["data"].get("label", "Node"), - "type": "custom", - "caption": center_record["data"].get("label", "Node"), - } + if not result["nds"]: + raise HTTPException(status_code=404, detail="Node not found") - # Collect all related nodes and relationships - related_nodes = [] - relationships = [] - seen_nodes = set() - seen_relationships = set() - - for record in result: - # Skip if no relationship found - if not record["rel_id"]: - continue - - # Add relationship if not seen - if record["rel_id"] not in seen_relationships: - if record["direction"] == "outgoing": - relationships.append( - { - "id": record["rel_id"], - "type": "straight", - "source": center_node["id"], - "target": record["other_node_id"], - "data": record["rel_data"], - "caption": record["rel_type"], - } - ) - else: # incoming - relationships.append( - { - "id": record["rel_id"], - "type": "straight", - "source": record["other_node_id"], - "target": center_node["id"], - "data": record["rel_data"], - "caption": record["rel_type"], - } - ) - seen_relationships.add(record["rel_id"]) - - # Add related node if not seen - if ( - record["other_node_id"] - and record["other_node_id"] not in seen_nodes - and record["other_node_id"] != center_node["id"] - ): - related_nodes.append( - { - "id": record["other_node_id"], - "labels": record["other_node_labels"], - "data": record["other_node_data"], - "label": record["other_node_data"].get("label", "Node"), - "type": "custom", - "caption": record["other_node_data"].get("label", "Node"), - } - ) - seen_nodes.add(record["other_node_id"]) - - # Combine center node with related nodes - all_nodes = [center_node] + related_nodes - - return {"nds": all_nodes, "rls": relationships} + return result @router.post("/{sketch_id}/import/analyze", response_model=FileParseResult) @@ -902,47 +701,47 @@ async def execute_import( # Filter only entities marked for inclusion entities_to_import = [m for m in entity_mappings if m.include] - # Import entities using GraphRepository - graph_repo = GraphRepository(neo4j_connection) - nodes_created = 0 - nodes_skipped = 0 - errors = [] + # Convert entity mappings to Pydantic objects + pydantic_nodes = [] + conversion_errors = [] for idx, mapping in enumerate(entities_to_import): - # Use data from mapping directly entity_type = mapping.entity_type label = mapping.label entity_data = mapping.data - # Flatten entity data for storage - flattened_data = flatten(entity_data) + # Get the Pydantic type from registry + DetectedType = TYPE_REGISTRY.get_lowercase(entity_type) + if not DetectedType: + conversion_errors.append(f"Entity {idx + 1} ({label}): Unknown type {entity_type}") + continue - # Remove fields that are passed as explicit parameters to avoid conflicts - flattened_data.pop('label', None) - flattened_data.pop('type', None) - flattened_data.pop('sketch_id', None) + # Clean empty values and add label + cleaned_data = clean_empty_values(entity_data) + cleaned_data["label"] = label - # Create node using GraphRepository + # Convert to Pydantic object try: - node_id = graph_repo.create_node_from_import( - node_type=entity_type, - label=label, - sketch_id=sketch_id, - **flattened_data, - ) - if node_id: - nodes_created += 1 - else: - nodes_skipped += 1 - errors.append(f"Entity {idx + 1} ({label}): Failed to create node") + pydantic_obj = DetectedType(**cleaned_data) + pydantic_nodes.append(pydantic_obj) except Exception as e: - error_msg = f"Entity {idx + 1} ({label}): {str(e)}" - errors.append(error_msg) - nodes_skipped += 1 + conversion_errors.append(f"Entity {idx + 1} ({label}): {str(e)}") + + # Batch create all nodes + graph_repo = GraphRepository(neo4j_connection) + try: + result = graph_repo.batch_create_nodes(nodes=pydantic_nodes, sketch_id=sketch_id) + nodes_created = result["nodes_created"] + batch_errors = result.get("errors", []) + all_errors = conversion_errors + batch_errors + except Exception as e: + raise HTTPException(status_code=500, detail=f"Batch import failed: {str(e)}") + + nodes_skipped = len(entities_to_import) - nodes_created return ImportExecuteResponse( - status="completed" if not errors else "completed_with_errors", + status="completed" if not all_errors else "completed_with_errors", nodes_created=nodes_created, nodes_skipped=nodes_skipped, - errors=errors[:50], # Limit to first 50 errors + errors=all_errors[:50], # Limit to first 50 errors ) diff --git a/flowsint-app/src/components/sketches/add-item-dialog.tsx b/flowsint-app/src/components/sketches/add-item-dialog.tsx index 0abea67e..69daa878 100644 --- a/flowsint-app/src/components/sketches/add-item-dialog.tsx +++ b/flowsint-app/src/components/sketches/add-item-dialog.tsx @@ -32,11 +32,9 @@ export default function AddItemDialog() { const setOpenFormDialog = useGraphStore((state) => state.setOpenFormDialog) const addNode = useGraphStore((state) => state.addNode) const addEdge = useGraphStore((state) => state.addEdge) - const replaceNodeId = useGraphStore((state) => state.replaceNodeId) + const replaceNode = useGraphStore((state) => state.replaceNode) const setActiveTab = useLayoutStore((state) => state.setActiveTab) const setImportModalOpen = useGraphSettingsStore((s) => s.setImportModalOpen) - const regenerateLayout = useGraphControls((s) => s.regenerateLayout) - const currentLayoutType = useGraphControls((s) => s.currentLayoutType) const getViewportCenter = useGraphControls((s) => s.getViewportCenter) const { id: sketch_id } = useParams({ strict: false }) @@ -148,10 +146,8 @@ export default function AddItemDialog() { sketch_id as string, JSON.stringify(newNode) ) - if (newNodeResponse.node && replaceNodeId) { - // Replace the temporary ID with the real ID from the API - replaceNodeId(tempId, newNodeResponse.node.id) - // If we have a related node, create the edge with the real ID + if (newNodeResponse.node && replaceNode) { + replaceNode(tempId, newNodeResponse.node.data) if (relatedNodeToAdd && tempEdgeId) { const relationPayload = { source: relatedNodeToAdd.id, diff --git a/flowsint-app/src/components/sketches/background-context-menu.tsx b/flowsint-app/src/components/sketches/background-context-menu.tsx index b50c65c3..0543dfe2 100644 --- a/flowsint-app/src/components/sketches/background-context-menu.tsx +++ b/flowsint-app/src/components/sketches/background-context-menu.tsx @@ -131,7 +131,7 @@ export default function BackgroundContextMenu({ { loading: `Deleting ${selectedNodes.length} node(s)...`, success: 'Nodes deleted successfully.', - error: 'Failed to delete selectedNodes.' + error: 'Failed to delete selected nodes.' } ) }, [selectedNodes, confirm, removeNodes, clearSelectedNodes, sketchId]) diff --git a/flowsint-app/src/components/sketches/selected-items-panel.tsx b/flowsint-app/src/components/sketches/selected-items-panel.tsx index 82c352e0..0a95ca79 100644 --- a/flowsint-app/src/components/sketches/selected-items-panel.tsx +++ b/flowsint-app/src/components/sketches/selected-items-panel.tsx @@ -116,7 +116,7 @@ const ActionBar = () => { { loading: `Deleting ${selectedNodes.length} node(s)...`, success: 'Nodes deleted successfully.', - error: 'Failed to delete selectedNodes.' + error: 'Failed to delete selected nodes.' } ) }, [selectedNodes, confirm, removeNodes, clearSelectedNodes, sketchId]) diff --git a/flowsint-app/src/stores/graph-store.ts b/flowsint-app/src/stores/graph-store.ts index 9fc67bea..128577f9 100644 --- a/flowsint-app/src/stores/graph-store.ts +++ b/flowsint-app/src/stores/graph-store.ts @@ -19,7 +19,7 @@ interface GraphState { updateGraphData: (nodes: GraphNode[], edges: GraphEdge[]) => void updateNode: (nodeId: string, updates: Partial) => void updateEdge: (edgeId: string, updates: Partial) => void - replaceNodeId: (oldId: string, newId: string) => void + replaceNode: (oldId: string, newData: NodeData) => void reset: () => void // === Selection & Current === @@ -188,19 +188,19 @@ export const useGraphStore = create()( set({ edges: updatedEdges, filteredNodes, filteredEdges }) }, - replaceNodeId: (oldId, newId) => { + replaceNode: (oldId, newData) => { const { nodes, edges, filters, currentNode } = get() // Update the node's ID and data.id const updatedNodes = nodes.map((node) => - node.id === oldId ? { ...node, id: newId, data: { ...node.data, id: newId } } : node + node.id === oldId ? { ...node, id: newData.id, data: newData } : node ) // Update all edges that reference this node const updatedEdges = edges.map((edge) => { if (edge.source === oldId) { - return { ...edge, source: newId } + return { ...edge, source: newData.id } } if (edge.target === oldId) { - return { ...edge, target: newId } + return { ...edge, target: newData.id } } return edge }) @@ -208,7 +208,7 @@ export const useGraphStore = create()( const filteredEdges = computeFilteredEdges(updatedEdges, filteredNodes) // Update currentNode if it matches the old ID const updatedCurrentNode = currentNode?.id === oldId - ? { ...currentNode, id: newId, data: { ...currentNode.data, id: newId } } + ? { ...currentNode, id: newData.id, data: newData } : currentNode set({ nodes: updatedNodes, edges: updatedEdges, filteredNodes, filteredEdges, currentNode: updatedCurrentNode }) }, diff --git a/flowsint-core/src/flowsint_core/core/enricher_base.py b/flowsint-core/src/flowsint_core/core/enricher_base.py index c222ea29..43b351c8 100644 --- a/flowsint-core/src/flowsint_core/core/enricher_base.py +++ b/flowsint-core/src/flowsint_core/core/enricher_base.py @@ -8,6 +8,7 @@ from .vault import VaultProtocol from .graph_service import GraphService, create_graph_service from ..utils import resolve_type + class InvalidEnricherParams(Exception): pass @@ -413,9 +414,7 @@ class Enricher(ABC): async def execute(self, values: List[str]) -> List[Dict[str, Any]]: if self.name() != "enricher_orchestrator": - Logger.info( - self.sketch_id, {"message": f"Enricher {self.name()} started."} - ) + Logger.info(self.sketch_id, {"message": f"Enricher {self.name()} started."}) try: await self.async_init() preprocessed = self.preprocess(values) @@ -440,240 +439,50 @@ class Enricher(ABC): ) return [] - def create_node( - self, node_type_or_obj, key_prop=None, key_value=None, **properties - ) -> None: + def create_node(self, node_obj, **properties) -> None: """ Create a single Neo4j node. - This method now uses the GraphService for improved performance and - better separation of concerns. - The following properties are automatically added to every node: - type: Lowercase version of node_type - sketch_id: Current sketch ID from enricher context - label: Automatically computed by FlowsintType, or defaults to key_value if not provided - created_at: ISO 8601 UTC timestamp (only on creation, not updates) - Best Practice - Use Pydantic object directly: - The simplest way is to pass a Pydantic object directly. The node type, - key property, and key value are automatically inferred: - + Use Pydantic object directly: ```python - # Best: pass the Pydantic object directly self.create_node(ip) - - # Also good if you need to override properties - self.create_node(domain, type="subdomain") ``` Args: - node_type_or_obj: Either a Pydantic object (FlowsintType), or node label string (e.g., "domain", "ip") - key_prop: Property name used as unique identifier (optional if passing Pydantic object) - key_value: Value of the key property (optional if passing Pydantic object) - **properties: Additional node properties or property overrides - - Note: - Uses MERGE semantics - if a node with the same (key_prop, sketch_id) - exists, it will be updated. The created_at field is only set on creation. + node_obj: Either a Pydantic object or node label string + **properties: Additional node properties or overrides """ - # Check if first argument is a Pydantic object - if isinstance(node_type_or_obj, BaseModel): - obj = node_type_or_obj - - # Infer node_type from class name (e.g., Ip -> "ip", Domain -> "domain") - node_type = obj.__class__.__name__.lower() - - # Get the primary field and its value - primary_field = self._get_primary_field(obj) - key_prop = primary_field - key_value = getattr(obj, primary_field) - - # If key_value is itself a Pydantic model, extract its primary value - if isinstance(key_value, BaseModel): - key_value = self._extract_primary_value(key_value) - - # Merge object properties with any overrides, but skip nested Pydantic objects - # Use model_dump(mode="json") to properly serialize Pydantic types (e.g., HttpUrl) - obj_dict = obj.model_dump(mode="json") if hasattr(obj, "model_dump") else obj.dict() - obj_properties = {} - for k, v in obj_dict.items(): - # Skip nested Pydantic objects (represented as dicts after model_dump) - if not isinstance(v, dict): - obj_properties[k] = v - obj_properties.update(properties) - properties = obj_properties - else: - # Legacy signature: node_type_or_obj is the node_type string - node_type = node_type_or_obj - - 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. - - DEPRECATED: This method is kept for backward compatibility. - New code should use GraphSerializer directly. - - Args: - properties: Dictionary of properties to serialize - - Returns: - Dictionary of serialized properties - """ - from .graph_serializer import GraphSerializer - - return GraphSerializer.serialize_properties(properties) + self._graph_service.create_node(node_obj=node_obj, **properties) def create_relationship( self, - from_type_or_obj, - from_key_or_to_obj, - from_value_or_rel_type=None, - to_type=None, - to_key=None, - to_value=None, - rel_type=None, + from_obj, + to_obj, + rel_label="IS_RELATED_TO", ) -> None: """ Create a relationship between two nodes. - This method now uses the GraphService for improved performance and - better separation of concerns. - Best Practice - Use Pydantic objects directly: - The simplest way is to pass two Pydantic objects and the relationship type: - ```python - # Best: pass Pydantic objects directly self.create_relationship(individual, domain, "HAS_DOMAIN") self.create_relationship(email, breach, "FOUND_IN_BREACH") ``` - Legacy Usage: - You can still use the explicit signature for backward compatibility: - - ```python - # Legacy: explicit signature - self.create_relationship( - "individual", "full_name", individual.full_name, - "domain", "domain", domain_name, - "HAS_DOMAIN" - ) - ``` - Args: - from_type_or_obj: Either a Pydantic object (source node) or source node label string - from_key_or_to_obj: Either a Pydantic object (target node) or source node key property - from_value_or_rel_type: Either relationship type string (if using objects) or source node key value - to_type: Target node label (only for legacy signature) - to_key: Target node key property (only for legacy signature) - to_value: Target node key value (only for legacy signature) - rel_type: Relationship type (only for legacy signature) + from_obj: Either a Pydantic object (source) or source node label + to_obj: Either a Pydantic object (target) or source node key property + rel_label: Either relationship type (Pydantic) or source node key value """ - # Check if using new signature (Pydantic objects) - if isinstance(from_type_or_obj, BaseModel) and isinstance(from_key_or_to_obj, BaseModel): - from_obj = from_type_or_obj - to_obj = from_key_or_to_obj - relationship_type = from_value_or_rel_type - - # Extract from_node info - from_node_type = from_obj.__class__.__name__.lower() - from_primary_field = self._get_primary_field(from_obj) - - # Use model_dump to properly serialize Pydantic types (e.g., HttpUrl) - from_obj_dict = from_obj.model_dump(mode="json") if hasattr(from_obj, "model_dump") else from_obj.dict() - from_key_value = from_obj_dict.get(from_primary_field) - - # If key_value is still a dict (nested Pydantic model), extract its primary value - if isinstance(from_key_value, dict): - # Get the raw nested object to extract its primary value - nested_obj = getattr(from_obj, from_primary_field) - if isinstance(nested_obj, BaseModel): - from_key_value = self._extract_primary_value(nested_obj) - - # Extract to_node info - to_node_type = to_obj.__class__.__name__.lower() - to_primary_field = self._get_primary_field(to_obj) - - # Use model_dump to properly serialize Pydantic types (e.g., HttpUrl) - to_obj_dict = to_obj.model_dump(mode="json") if hasattr(to_obj, "model_dump") else to_obj.dict() - to_key_value = to_obj_dict.get(to_primary_field) - - # If key_value is still a dict (nested Pydantic model), extract its primary value - if isinstance(to_key_value, dict): - # Get the raw nested object to extract its primary value - nested_obj = getattr(to_obj, to_primary_field) - if isinstance(nested_obj, BaseModel): - to_key_value = self._extract_primary_value(nested_obj) - - self._graph_service.create_relationship( - from_type=from_node_type, - from_key=from_primary_field, - from_value=from_key_value, - to_type=to_node_type, - to_key=to_primary_field, - to_value=to_key_value, - rel_type=relationship_type, - ) - else: - # Legacy signature - self._graph_service.create_relationship( - from_type=from_type_or_obj, - from_key=from_key_or_to_obj, - from_value=from_value_or_rel_type, - to_type=to_type, - to_key=to_key, - to_value=to_value, - rel_type=rel_type, - ) - - def _get_primary_field(self, obj: BaseModel) -> str: - """Helper method to get the primary field of a Pydantic object.""" - # Access model_fields from the class, not the instance - model_fields = obj.__class__.model_fields - - # Find the primary field (marked with json_schema_extra={"primary": True}) - primary_field = None - for field_name, field_info in model_fields.items(): - if field_info.json_schema_extra and field_info.json_schema_extra.get("primary"): - primary_field = field_name - break - - # Fallback: use first required field or first field - if primary_field is None: - for field_name, field_info in model_fields.items(): - if field_info.is_required(): - primary_field = field_name - break - if primary_field is None: - primary_field = next(iter(model_fields.keys())) - - return primary_field - - def _extract_primary_value(self, obj: BaseModel) -> Any: - """ - Extract the primitive value from a Pydantic object recursively. - If the primary field is itself a Pydantic object, recursively extract its primary value. - Uses model_dump to properly serialize Pydantic types like HttpUrl. - """ - primary_field = self._get_primary_field(obj) - - # Use model_dump to properly serialize Pydantic types (e.g., HttpUrl) - obj_dict = obj.model_dump(mode="json") if hasattr(obj, "model_dump") else obj.dict() - value = obj_dict.get(primary_field) - - # If the value is still a dict (nested Pydantic model), recursively extract - if isinstance(value, dict): - # Get the raw nested object to recursively extract - nested_obj = getattr(obj, primary_field) - if isinstance(nested_obj, BaseModel): - return self._extract_primary_value(nested_obj) - - return value + self._graph_service.create_relationship( + from_obj=from_obj, to_obj=to_obj, rel_label=rel_label + ) def log_graph_message(self, message: str) -> None: """ diff --git a/flowsint-core/src/flowsint_core/core/graph_db.py b/flowsint-core/src/flowsint_core/core/graph_db.py index 0d7e200c..8e557877 100644 --- a/flowsint-core/src/flowsint_core/core/graph_db.py +++ b/flowsint-core/src/flowsint_core/core/graph_db.py @@ -91,9 +91,17 @@ class Neo4jConnection: List of result records as dictionaries """ with self._driver.session() as session: - result = session.run(query, parameters or {}) + cleaned_params = self._clean_parameters(parameters) + result = session.run(query, cleaned_params) return result.data() + @staticmethod + def _clean_parameters(parameters: Dict[str, Any]) -> Dict[str, Any]: + """Remove None keys from parameters dict to avoid Neo4j errors.""" + if not parameters: + return {} + return {k: v for k, v in parameters.items() if k is not None} + def execute_write(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]: """ Execute a write query within a write transaction. @@ -106,7 +114,8 @@ class Neo4jConnection: List of result records as dictionaries """ def _execute(tx): - result = tx.run(query, parameters or {}) + cleaned_params = self._clean_parameters(parameters) + result = tx.run(query, cleaned_params) return result.data() with self._driver.session() as session: @@ -121,7 +130,8 @@ class Neo4jConnection: """ def _execute_batch(tx): for query, params in queries: - tx.run(query, params or {}) + cleaned_params = self._clean_parameters(params) + tx.run(query, cleaned_params) with self._driver.session() as session: session.execute_write(_execute_batch) diff --git a/flowsint-core/src/flowsint_core/core/graph_repository.py b/flowsint-core/src/flowsint_core/core/graph_repository.py index e12ae02b..9bec397a 100644 --- a/flowsint-core/src/flowsint_core/core/graph_repository.py +++ b/flowsint-core/src/flowsint_core/core/graph_repository.py @@ -5,11 +5,11 @@ 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 typing import Dict, Any, List, Optional, Tuple, Union from datetime import datetime, timezone +from pydantic import BaseModel from .graph_db import Neo4jConnection from .graph_serializer import GraphSerializer -from flowsint_types.registry import get_type class GraphRepository: @@ -33,142 +33,105 @@ class GraphRepository: self._batch_size = 100 @staticmethod - def _get_primary_field_for_type(type_name: str) -> Optional[str]: - """ - Get the primary field name for a given type. + def _get_primary_field(obj: BaseModel) -> str: + """Get the primary field of a Pydantic object.""" + model_fields = obj.__class__.model_fields - Args: - type_name: The type name (e.g., "Ip", "Domain", or lowercase variants) - - Returns: - The primary field name, or None if type not found - """ - # Get the type class from registry - type_class = get_type(type_name, case_sensitive=False) - if not type_class: - return None - - # Find the primary field (marked with json_schema_extra={"primary": True}) - model_fields = type_class.model_fields + primary_field = None for field_name, field_info in model_fields.items(): - if field_info.json_schema_extra and field_info.json_schema_extra.get("primary"): - return field_name + if field_info.json_schema_extra and field_info.json_schema_extra.get( + "primary" + ): + primary_field = field_name + break - # Fallback: use first required field or first field - for field_name, field_info in model_fields.items(): - if field_info.is_required(): - return field_name + if primary_field is None: + for field_name, field_info in model_fields.items(): + if field_info.is_required(): + primary_field = field_name + break + if primary_field is None: + primary_field = next(iter(model_fields.keys())) - # Last resort: first field - return next(iter(model_fields.keys())) if model_fields else None + return primary_field + + @staticmethod + def _extract_primary_value(obj: BaseModel) -> Any: + """ + Extract the primitive value from a Pydantic object recursively. + Uses model_dump to properly serialize Pydantic types like HttpUrl. + """ + primary_field = GraphRepository._get_primary_field(obj) + obj_dict = ( + obj.model_dump(mode="json") if hasattr(obj, "model_dump") else obj.dict() + ) + value = obj_dict.get(primary_field) + + if isinstance(value, dict): + nested_obj = getattr(obj, primary_field) + if isinstance(nested_obj, BaseModel): + return GraphRepository._extract_primary_value(nested_obj) + + return value def create_node( self, - node_type: str, - key_prop: str, - key_value: str, - sketch_id: str, - **properties: Any - ) -> None: + node_obj: BaseModel, + sketch_id: Optional[str] = None, + **properties: Any, + ) -> Optional[str]: """ Create or update a single node in Neo4j. + Supports 1 signature: + 1. Pydantic object: create_node(obj, sketch_id="...", **overrides) + 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 + node_obj: a Pydantic object + sketch_id: Investigation sketch ID (required) + **properties: Additional node properties or overrides + + Returns: + Element ID of created/updated node """ if not self._connection: - return + return None - # 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 (exclude sketch_id as it's in MERGE) - set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys() if prop != "sketch_id"] - params = { - key_prop: key_value, - "created_at": datetime.now(timezone.utc).isoformat(), - **serialized_props - } - - # Build and execute query - MERGE on both key_prop AND sketch_id for uniqueness - # Use ON CREATE SET to only set created_at when creating, not updating - query = f""" - MERGE (n:{node_type} {{{key_prop}: ${key_prop}, sketch_id: $sketch_id}}) - ON CREATE SET n.created_at = $created_at - SET {', '.join(set_clauses)} - """ - - self._connection.execute_write(query, params) + query, params = self._build_node_query(node_obj, sketch_id, **properties) + result = self._connection.query(query, params) + return result[0]["id"] if result else None 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 + from_obj: BaseModel, + to_obj: BaseModel, + rel_label: Optional[str] = None, + sketch_id: Optional[str] = None, + **properties: Any, ) -> None: """ Create a relationship between two nodes. + Supports one signature: + - Pydantic objects: create_relationship(obj1, obj2, "REL_TYPE", sketch_id="...") + 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 + from_obj: a Pydantic object (source) + to_obj: a Pydantic object (target) + rel_label: Relationship label (ex: "IS_CONNECTED_TO") + sketch_id: Investigation sketch ID (required) **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}" - - # MATCH nodes by both key and sketch_id to ensure we're connecting nodes from the same sketch - query = f""" - MATCH (from:{from_type} {{{from_key}: $from_value, sketch_id: $sketch_id}}) - MATCH (to:{to_type} {{{to_key}: $to_value, sketch_id: $sketch_id}}) - MERGE (from)-[:{rel_type} {rel_props}]->(to) - """ - - params = { - "from_value": from_value, - "to_value": to_value, - **serialized_props - } + query, params = self._build_relationship_query( + from_obj, to_obj, rel_label, sketch_id, **properties + ) self._connection.execute_write(query, params) - def add_to_batch( - self, - operation_type: str, - **kwargs: Any - ) -> None: + def add_to_batch(self, operation_type: str, **kwargs: Any) -> None: """ Add an operation to the batch queue. @@ -191,24 +154,64 @@ class GraphRepository: def _build_node_query( self, - node_type: str, - key_prop: str, - key_value: str, + node_obj: BaseModel, sketch_id: str, - **properties: Any + **properties: Any, ) -> Tuple[str, Dict[str, Any]]: """Build a node creation query.""" + + node_type = node_obj.__class__.__name__.lower() + primary_field = self._get_primary_field(node_obj) + key_prop = primary_field + key_value = getattr(node_obj, primary_field, None) + + # If primary key value is None or empty, fallback to label + if key_value is None or (isinstance(key_value, str) and not key_value): + key_prop = "label" + key_value = properties.get("label") or getattr(node_obj, "label", "Node") + + if isinstance(key_value, BaseModel): + key_value = self._extract_primary_value(key_value) + + obj_dict = ( + node_obj.model_dump(mode="json") + if hasattr(node_obj, "model_dump") + else node_obj.dict() + ) + # Extract only non-dict values, skip None keys + obj_properties = { + k: (v if v is not None else "") + for k, v in obj_dict.items() + if k is not None and not isinstance(v, dict) + } + obj_properties.update(properties) + properties = obj_properties + + if not sketch_id: + raise ValueError("sketch_id is required") + if not key_prop or key_prop is None: + raise ValueError(f"key_prop cannot be None for node type {node_type}") + 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) + label = serialized_props.get("label", key_value) + serialized_props["label"] = label - # Build SET clauses (exclude sketch_id as it's in MERGE) - set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys() if prop != "sketch_id"] + # Clean None keys from serialized_props + serialized_props = {k: v for k, v in serialized_props.items() if k is not None} + + set_clauses = [ + f"n.{prop} = ${prop}" + for prop in serialized_props.keys() + if prop != "sketch_id" + ] + + # Build params, ensuring no None keys params = { key_prop: key_value, "created_at": datetime.now(timezone.utc).isoformat(), - **serialized_props + **serialized_props, } # MERGE on both key_prop AND sketch_id for uniqueness per sketch @@ -217,44 +220,87 @@ class GraphRepository: MERGE (n:{node_type} {{{key_prop}: ${key_prop}, sketch_id: $sketch_id}}) ON CREATE SET n.created_at = $created_at SET {', '.join(set_clauses)} + RETURN elementId(n) as id """ 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 + from_obj: BaseModel, + to_obj: BaseModel, + rel_label: Optional[str] = None, + sketch_id: Optional[str] = None, + **properties: Any, ) -> Tuple[str, Dict[str, Any]]: """Build a relationship creation query.""" + # From object + from_node_type = from_obj.__class__.__name__.lower() + from_primary_field = self._get_primary_field(from_obj) + from_obj_dict = ( + from_obj.model_dump(mode="json") + if hasattr(from_obj, "model_dump") + else from_obj.dict() + ) + from_key_value = from_obj_dict.get(from_primary_field) + + if isinstance(from_key_value, dict): + nested_obj = getattr(from_obj, from_primary_field) + if isinstance(nested_obj, BaseModel): + from_key_value = self._extract_primary_value(nested_obj) + # To object + to_node_type = to_obj.__class__.__name__.lower() + to_primary_field = self._get_primary_field(to_obj) + to_obj_dict = ( + to_obj.model_dump(mode="json") + if hasattr(to_obj, "model_dump") + else to_obj.dict() + ) + to_key_value = to_obj_dict.get(to_primary_field) + + if isinstance(to_key_value, dict): + nested_obj = getattr(to_obj, to_primary_field) + if isinstance(nested_obj, BaseModel): + to_key_value = self._extract_primary_value(nested_obj) + + from_type = from_node_type + from_key = from_primary_field + from_value = from_key_value + to_type = to_node_type + to_key = to_primary_field + to_value = to_key_value + rel_type = rel_label + + if not sketch_id: + raise ValueError("sketch_id is required") + if not from_key or from_key is None: + raise ValueError( + f"from_key cannot be None for relationship type {rel_type}" + ) + if not to_key or to_key is None: + raise ValueError(f"to_key cannot be None for relationship type {rel_type}") + serialized_props = GraphSerializer.serialize_properties(properties) serialized_props["sketch_id"] = sketch_id + # Clean None keys + serialized_props = {k: v for k, v in serialized_props.items() if k is not None} + 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}" - # MATCH nodes by both key and sketch_id to ensure we're connecting nodes from the same sketch query = f""" MATCH (from:{from_type} {{{from_key}: $from_value, sketch_id: $sketch_id}}) MATCH (to:{to_type} {{{to_key}: $to_value, sketch_id: $sketch_id}}) MERGE (from)-[:{rel_type} {rel_props}]->(to) """ - params = { - "from_value": from_value, - "to_value": to_value, - **serialized_props - } + params = {"from_value": from_value, "to_value": to_value, **serialized_props} + # Clean None keys from params + params = {k: v for k, v in params.items() if k is not None} return query, params @@ -287,121 +333,122 @@ class GraphRepository: raise ValueError("Batch size must be at least 1") self._batch_size = size - def create_node_from_import( + def batch_create_nodes( self, - node_type: str, - label: str, + nodes: List[BaseModel], sketch_id: str, - **properties: Any - ) -> Optional[str]: + ) -> Dict[str, Any]: """ - Create a node from import operation using the type's primary key. - - This method uses the same MERGE strategy as enrichers: it identifies nodes - by their primary key field (e.g., "address" for IP, "domain" for Domain) - combined with sketch_id, preventing duplicates between imports and enrichers. + Create multiple nodes in a single batch transaction. Args: - node_type: Node label (e.g., "Domain", "Ip") - will be normalized to lowercase - label: Human-readable label for the node + nodes: List of Pydantic model instances to create sketch_id: Investigation sketch ID - **properties: All additional node properties Returns: - Element ID of the created/found node or None + Dictionary with: + - nodes_created: Number of successfully created nodes + - node_ids: List of created node element IDs + - errors: List of error messages for failed nodes + """ + if not self._connection: + return {"nodes_created": 0, "node_ids": [], "errors": ["No database connection"]} + + if not nodes: + return {"nodes_created": 0, "node_ids": [], "errors": []} + + # Build all queries + batch_operations = [] + errors = [] + + for idx, node_obj in enumerate(nodes): + try: + query, params = self._build_node_query( + node_obj=node_obj, + sketch_id=sketch_id, + ) + batch_operations.append((query, params)) + except Exception as e: + errors.append(f"Node {idx}: {str(e)}") + + # Execute batch + if not batch_operations: + return {"nodes_created": 0, "node_ids": [], "errors": errors} + + try: + # Execute all operations in a single transaction + results = self._connection.execute_batch(batch_operations) + + # Extract node IDs from results + node_ids = [] + if results: + for result in results: + if result and len(result) > 0 and "id" in result[0]: + node_ids.append(result[0]["id"]) + + return { + "nodes_created": len(node_ids), + "node_ids": node_ids, + "errors": errors, + } + except Exception as e: + errors.append(f"Batch execution failed: {str(e)}") + return {"nodes_created": 0, "node_ids": [], "errors": errors} + + + def update_node( + self, + element_id: str, + node_obj: BaseModel, + sketch_id: str, + ) -> Optional[str]: + """ + Update an existing node's properties using a Pydantic model. + + Args: + element_id: Neo4j element ID of the node to update + node_obj: Pydantic model instance with updated data + sketch_id: Investigation sketch ID (for safety) + + Returns: + Element ID of the updated node or None if not found """ if not self._connection: return None - # Normalize node_type to lowercase (same as enrichers) - node_type = node_type.lower() + # Extract properties from Pydantic object + obj_dict = ( + node_obj.model_dump(mode="json") + if hasattr(node_obj, "model_dump") + else node_obj.dict() + ) - # Get the primary key field for this type - key_prop = self._get_primary_field_for_type(node_type) - if not key_prop: - # Fallback: if type not found in registry, use "label" as key - key_prop = "label" - key_value = label - else: - # Extract the primary key value from properties - key_value = properties.get(key_prop) - if not key_value: - # Fallback: if primary key not in properties, use label - key_value = label - - # Serialize and prepare all properties - serialized_props = GraphSerializer.serialize_properties(properties) - serialized_props["type"] = node_type.lower() - serialized_props["sketch_id"] = sketch_id - serialized_props["label"] = label - serialized_props["caption"] = label - - # Build SET clauses (exclude sketch_id as it's in MERGE, but include key_prop) - set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys() - if prop != "sketch_id"] - - params = { - key_prop: key_value, - "created_at": datetime.now(timezone.utc).isoformat(), - **serialized_props + # Extract only non-dict values, skip None keys + properties = { + k: (v if v is not None else "") + for k, v in obj_dict.items() + if k is not None and not isinstance(v, dict) } - # MERGE on primary key + sketch_id (same as enrichers) + # Serialize properties + serialized_props = GraphSerializer.serialize_properties(properties) + serialized_props["type"] = node_obj.__class__.__name__.lower() + + # Build SET clauses + set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys()] + query = f""" - MERGE (n:{node_type} {{{key_prop}: ${key_prop}, sketch_id: $sketch_id}}) - ON CREATE SET n.created_at = $created_at + MATCH (n) + WHERE elementId(n) = $element_id AND n.sketch_id = $sketch_id SET {', '.join(set_clauses)} RETURN elementId(n) as id """ + params = {"element_id": element_id, "sketch_id": sketch_id, **serialized_props} result = self._connection.query(query, params) return result[0]["id"] if result else None - def update_node( - self, - node_type: str, - key_prop: str, - key_value: str, - sketch_id: str, - **properties: Any - ) -> Optional[Dict[str, Any]]: - """ - Update an existing node's properties. - - 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: Properties to update - - Returns: - Updated node properties or None if not found - """ - if not self._connection: - return None - - # Serialize properties - serialized_props = GraphSerializer.serialize_properties(properties) - - # Build SET clauses - set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys()] - params = {key_prop: key_value, "sketch_id": sketch_id, **serialized_props} - - query = f""" - MATCH (n:{node_type} {{{key_prop}: ${key_prop}, sketch_id: $sketch_id}}) - SET {', '.join(set_clauses)} - RETURN properties(n) as node - """ - - result = self._connection.query(query, params) - return result[0]["node"] if result else None - - def delete_nodes( - self, - node_ids: List[str], - sketch_id: str - ) -> int: + def delete_nodes(self, node_ids: List[str], sketch_id: str) -> int: """ Delete nodes by their element IDs. @@ -424,16 +471,11 @@ class GraphRepository: """ result = self._connection.query( - query, - {"node_ids": node_ids, "sketch_id": sketch_id} + query, {"node_ids": node_ids, "sketch_id": sketch_id} ) return result[0]["deleted_count"] if result else 0 - def delete_relationships( - self, - relationship_ids: List[str], - sketch_id: str - ) -> int: + def delete_relationships(self, relationship_ids: List[str], sketch_id: str) -> int: """ Delete relationships by their element IDs. @@ -456,8 +498,7 @@ class GraphRepository: """ result = self._connection.query( - query, - {"relationship_ids": relationship_ids, "sketch_id": sketch_id} + query, {"relationship_ids": relationship_ids, "sketch_id": sketch_id} ) return result[0]["deleted_count"] if result else 0 @@ -484,9 +525,7 @@ class GraphRepository: return result[0]["deleted_count"] if result else 0 def get_sketch_graph( - self, - sketch_id: str, - limit: int = 100000 + self, sketch_id: str, limit: int = 100000 ) -> Dict[str, List[Dict[str, Any]]]: """ Get all nodes and relationships for a sketch. @@ -509,8 +548,7 @@ class GraphRepository: LIMIT $limit """ nodes_result = self._connection.query( - nodes_query, - {"sketch_id": sketch_id, "limit": limit} + nodes_query, {"sketch_id": sketch_id, "limit": limit} ) if not nodes_result: @@ -526,15 +564,9 @@ class GraphRepository: RETURN elementId(r) as id, type(r) as type, elementId(a) as source, elementId(b) as target, properties(r) as data """ - rels_result = self._connection.query( - rels_query, - {"node_ids": node_ids} - ) + rels_result = self._connection.query(rels_query, {"node_ids": node_ids}) - return { - "nodes": nodes_result, - "relationships": rels_result or [] - } + return {"nodes": nodes_result, "relationships": rels_result or []} def create_relationship_by_element_id( self, @@ -542,7 +574,7 @@ class GraphRepository: to_element_id: str, rel_type: str, sketch_id: str, - **properties: Any + **properties: Any, ) -> Optional[Dict[str, Any]]: """ Create a relationship between two nodes using their element IDs. @@ -576,47 +608,15 @@ class GraphRepository: params = { "from_id": from_element_id, "to_id": to_element_id, - **serialized_props + **serialized_props, } result = self._connection.query(query, params) return result[0]["rel"] if result else None - def update_node_by_element_id( - self, - element_id: str, - sketch_id: str, - **properties: Any - ) -> Optional[Dict[str, Any]]: - """ - Update a node by its element ID. - - Args: - element_id: Neo4j element ID - sketch_id: Investigation sketch ID (for safety) - **properties: Properties to update - - Returns: - Updated node properties or None if not found - """ - if not self._connection: - return None - - serialized_props = GraphSerializer.serialize_properties(properties) - set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys()] - - query = f""" - MATCH (n) - WHERE elementId(n) = $element_id AND n.sketch_id = $sketch_id - SET {', '.join(set_clauses)} - RETURN properties(n) as node - """ - - params = {"element_id": element_id, "sketch_id": sketch_id, **serialized_props} - result = self._connection.query(query, params) - return result[0]["node"] if result else None - - def query(self, cypher: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]: + def query( + self, cypher: str, parameters: Dict[str, Any] = None + ) -> List[Dict[str, Any]]: """ Execute a custom Cypher query. @@ -633,9 +633,7 @@ class GraphRepository: return self._connection.query(cypher, parameters) def update_nodes_positions( - self, - positions: List[Dict[str, Any]], - sketch_id: str + self, positions: List[Dict[str, Any]], sketch_id: str ) -> int: """ Update positions (x, y) for multiple nodes in batch. @@ -658,18 +656,13 @@ class GraphRepository: RETURN count(n) as updated_count """ - params = { - "positions": positions, - "sketch_id": sketch_id - } + params = {"positions": positions, "sketch_id": sketch_id} result = self._connection.query(query, params) return result[0]["updated_count"] if result else 0 def get_nodes_by_ids( - self, - node_ids: List[str], - sketch_id: str + self, node_ids: List[str], sketch_id: str ) -> List[Dict[str, Any]]: """ Get nodes by their element IDs. @@ -692,12 +685,237 @@ class GraphRepository: """ result = self._connection.query( - query, - {"node_ids": node_ids, "sketch_id": sketch_id} + query, {"node_ids": node_ids, "sketch_id": sketch_id} ) return [record["data"] for record in result] if result else [] + def merge_nodes( + self, + old_node_ids: List[str], + new_node_data: Dict[str, Any], + new_node_id: Optional[str], + sketch_id: str, + ) -> Optional[str]: + """ + Merge multiple nodes into one, transferring all relationships. + + Args: + old_node_ids: List of element IDs of nodes to merge + new_node_data: Properties for the merged node + new_node_id: Optional element ID if reusing an existing node + sketch_id: Investigation sketch ID + + Returns: + Element ID of the merged node + """ + if not self._connection or not old_node_ids: + return None + + node_type = new_node_data.get("type", "Node") + properties = GraphSerializer.serialize_properties(new_node_data) + properties["sketch_id"] = sketch_id + + is_reusing_node = new_node_id and new_node_id in old_node_ids + + if is_reusing_node: + set_clause = ", ".join(f"n.{key} = ${key}" for key in properties.keys()) + create_query = f""" + MATCH (n) + WHERE elementId(n) = $nodeId AND n.sketch_id = $sketch_id + SET {set_clause} + RETURN elementId(n) as newElementId + """ + params = {"nodeId": new_node_id, "sketch_id": sketch_id, **properties} + else: + properties["created_at"] = datetime.now(timezone.utc).isoformat() + create_query = f""" + CREATE (n:`{node_type}`) + SET n = $properties + RETURN elementId(n) as newElementId + """ + params = {"properties": properties} + + result = self._connection.query(create_query, params) + if not result: + return None + + new_node_element_id = result[0]["newElementId"] + + copy_relationships_query = """ + MATCH (new) WHERE elementId(new) = $newElementId + + UNWIND $oldNodeIds AS oldNodeId + MATCH (old) WHERE elementId(old) = oldNodeId AND old.sketch_id = $sketch_id + + WITH new, collect(old) as oldNodes + UNWIND oldNodes as old + MATCH (src)-[r]->(old) + WHERE elementId(src) NOT IN $oldNodeIds AND elementId(src) <> $newElementId + WITH new, src, type(r) as relType, properties(r) as relProps, r + MERGE (src)-[newRel:RELATED_TO {sketch_id: $sketch_id}]->(new) + SET newRel = relProps + + WITH new, $oldNodeIds as oldNodeIds + UNWIND oldNodeIds AS oldNodeId + MATCH (old) WHERE elementId(old) = oldNodeId AND old.sketch_id = $sketch_id + + MATCH (old)-[r]->(dst) + WHERE elementId(dst) NOT IN oldNodeIds AND elementId(dst) <> $newElementId + WITH new, dst, type(r) as relType, properties(r) as relProps + MERGE (new)-[newRel:RELATED_TO {sketch_id: $sketch_id}]->(dst) + SET newRel = relProps + """ + + self._connection.query( + copy_relationships_query, + { + "newElementId": new_node_element_id, + "oldNodeIds": old_node_ids, + "sketch_id": sketch_id, + }, + ) + + nodes_to_delete = [nid for nid in old_node_ids if nid != new_node_element_id] + if nodes_to_delete: + delete_query = """ + UNWIND $nodeIds AS nodeId + MATCH (old) + WHERE elementId(old) = nodeId AND old.sketch_id = $sketch_id + DETACH DELETE old + """ + self._connection.query( + delete_query, {"nodeIds": nodes_to_delete, "sketch_id": sketch_id} + ) + + return new_node_element_id + + def get_related_nodes(self, node_id: str, sketch_id: str) -> Dict[str, Any]: + """ + Get a node and all its direct relationships and connected nodes. + + Args: + node_id: Element ID of the center node + sketch_id: Investigation sketch ID + + Returns: + Dictionary with 'nds' (nodes) and 'rls' (relationships) lists + """ + if not self._connection: + return {"nds": [], "rls": []} + + center_query = """ + MATCH (n) + WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id + RETURN elementId(n) as id, labels(n) as labels, properties(n) as data + """ + + center_result = self._connection.query( + center_query, {"sketch_id": sketch_id, "node_id": node_id} + ) + + if not center_result: + return {"nds": [], "rls": []} + + relationships_query = """ + MATCH (n) + WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id + OPTIONAL MATCH (n)-[r]->(other) + WHERE other.sketch_id = $sketch_id + OPTIONAL MATCH (other)-[r2]->(n) + WHERE other.sketch_id = $sketch_id + RETURN + elementId(r) as rel_id, + type(r) as rel_type, + properties(r) as rel_data, + elementId(other) as other_node_id, + labels(other) as other_node_labels, + properties(other) as other_node_data, + 'outgoing' as direction + UNION + MATCH (n) + WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id + OPTIONAL MATCH (other)-[r]->(n) + WHERE other.sketch_id = $sketch_id + RETURN + elementId(r) as rel_id, + type(r) as rel_type, + properties(r) as rel_data, + elementId(other) as other_node_id, + labels(other) as other_node_labels, + properties(other) as other_node_data, + 'incoming' as direction + """ + + result = self._connection.query( + relationships_query, {"sketch_id": sketch_id, "node_id": node_id} + ) + + center_record = center_result[0] + center_node = { + "id": center_record["id"], + "labels": center_record["labels"], + "data": center_record["data"], + "label": center_record["data"].get("label", "Node"), + "type": "custom", + "caption": center_record["data"].get("label", "Node"), + } + + related_nodes = [] + relationships = [] + seen_nodes = set() + seen_relationships = set() + + for record in result: + if not record["rel_id"]: + continue + + if record["rel_id"] not in seen_relationships: + if record["direction"] == "outgoing": + relationships.append( + { + "id": record["rel_id"], + "type": "straight", + "source": center_node["id"], + "target": record["other_node_id"], + "data": record["rel_data"], + "caption": record["rel_type"], + } + ) + else: + relationships.append( + { + "id": record["rel_id"], + "type": "straight", + "source": record["other_node_id"], + "target": center_node["id"], + "data": record["rel_data"], + "caption": record["rel_type"], + } + ) + seen_relationships.add(record["rel_id"]) + + if ( + record["other_node_id"] + and record["other_node_id"] not in seen_nodes + and record["other_node_id"] != center_node["id"] + ): + related_nodes.append( + { + "id": record["other_node_id"], + "labels": record["other_node_labels"], + "data": record["other_node_data"], + "label": record["other_node_data"].get("label", "Node"), + "type": "custom", + "caption": record["other_node_data"].get("label", "Node"), + } + ) + seen_nodes.add(record["other_node_id"]) + + all_nodes = [center_node] + related_nodes + + return {"nds": all_nodes, "rls": relationships} + def __enter__(self): """Context manager entry.""" return self diff --git a/flowsint-core/src/flowsint_core/core/graph_serializer.py b/flowsint-core/src/flowsint_core/core/graph_serializer.py index 63f658b6..1546286f 100644 --- a/flowsint-core/src/flowsint_core/core/graph_serializer.py +++ b/flowsint-core/src/flowsint_core/core/graph_serializer.py @@ -26,7 +26,7 @@ class GraphSerializer: - Nested objects (flattened with prefixed keys) - Lists (converted to primitive types) - Dictionaries (flattened with prefixed keys) - - None values (converted to empty strings) + - None values (skipped) Args: properties: Dictionary of properties to serialize @@ -37,6 +37,8 @@ class GraphSerializer: serialized = {} for key, value in properties.items(): + if key is None: + continue if value is None: serialized[key] = "" elif GraphSerializer._is_pydantic_model(value): @@ -84,21 +86,24 @@ class GraphSerializer: # Try Pydantic v2 first, then v1 if hasattr(model, "model_dump"): - data = model.model_dump() + data = model.model_dump(mode="json") elif hasattr(model, "dict"): data = model.dict() else: # Fallback to __dict__ - data = model.__dict__ + data = {k: v for k, v in model.__dict__.items() if k is not None} 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) + if nested_key is None: + continue + new_key = f"{prefix}_{nested_key}" + if nested_value is None: + flattened[new_key] = "" + elif 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 @@ -117,12 +122,15 @@ class GraphSerializer: 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) + if dict_key is None: + continue + new_key = f"{prefix}_{dict_key}" + if dict_value is None: + flattened[new_key] = "" + elif isinstance(dict_value, (str, int, float, bool)): + flattened[new_key] = dict_value + else: + flattened[new_key] = str(dict_value) return flattened diff --git a/flowsint-core/src/flowsint_core/core/graph_service.py b/flowsint-core/src/flowsint_core/core/graph_service.py index 3862fd3a..d622fa7c 100644 --- a/flowsint-core/src/flowsint_core/core/graph_service.py +++ b/flowsint-core/src/flowsint_core/core/graph_service.py @@ -5,8 +5,9 @@ This module provides a service layer for graph operations, integrating repository and logging functionality. """ -from typing import Dict, Any, Optional, Protocol +from typing import Dict, Any, Optional, Protocol, Union from uuid import UUID +from pydantic import BaseModel from .graph_repository import GraphRepository from .graph_db import Neo4jConnection @@ -33,7 +34,7 @@ class GraphService: sketch_id: str, neo4j_connection: Optional[Neo4jConnection] = None, logger: Optional[LoggerProtocol] = None, - enable_batching: bool = True + enable_batching: bool = True, ): """ Initialize the graph service. @@ -59,94 +60,66 @@ class GraphService: """Get the underlying repository.""" return self._repository - def create_node( - self, - node_type: str, - key_prop: str, - key_value: str, - **properties: Any - ) -> None: + def create_node(self, node_obj: BaseModel, **properties: Any) -> None: """ Create or update a node in the graph. - Automatically adds the following properties: - - type: Lowercase version of node_type - - sketch_id: Current sketch ID - - label: Defaults to key_value if not provided - - created_at: ISO 8601 UTC timestamp (only on creation via ON CREATE SET) + Supports one signatures: + - Pydantic object: create_node(obj, **overrides) 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 + from_obj: a Pydantic object + **properties: Additional node properties or overrides """ if self._enable_batching: self._repository.add_to_batch( "node", - node_type=node_type, - key_prop=key_prop, - key_value=key_value, + node_obj=node_obj, sketch_id=self._sketch_id, - **properties + **properties, ) else: self._repository.create_node( - node_type=node_type, - key_prop=key_prop, - key_value=key_value, + node_obj=node_obj, sketch_id=self._sketch_id, - **properties + **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 + from_obj: BaseModel, + to_obj: BaseModel, + rel_label: Optional[str] = None, + **properties: Any, ) -> None: """ Create a relationship between two nodes. + Supports 1 signature: + - Pydantic objects: create_relationship(obj1, obj2, "REL_TYPE") + 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 + from_obj: Either a Pydantic object (source) + to_obj: Either a Pydantic object (target) + rel_label: Relationship label (ex: "IS_CONNECTED_TO") **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, + from_obj=from_obj, + to_obj=to_obj, + rel_label=rel_label, sketch_id=self._sketch_id, - **properties + **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, + from_obj=from_obj, + to_obj=to_obj, + rel_label=rel_label, sketch_id=self._sketch_id, - **properties + **properties, ) def log_graph_message(self, message: str) -> None: @@ -199,7 +172,7 @@ class GraphService: def create_graph_service( sketch_id: str, neo4j_connection: Optional[Neo4jConnection] = None, - enable_batching: bool = True + enable_batching: bool = True, ) -> GraphService: """ Factory function to create a GraphService instance. @@ -219,5 +192,5 @@ def create_graph_service( sketch_id=sketch_id, neo4j_connection=neo4j_connection, logger=Logger, - enable_batching=enable_batching + enable_batching=enable_batching, ) diff --git a/flowsint-core/tests/enrichers/base.py b/flowsint-core/tests/core/base.py similarity index 100% rename from flowsint-core/tests/enrichers/base.py rename to flowsint-core/tests/core/base.py diff --git a/flowsint-core/tests/core/test_enricher_base_simplified_api.py b/flowsint-core/tests/core/test_enricher_base_simplified_api.py new file mode 100644 index 00000000..9b741eca --- /dev/null +++ b/flowsint-core/tests/core/test_enricher_base_simplified_api.py @@ -0,0 +1,52 @@ +"""Test simplified API for create_node and create_relationship.""" + +import pytest +from flowsint_core.core.enricher_base import Enricher +from flowsint_types.domain import Domain +from flowsint_types.email import Email +from flowsint_types.individual import Individual +from typing import List + + +class MockEnricher(Enricher): + """Simple enricher for testing.""" + + InputType = Domain + OutputType = Domain + + @classmethod + def name(cls) -> str: + return "test_enricher" + + @classmethod + def category(cls) -> str: + return "Test" + + @classmethod + def key(cls) -> str: + return "domain" + + async def scan(self, data: List[InputType]) -> List[OutputType]: + return data + + +def test_create_relationship_with_pydantic_objects(): + """Test that create_relationship works with Pydantic objects.""" + enricher = MockEnricher(sketch_id="test", scan_id="test") + + # Create objects + individual = Individual(first_name="John", last_name="Doe", full_name="John Doe") + domain = Domain(domain="example.com") + + # This should not raise an error + enricher.create_relationship(individual, domain, "HAS_DOMAIN") + + +def test_create_node_with_property_override(): + """Test that property overrides work with Pydantic objects.""" + enricher = MockEnricher(sketch_id="test", scan_id="test") + + domain = Domain(domain="example.com") + + # Should be able to override properties + enricher.create_node(domain, type="subdomain") diff --git a/flowsint-core/tests/core/test_graph_repository.py b/flowsint-core/tests/core/test_graph_repository.py new file mode 100644 index 00000000..a603071b --- /dev/null +++ b/flowsint-core/tests/core/test_graph_repository.py @@ -0,0 +1,260 @@ +"""Test GraphRepository batch operations.""" + +import pytest +from unittest.mock import Mock, MagicMock +from flowsint_core.core.graph_repository import GraphRepository +from flowsint_types.domain import Domain +from flowsint_types.ip import Ip +from flowsint_types.email import Email + + +@pytest.fixture +def mock_connection(): + """Create a mock Neo4j connection.""" + conn = Mock() + conn.execute_batch = Mock(return_value=[]) + conn.query = Mock(return_value=[]) + return conn + + +@pytest.fixture +def graph_repo(mock_connection): + """Create a GraphRepository with mocked connection.""" + return GraphRepository(neo4j_connection=mock_connection) + + +def test_batch_create_nodes_empty_list(graph_repo): + """Test batch_create_nodes with empty list.""" + result = graph_repo.batch_create_nodes(nodes=[], sketch_id="test-sketch") + + assert result["nodes_created"] == 0 + assert result["node_ids"] == [] + assert result["errors"] == [] + + +def test_batch_create_nodes_single_node(graph_repo, mock_connection): + """Test batch_create_nodes with a single node.""" + domain = Domain(domain="example.com") + + # Mock the batch execution to return a node ID + mock_connection.execute_batch.return_value = [ + [{"id": "element-id-123"}] + ] + + result = graph_repo.batch_create_nodes( + nodes=[domain], + sketch_id="test-sketch" + ) + + assert result["nodes_created"] == 1 + assert len(result["node_ids"]) == 1 + assert result["node_ids"][0] == "element-id-123" + assert result["errors"] == [] + + # Verify execute_batch was called once + assert mock_connection.execute_batch.call_count == 1 + + +def test_batch_create_nodes_multiple_nodes(graph_repo, mock_connection): + """Test batch_create_nodes with multiple nodes of different types.""" + domain = Domain(domain="example.com") + ip = Ip(address="192.168.1.1") + email = Email(email="test@example.com") + + # Mock the batch execution to return multiple node IDs + mock_connection.execute_batch.return_value = [ + [{"id": "element-id-1"}], + [{"id": "element-id-2"}], + [{"id": "element-id-3"}] + ] + + result = graph_repo.batch_create_nodes( + nodes=[domain, ip, email], + sketch_id="test-sketch" + ) + + assert result["nodes_created"] == 3 + assert len(result["node_ids"]) == 3 + assert result["errors"] == [] + + # Verify execute_batch was called once with 3 operations + assert mock_connection.execute_batch.call_count == 1 + batch_operations = mock_connection.execute_batch.call_args[0][0] + assert len(batch_operations) == 3 + + +def test_batch_create_nodes_with_validation_errors(graph_repo, mock_connection): + """Test batch_create_nodes when some nodes have validation errors.""" + valid_domain = Domain(domain="example.com") + + # Mock execute_batch to succeed for valid nodes + mock_connection.execute_batch.return_value = [ + [{"id": "element-id-1"}] + ] + + result = graph_repo.batch_create_nodes( + nodes=[valid_domain], + sketch_id="test-sketch" + ) + + assert result["nodes_created"] == 1 + assert len(result["errors"]) == 0 + + +def test_batch_create_nodes_batch_execution_failure(graph_repo, mock_connection): + """Test batch_create_nodes when the batch execution fails.""" + domain = Domain(domain="example.com") + + # Mock execute_batch to raise an exception + mock_connection.execute_batch.side_effect = Exception("Database connection error") + + result = graph_repo.batch_create_nodes( + nodes=[domain], + sketch_id="test-sketch" + ) + + assert result["nodes_created"] == 0 + assert result["node_ids"] == [] + assert len(result["errors"]) > 0 + assert "Database connection error" in result["errors"][0] + + +def test_batch_create_nodes_no_connection(): + """Test batch_create_nodes when there's no database connection.""" + # Create a repository with no connection by manually setting it to None + repo = GraphRepository(neo4j_connection=Mock()) + repo._connection = None # Force no connection + domain = Domain(domain="example.com") + + result = repo.batch_create_nodes( + nodes=[domain], + sketch_id="test-sketch" + ) + + assert result["nodes_created"] == 0 + assert result["node_ids"] == [] + assert len(result["errors"]) == 1 + assert "No database connection" in result["errors"][0] + + +def test_batch_create_nodes_with_label_fallback(graph_repo, mock_connection): + """Test batch_create_nodes with nodes that need label fallback.""" + from pydantic import BaseModel + + class CustomNode(BaseModel): + """Node type without a clear primary field.""" + name: str = "" + label: str = "Custom Node" + + node = CustomNode(name="", label="Fallback Label") + + # Mock the batch execution + mock_connection.execute_batch.return_value = [ + [{"id": "element-id-fallback"}] + ] + + result = graph_repo.batch_create_nodes( + nodes=[node], + sketch_id="test-sketch" + ) + + # Should succeed with fallback to label + assert result["nodes_created"] == 1 + assert len(result["node_ids"]) == 1 + + +def test_batch_create_nodes_large_batch(graph_repo, mock_connection): + """Test batch_create_nodes with a large number of nodes.""" + # Create 1000 domains + domains = [Domain(domain=f"example{i}.com") for i in range(1000)] + + # Mock the batch execution to return 1000 node IDs + mock_connection.execute_batch.return_value = [ + [{"id": f"element-id-{i}"}] for i in range(1000) + ] + + result = graph_repo.batch_create_nodes( + nodes=domains, + sketch_id="test-sketch" + ) + + assert result["nodes_created"] == 1000 + assert len(result["node_ids"]) == 1000 + assert result["errors"] == [] + + # Verify execute_batch was called once (single transaction) + assert mock_connection.execute_batch.call_count == 1 + batch_operations = mock_connection.execute_batch.call_args[0][0] + assert len(batch_operations) == 1000 + + +# Tests for update_node + + +def test_update_node_success(graph_repo, mock_connection): + """Test update_node with a valid Pydantic object.""" + domain = Domain(domain="updated-example.com", label="Updated Domain") + + # Mock the query to return the element ID + mock_connection.query.return_value = [{"id": "element-id-123"}] + + result = graph_repo.update_node( + element_id="element-id-123", + node_obj=domain, + sketch_id="test-sketch" + ) + + assert result == "element-id-123" + assert mock_connection.query.call_count == 1 + + +def test_update_node_not_found(graph_repo, mock_connection): + """Test update_node when node doesn't exist.""" + domain = Domain(domain="example.com") + + # Mock the query to return empty result + mock_connection.query.return_value = [] + + result = graph_repo.update_node( + element_id="non-existent-id", + node_obj=domain, + sketch_id="test-sketch" + ) + + assert result is None + + +def test_update_node_no_connection(): + """Test update_node when there's no database connection.""" + repo = GraphRepository(neo4j_connection=Mock()) + repo._connection = None + domain = Domain(domain="example.com") + + result = repo.update_node( + element_id="element-id-123", + node_obj=domain, + sketch_id="test-sketch" + ) + + assert result is None + + +def test_update_node_different_types(graph_repo, mock_connection): + """Test update_node with different Pydantic types.""" + ip = Ip(address="10.0.0.1", label="Updated IP") + + # Mock the query + mock_connection.query.return_value = [{"id": "element-id-456"}] + + result = graph_repo.update_node( + element_id="element-id-456", + node_obj=ip, + sketch_id="test-sketch" + ) + + assert result == "element-id-456" + + # Verify the query contains the correct type + query_call = mock_connection.query.call_args + params = query_call[0][1] + assert params["type"] == "ip" diff --git a/flowsint-core/tests/test_logger_singleton.py b/flowsint-core/tests/core/test_logger_singleton.py similarity index 98% rename from flowsint-core/tests/test_logger_singleton.py rename to flowsint-core/tests/core/test_logger_singleton.py index 29d7d92d..31ff59b0 100644 --- a/flowsint-core/tests/test_logger_singleton.py +++ b/flowsint-core/tests/core/test_logger_singleton.py @@ -35,9 +35,11 @@ def mock_db_session(): @pytest.fixture def mock_get_db(mock_db_session): """Mock the get_db generator.""" - with patch('flowsint_core.core.logger.get_db') as mock: - mock.return_value = iter([mock_db_session]) - yield mock + def mock_generator(): + yield mock_db_session + + with patch('flowsint_core.core.logger.get_db', mock_generator): + yield mock_generator @pytest.fixture diff --git a/flowsint-core/tests/enrichers/test_enricher_base_simplified_api.py b/flowsint-core/tests/enrichers/test_enricher_base_simplified_api.py deleted file mode 100644 index 697d1485..00000000 --- a/flowsint-core/tests/enrichers/test_enricher_base_simplified_api.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Test simplified API for create_node and create_relationship.""" - -import pytest -from flowsint_core.core.enricher_base import Enricher -from flowsint_types.domain import Domain -from flowsint_types.email import Email -from flowsint_types.individual import Individual -from typing import List - - -class MockEnricher(Enricher): - """Simple enricher for testing.""" - - InputType = Domain - OutputType = Domain - - @classmethod - def name(cls) -> str: - return "test_enricher" - - @classmethod - def category(cls) -> str: - return "Test" - - @classmethod - def key(cls) -> str: - return "domain" - - async def scan(self, data: List[InputType]) -> List[OutputType]: - return data - - -def test_create_node_with_pydantic_object(): - """Test that create_node works with Pydantic objects.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - # Create a domain object - domain = Domain(domain="example.com") - - # This should not raise an error - enricher.create_node(domain) - - # Verify the helper method works - primary_field = enricher._get_primary_field(domain) - assert primary_field == "domain" - - -def test_create_relationship_with_pydantic_objects(): - """Test that create_relationship works with Pydantic objects.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - # Create objects - individual = Individual(first_name="John", last_name="Doe", full_name="John Doe") - domain = Domain(domain="example.com") - - # This should not raise an error - enricher.create_relationship(individual, domain, "HAS_DOMAIN") - - -def test_create_node_legacy_signature(): - """Test that legacy create_node signature still works.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - domain = Domain(domain="example.com") - - # Legacy signature should still work - enricher.create_node("domain", "domain", "example.com", **domain.__dict__) - - -def test_create_relationship_legacy_signature(): - """Test that legacy create_relationship signature still works.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - # Legacy signature should still work - enricher.create_relationship( - "individual", - "full_name", - "John Doe", - "domain", - "domain", - "example.com", - "HAS_DOMAIN", - ) - - -def test_get_primary_field(): - """Test the _get_primary_field helper method.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - # Test with Domain (has primary field marked) - domain = Domain(domain="example.com") - assert enricher._get_primary_field(domain) == "domain" - - # Test with Email (has primary field marked) - email = Email(email="test@example.com") - assert enricher._get_primary_field(email) == "email" - - # Test with Individual (no primary field marked, falls back to first required field) - individual = Individual(first_name="John", last_name="Doe", full_name="John Doe") - assert enricher._get_primary_field(individual) == "full_name" - - -def test_create_node_with_property_override(): - """Test that property overrides work with Pydantic objects.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - domain = Domain(domain="example.com") - - # Should be able to override properties - enricher.create_node(domain, type="subdomain") diff --git a/flowsint-core/tests/test_enricher_base_simplified_api.py b/flowsint-core/tests/test_enricher_base_simplified_api.py deleted file mode 100644 index 9569f64e..00000000 --- a/flowsint-core/tests/test_enricher_base_simplified_api.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Test simplified API for create_node and create_relationship.""" -import pytest -from flowsint_core.core.enricher_base import Enricher -from flowsint_types.domain import Domain -from flowsint_types.email import Email -from flowsint_types.individual import Individual -from typing import List - - -class MockEnricher(Enricher): - """Simple enricher for testing.""" - - InputType = Domain - OutputType = Domain - - @classmethod - def name(cls) -> str: - return "test_enricher" - - @classmethod - def category(cls) -> str: - return "Test" - - @classmethod - def key(cls) -> str: - return "domain" - - async def scan(self, data: List[InputType]) -> List[OutputType]: - return data - - -def test_create_node_with_pydantic_object(): - """Test that create_node works with Pydantic objects.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - # Create a domain object - domain = Domain(domain="example.com") - - # This should not raise an error - enricher.create_node(domain) - - # Verify the helper method works - primary_field = enricher._get_primary_field(domain) - assert primary_field == "domain" - - -def test_create_relationship_with_pydantic_objects(): - """Test that create_relationship works with Pydantic objects.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - # Create objects - individual = Individual( - first_name="John", - last_name="Doe", - full_name="John Doe" - ) - domain = Domain(domain="example.com") - - # This should not raise an error - enricher.create_relationship(individual, domain, "HAS_DOMAIN") - - -def test_create_node_legacy_signature(): - """Test that legacy create_node signature still works.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - domain = Domain(domain="example.com") - - # Legacy signature should still work - enricher.create_node("domain", "domain", "example.com", **domain.__dict__) - - -def test_create_relationship_legacy_signature(): - """Test that legacy create_relationship signature still works.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - # Legacy signature should still work - enricher.create_relationship( - "individual", - "full_name", - "John Doe", - "domain", - "domain", - "example.com", - "HAS_DOMAIN" - ) - - -def test_get_primary_field(): - """Test the _get_primary_field helper method.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - # Test with Domain (has primary field marked) - domain = Domain(domain="example.com") - assert enricher._get_primary_field(domain) == "domain" - - # Test with Email (has primary field marked) - email = Email(email="test@example.com") - assert enricher._get_primary_field(email) == "email" - - # Test with Individual (no primary field marked, falls back to first required field) - individual = Individual( - first_name="John", - last_name="Doe", - full_name="John Doe" - ) - assert enricher._get_primary_field(individual) == "full_name" - - -def test_create_node_with_property_override(): - """Test that property overrides work with Pydantic objects.""" - enricher = MockEnricher(sketch_id="test", scan_id="test") - - domain = Domain(domain="example.com") - - # Should be able to override properties - enricher.create_node(domain, type="subdomain") diff --git a/flowsint-types/src/flowsint_types/individual.py b/flowsint-types/src/flowsint_types/individual.py index 413a0615..d62cc1a8 100644 --- a/flowsint-types/src/flowsint_types/individual.py +++ b/flowsint-types/src/flowsint_types/individual.py @@ -411,6 +411,7 @@ class Individual(FlowsintType): self.label = self.first_name elif self.last_name: self.label = self.last_name + self.full_name = self.label return self @classmethod diff --git a/flowsint-types/src/flowsint_types/username.py b/flowsint-types/src/flowsint_types/username.py index 5c57415d..a6326ff6 100644 --- a/flowsint-types/src/flowsint_types/username.py +++ b/flowsint-types/src/flowsint_types/username.py @@ -13,14 +13,16 @@ class Username(FlowsintType): ..., description="Username or handle string", title="Username value", - json_schema_extra={"primary": True} + json_schema_extra={"primary": True}, + ) + platform: Optional[str] = Field( + None, description="Platform name, e.g., 'twitter'", title="Username platform" ) - platform: Optional[str] = Field(None, description="Platform name, e.g., 'twitter'", title="Username platform") last_seen: Optional[str] = Field( None, description="Last time this username was observed", title="Last seen at" ) - @field_validator('value') + @field_validator("value") @classmethod def validate_username(cls, v: str) -> str: """Validate username format. @@ -31,13 +33,15 @@ class Username(FlowsintType): - Underscores (_) - Hyphens (-) """ + if v.startswith("@"): + v = v[1:] # We remove it if not re.match(r"^[a-zA-Z0-9_-]{3,80}$", v): raise ValueError( f"Invalid username: {v}. Must be 3-80 characters and contain only letters, numbers, underscores, and hyphens." ) return v - @model_validator(mode='after') + @model_validator(mode="after") def compute_label(self) -> Self: self.label = f"{self.value}" return self diff --git a/flowsint-types/tests/test_labels.py b/flowsint-types/tests/test_labels.py index 9f5aad13..11cac887 100644 --- a/flowsint-types/tests/test_labels.py +++ b/flowsint-types/tests/test_labels.py @@ -1,8 +1,21 @@ from flowsint_types import ( - Domain, Ip, Individual, Email, Phone, Organization, - Username, Credential, CryptoWallet, CryptoNFT, - CryptoWalletTransaction, SocialAccount, Website, - Port, CIDR, ASN, Location + Domain, + Ip, + Individual, + Email, + Phone, + Organization, + Username, + Credential, + CryptoWallet, + CryptoNFT, + CryptoWalletTransaction, + SocialAccount, + Website, + Port, + CIDR, + ASN, + Location, ) @@ -19,6 +32,7 @@ def test_domain_label(): def test_individual_label(): individual = Individual(first_name="John", last_name="Doe") assert individual.label == "John Doe" + assert individual.full_name == "John Doe" def test_email_label(): @@ -52,6 +66,10 @@ def test_organization_label_with_nom_raison_sociale(): def test_username_label(): username = Username(value="johndoe") assert username.label == "johndoe" + +def test_username_label_2(): + username = Username(value="@johndoe") + assert username.label == "johndoe" def test_credential_label(): @@ -77,7 +95,7 @@ def test_crypto_nft_label_with_name(): wallet=wallet, contract_address="0x123d35Cc6634C0532925a3b844Bc454e4438f123", token_id="1234", - name="Cool NFT" + name="Cool NFT", ) assert nft.label == "Cool NFT" @@ -88,7 +106,7 @@ def test_crypto_nft_label_with_collection(): wallet=wallet, contract_address="0x123d35Cc6634C0532925a3b844Bc454e4438f123", token_id="1234", - collection_name="Bored Apes" + collection_name="Bored Apes", ) assert nft.label == "Bored Apes #1234" @@ -98,27 +116,21 @@ def test_crypto_nft_label_fallback_uid(): nft = CryptoNFT( wallet=wallet, contract_address="0x123d35Cc6634C0532925a3b844Bc454e4438f123", - token_id="1234" + token_id="1234", ) assert nft.label == "0x123d35Cc6634C0532925a3b844Bc454e4438f123:1234" def test_crypto_wallet_transaction_label_with_hash(): source_wallet = CryptoWallet(address="0x742d35Cc6634C0532925a3b844Bc454e4438f44e") - transaction = CryptoWalletTransaction( - source=source_wallet, - hash="0xabc123def456" - ) + transaction = CryptoWalletTransaction(source=source_wallet, hash="0xabc123def456") assert transaction.label == "0xabc123def456" def test_crypto_wallet_transaction_label_with_source_and_target(): source_wallet = CryptoWallet(address="0x742d35Cc6634C0532925a3b844Bc454e4438f44e") target_wallet = CryptoWallet(address="0x123d35Cc6634C0532925a3b844Bc454e4438f123") - transaction = CryptoWalletTransaction( - source=source_wallet, - target=target_wallet - ) + transaction = CryptoWalletTransaction(source=source_wallet, target=target_wallet) assert transaction.label == "Transaction from 0x742d35... to 0x123d35..." @@ -131,19 +143,14 @@ def test_crypto_wallet_transaction_label_source_only(): def test_social_account_label_with_display_name(): username = Username(value="johndoe") account = SocialAccount( - username=username, - display_name="John Doe", - platform="twitter" + username=username, display_name="John Doe", platform="twitter" ) assert account.label == "John Doe (@johndoe)" def test_social_account_label_without_display_name(): username = Username(value="johndoe") - account = SocialAccount( - username=username, - platform="twitter" - ) + account = SocialAccount(username=username, platform="twitter") assert account.label == "@johndoe" @@ -187,9 +194,6 @@ def test_asn_label_without_name(): def test_location_label(): location = Location( - address="123 Main St", - city="Paris", - country="France", - zip="75001" + address="123 Main St", city="Paris", country="France", zip="75001" ) assert location.label == "123 Main St, Paris, France" diff --git a/flowsint-types/tests/test_registry.py b/flowsint-types/tests/test_registry.py index 4c968394..ec269500 100644 --- a/flowsint-types/tests/test_registry.py +++ b/flowsint-types/tests/test_registry.py @@ -1,4 +1,4 @@ -from flowsint_types import Domain, Ip, get_type +from flowsint_types import Domain, Individual, Ip, get_type from flowsint_types.registry import TYPE_REGISTRY