github-starred/flowsint
2
0
Fork 0
mirror of https://github.com/reconurge/flowsint.git synced 2026-07-23 14:24:58 -05:00
Code Issues 980 Packages Projects Releases 14 Wiki Activity

Merge pull request #92 from reconurge/feat/only-support-pydantic

Browse Source 
feat: only support pydantic models for manipulation
This commit is contained in:
dextmorgn
2025-12-05 15:51:55 +01:00
committed by GitHub
parent 240693fb80 5bb7fae4bd
commit f58ad2ae38
21 changed files with 1089 additions and 1180 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+5 -5
View File
@@ -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'
flowsint-api/app/api/routes/sketches.py
+119 -320
View File
@@ -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
)
flowsint-app/src/components/sketches/add-item-dialog.tsx
+3 -7
View File
@@ -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,
flowsint-app/src/components/sketches/background-context-menu.tsx
+1 -1
View File
@@ -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])
flowsint-app/src/components/sketches/selected-items-panel.tsx
+1 -1
View File
@@ -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])
flowsint-app/src/stores/graph-store.ts
+6 -6
View File
@@ -19,7 +19,7 @@ interface GraphState {
updateGraphData: (nodes: GraphNode[], edges: GraphEdge[]) => void
updateNode: (nodeId: string, updates: Partial<NodeData>) => void
updateEdge: (edgeId: string, updates: Partial<GraphEdge>) => 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<GraphState>()(
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<GraphState>()(
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 })
},
flowsint-core/src/flowsint_core/core/enricher_base.py
+16 -207
View File
@@ -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:
"""
flowsint-core/src/flowsint_core/core/graph_db.py
+13 -3
View File
@@ -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)
flowsint-core/src/flowsint_core/core/graph_repository.py
+512 -294
View File
File diff suppressed because it is too large Load Diff
flowsint-core/src/flowsint_core/core/graph_serializer.py
+24 -16
View File
@@ -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
flowsint-core/src/flowsint_core/core/graph_service.py
+32 -59
View File
@@ -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,
)
flowsint-core/tests/enrichers/base.py → flowsint-core/tests/core/base.py
View File
flowsint-core/tests/core/test_enricher_base_simplified_api.py
+52
View File
@@ -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")
flowsint-core/tests/core/test_graph_repository.py
+260
View File
@@ -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"
flowsint-core/tests/test_logger_singleton.py → flowsint-core/tests/core/test_logger_singleton.py
+5 -3
View File
@@ -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
flowsint-core/tests/enrichers/test_enricher_base_simplified_api.py
-110
View File
@@ -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")
flowsint-core/tests/test_enricher_base_simplified_api.py
-117
View File
@@ -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")
flowsint-types/src/flowsint_types/individual.py
+1
View File
@@ -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
flowsint-types/src/flowsint_types/username.py
+8 -4
View File
@@ -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
flowsint-types/tests/test_labels.py
+30 -26
View File
@@ -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"
flowsint-types/tests/test_registry.py
+1 -1
View File
@@ -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