Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions flowsint-core/src/flowsint_core/core/graph/serializer.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from typing import Any, Callable, Dict, List, Optional, Type, Union

from pydantic import ValidationError
from pydantic import BaseModel, ValidationError

from flowsint_core.utils import flatten, unflatten
from flowsint_types import FlowsintType
Expand Down Expand Up @@ -160,8 +160,20 @@ def graph_node_to_neo4j_dict(node: GraphNode) -> Dict[str, Any]:

Serializes the model to JSON-compatible types and flattens nested
structures into dot-notation keys for Neo4j property storage.

nodeProperties uses exclude_unset: repository._build_node_query does
`SET n += $props`, so a full dump of a partial FlowsintType would
null out fields a prior enricher already set on the same node.
"""
neo4j_dict = node.model_dump(mode="json", serialize_as_any=True)
neo4j_dict = node.model_dump(
mode="json", serialize_as_any=True, exclude={"nodeProperties"}
)
node_properties = node.nodeProperties
neo4j_dict["nodeProperties"] = (
node_properties.model_dump(mode="json", exclude_unset=True)
if isinstance(node_properties, BaseModel)
else node_properties
)
neo4j_dict_flatten = flatten(neo4j_dict, remove_empty=False)
neo4j_dict_flatten.pop(
"nodeProperties.nodeLabel", None
Expand Down
43 changes: 42 additions & 1 deletion flowsint-core/tests/core/graph/test_graph_serializer.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
GraphSerializer,
NodeMetadata,
)
from flowsint_types import Domain, Ip
from flowsint_types import Domain, Ip, Website


def test_serializer():
Expand Down Expand Up @@ -112,6 +112,47 @@ def test_serialize_from_flowsint_type():
assert "T" in neo4j_dict["nodeMetadata.created_at"]


class TestNodePropertiesExcludeUnset:
def test_thin_website_only_sends_fields_it_set(self):
website = Website(url="https://example.com/")
node = GraphNode(
id="id",
nodeLabel="https://example.com/",
nodeType="website",
nodeProperties=website,
nodeMetadata=NodeMetadata(),
)
neo4j_dict = GraphSerializer.graph_node_to_neo4j_dict(node)

assert neo4j_dict["nodeProperties.url"] == "https://example.com/"
for unset_field in (
"active",
"status_code",
"headers",
"title",
"description",
"content",
"technologies",
"redirects",
"domain",
):
assert f"nodeProperties.{unset_field}" not in neo4j_dict

def test_richer_website_keeps_explicitly_set_fields(self):
website = Website(url="https://example.com/", active=False, status_code=404)
node = GraphNode(
id="id",
nodeLabel="https://example.com/",
nodeType="website",
nodeProperties=website,
nodeMetadata=NodeMetadata(),
)
neo4j_dict = GraphSerializer.graph_node_to_neo4j_dict(node)

assert neo4j_dict["nodeProperties.active"] is False
assert neo4j_dict["nodeProperties.status_code"] == 404


class TestCleanEmptyValues:
def test_removes_empty_strings(self):
data = {"key1": "value", "key2": "", "key3": "another"}
Expand Down
Loading