Skip to content

Commit

Permalink
Add db specific node class registry
Browse files Browse the repository at this point in the history
  • Loading branch information
Marius Conjeaud committed May 30, 2024
1 parent 7c6662a commit 2f4054c
Show file tree
Hide file tree
Showing 7 changed files with 308 additions and 68 deletions.
88 changes: 57 additions & 31 deletions neomodel/async_/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ class AsyncDatabase(local):
"""

_NODE_CLASS_REGISTRY = {}
_DB_SPECIFIC_CLASS_REGISTRY = {}

def __init__(self):
self._active_transaction = None
Expand Down Expand Up @@ -352,13 +353,42 @@ def _object_resolution(self, object_to_resolve):
# Consequently, the type checking was changed for both
# Node, Relationship objects
if isinstance(object_to_resolve, Node):
return self._NODE_CLASS_REGISTRY[
frozenset(object_to_resolve.labels)
].inflate(object_to_resolve)
_labels = frozenset(object_to_resolve.labels)
if _labels in self._NODE_CLASS_REGISTRY:
return self._NODE_CLASS_REGISTRY[_labels].inflate(object_to_resolve)
elif (
self._database_name is not None
and self._database_name in self._DB_SPECIFIC_CLASS_REGISTRY
and _labels in self._DB_SPECIFIC_CLASS_REGISTRY[self._database_name]
):
return self._DB_SPECIFIC_CLASS_REGISTRY[self._database_name][
_labels
].inflate(object_to_resolve)
else:
raise NodeClassNotDefined(
object_to_resolve,
self._NODE_CLASS_REGISTRY,
self._DB_SPECIFIC_CLASS_REGISTRY,
)

if isinstance(object_to_resolve, Relationship):
rel_type = frozenset([object_to_resolve.type])
return self._NODE_CLASS_REGISTRY[rel_type].inflate(object_to_resolve)
if rel_type in self._NODE_CLASS_REGISTRY:
return self._NODE_CLASS_REGISTRY[rel_type].inflate(object_to_resolve)
elif (
self._database_name is not None
and self._database_name in self._DB_SPECIFIC_CLASS_REGISTRY
and rel_type in self._DB_SPECIFIC_CLASS_REGISTRY[self._database_name]
):
return self._DB_SPECIFIC_CLASS_REGISTRY[self._database_name][

Check warning on line 383 in neomodel/async_/core.py

View check run for this annotation

Codecov / codecov/patch

neomodel/async_/core.py#L383

Added line #L383 was not covered by tests
rel_type
].inflate(object_to_resolve)
else:
raise RelationshipClassNotDefined(
object_to_resolve,
self._NODE_CLASS_REGISTRY,
self._DB_SPECIFIC_CLASS_REGISTRY,
)

if isinstance(object_to_resolve, Path):
from neomodel.async_.path import AsyncNeomodelPath
Expand Down Expand Up @@ -388,30 +418,13 @@ def _result_resolution(self, result_list):
# Object resolution occurs in-place
for a_result_item in enumerate(result_list):
for a_result_attribute in enumerate(a_result_item[1]):
try:
# Primitive types should remain primitive types,
# Nodes to be resolved to native objects
resolved_object = a_result_attribute[1]

resolved_object = self._object_resolution(resolved_object)

result_list[a_result_item[0]][
a_result_attribute[0]
] = resolved_object

except KeyError as exc:
# Not being able to match the label set of a node with a known object results
# in a KeyError in the internal dictionary used for resolution. If it is impossible
# to match, then raise an exception with more details about the error.
if isinstance(a_result_attribute[1], Node):
raise NodeClassNotDefined(
a_result_attribute[1], self._NODE_CLASS_REGISTRY
) from exc

if isinstance(a_result_attribute[1], Relationship):
raise RelationshipClassNotDefined(
a_result_attribute[1], self._NODE_CLASS_REGISTRY
) from exc
# Primitive types should remain primitive types,
# Nodes to be resolved to native objects
resolved_object = a_result_attribute[1]

resolved_object = self._object_resolution(resolved_object)

result_list[a_result_item[0]][a_result_attribute[0]] = resolved_object

return result_list

Expand Down Expand Up @@ -1083,10 +1096,23 @@ def build_class_registry(cls):
possible_label_combinations.append(base_label_set)

for label_set in possible_label_combinations:
if label_set not in adb._NODE_CLASS_REGISTRY:
adb._NODE_CLASS_REGISTRY[label_set] = cls
if not hasattr(cls, "__target_databases__"):
if label_set not in adb._NODE_CLASS_REGISTRY:
adb._NODE_CLASS_REGISTRY[label_set] = cls
else:
raise NodeClassAlreadyDefined(
cls, adb._NODE_CLASS_REGISTRY, adb._DB_SPECIFIC_CLASS_REGISTRY
)
else:
raise NodeClassAlreadyDefined(cls, adb._NODE_CLASS_REGISTRY)
for database in cls.__target_databases__:
if database not in adb._DB_SPECIFIC_CLASS_REGISTRY:
adb._DB_SPECIFIC_CLASS_REGISTRY[database] = {}
if label_set not in adb._DB_SPECIFIC_CLASS_REGISTRY[database]:
adb._DB_SPECIFIC_CLASS_REGISTRY[database][label_set] = cls
else:
raise NodeClassAlreadyDefined(
cls, adb._NODE_CLASS_REGISTRY, adb._DB_SPECIFIC_CLASS_REGISTRY
)


NodeBase = NodeMeta("NodeBase", (AsyncPropertyManager,), {"__abstract_node__": True})
Expand Down
5 changes: 4 additions & 1 deletion neomodel/async_/relationship_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -457,7 +457,10 @@ def __init__(
is_parent = issubclass(model_from_registry, model)
if is_direct_subclass(model, AsyncStructuredRel) and not is_parent:
raise RelationshipClassRedefined(
relation_type, adb._NODE_CLASS_REGISTRY, model
relation_type,
adb._NODE_CLASS_REGISTRY,
adb._DB_SPECIFIC_CLASS_REGISTRY,
model,
)
else:
adb._NODE_CLASS_REGISTRY[label_set] = model
Expand Down
34 changes: 30 additions & 4 deletions neomodel/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,17 +38,27 @@ class ModelDefinitionException(NeomodelException):
Abstract exception to handle error conditions related to the node-to-class registry.
"""

def __init__(self, db_node_rel_class, current_node_class_registry):
def __init__(
self,
db_node_rel_class,
current_node_class_registry,
current_db_specific_node_class_registry,
):
"""
Initialises the exception with the database node that caused the missmatch.
:param db_node_rel_class: Depending on the concrete class, this is either a Neo4j driver node object
from the DBMS, or a data model class from an application's hierarchy.
:param current_node_class_registry: Dictionary that maps frozenset of
node labels to model classes
:param current_db_specific_node_class_registry: Dictionary that maps frozenset of
node labels to model classes for specific databases
"""
self.db_node_rel_class = db_node_rel_class
self.current_node_class_registry = current_node_class_registry
self.current_db_specific_node_class_registry = (
current_db_specific_node_class_registry
)

def _get_node_class_registry_formatted(self):
"""
Expand All @@ -57,13 +67,23 @@ def _get_node_class_registry_formatted(self):
:return: str
"""
ncr_items = list(
output = "\n".join(
map(
lambda x: f"{','.join(x[0])} --> {x[1]}",
self.current_node_class_registry.items(),
)
)
return "\n".join(ncr_items)
for db, db_registry in self.current_db_specific_node_class_registry.items():
output += f"\n\nDatabase-specific: {db}\n"
output += "\n".join(

Check warning on line 78 in neomodel/exceptions.py

View check run for this annotation

Codecov / codecov/patch

neomodel/exceptions.py#L77-L78

Added lines #L77 - L78 were not covered by tests
list(
map(
lambda x: f"{','.join(x[0])} --> {x[1]}",
db_registry.items(),
)
)
)
return output


class NodeClassNotDefined(ModelDefinitionException):
Expand Down Expand Up @@ -102,6 +122,7 @@ def __init__(
self,
db_rel_class_type,
current_node_class_registry,
current_db_specific_node_class_registry,
remapping_to_class,
):
"""
Expand All @@ -110,11 +131,16 @@ def __init__(
:param db_rel_class_type: The type of the relationship that caused the error.
:type db_rel_class_type: str (The label of the relationship that caused the error)
:param current_node_class_registry: The current db object's node-class registry.
:param current_db_specific_node_class_registry: The current db object's node-class registry for specific databases.
:type current_node_class_registry: dict
:param remapping_to_class: The relationship class the relationship type was attempted to be redefined to.
:type remapping_to_class: class
"""
super().__init__(db_rel_class_type, current_node_class_registry)
super().__init__(
db_rel_class_type,
current_node_class_registry,
current_db_specific_node_class_registry,
)
self.remapping_to_class = remapping_to_class

def __str__(self):
Expand Down
88 changes: 57 additions & 31 deletions neomodel/sync_/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ class Database(local):
"""

_NODE_CLASS_REGISTRY = {}
_DB_SPECIFIC_CLASS_REGISTRY = {}

def __init__(self):
self._active_transaction = None
Expand Down Expand Up @@ -350,13 +351,42 @@ def _object_resolution(self, object_to_resolve):
# Consequently, the type checking was changed for both
# Node, Relationship objects
if isinstance(object_to_resolve, Node):
return self._NODE_CLASS_REGISTRY[
frozenset(object_to_resolve.labels)
].inflate(object_to_resolve)
_labels = frozenset(object_to_resolve.labels)
if _labels in self._NODE_CLASS_REGISTRY:
return self._NODE_CLASS_REGISTRY[_labels].inflate(object_to_resolve)
elif (
self._database_name is not None
and self._database_name in self._DB_SPECIFIC_CLASS_REGISTRY
and _labels in self._DB_SPECIFIC_CLASS_REGISTRY[self._database_name]
):
return self._DB_SPECIFIC_CLASS_REGISTRY[self._database_name][
_labels
].inflate(object_to_resolve)
else:
raise NodeClassNotDefined(
object_to_resolve,
self._NODE_CLASS_REGISTRY,
self._DB_SPECIFIC_CLASS_REGISTRY,
)

if isinstance(object_to_resolve, Relationship):
rel_type = frozenset([object_to_resolve.type])
return self._NODE_CLASS_REGISTRY[rel_type].inflate(object_to_resolve)
if rel_type in self._NODE_CLASS_REGISTRY:
return self._NODE_CLASS_REGISTRY[rel_type].inflate(object_to_resolve)
elif (
self._database_name is not None
and self._database_name in self._DB_SPECIFIC_CLASS_REGISTRY
and rel_type in self._DB_SPECIFIC_CLASS_REGISTRY[self._database_name]
):
return self._DB_SPECIFIC_CLASS_REGISTRY[self._database_name][

Check warning on line 381 in neomodel/sync_/core.py

View check run for this annotation

Codecov / codecov/patch

neomodel/sync_/core.py#L381

Added line #L381 was not covered by tests
rel_type
].inflate(object_to_resolve)
else:
raise RelationshipClassNotDefined(
object_to_resolve,
self._NODE_CLASS_REGISTRY,
self._DB_SPECIFIC_CLASS_REGISTRY,
)

if isinstance(object_to_resolve, Path):
from neomodel.sync_.path import NeomodelPath
Expand Down Expand Up @@ -386,30 +416,13 @@ def _result_resolution(self, result_list):
# Object resolution occurs in-place
for a_result_item in enumerate(result_list):
for a_result_attribute in enumerate(a_result_item[1]):
try:
# Primitive types should remain primitive types,
# Nodes to be resolved to native objects
resolved_object = a_result_attribute[1]

resolved_object = self._object_resolution(resolved_object)

result_list[a_result_item[0]][
a_result_attribute[0]
] = resolved_object

except KeyError as exc:
# Not being able to match the label set of a node with a known object results
# in a KeyError in the internal dictionary used for resolution. If it is impossible
# to match, then raise an exception with more details about the error.
if isinstance(a_result_attribute[1], Node):
raise NodeClassNotDefined(
a_result_attribute[1], self._NODE_CLASS_REGISTRY
) from exc

if isinstance(a_result_attribute[1], Relationship):
raise RelationshipClassNotDefined(
a_result_attribute[1], self._NODE_CLASS_REGISTRY
) from exc
# Primitive types should remain primitive types,
# Nodes to be resolved to native objects
resolved_object = a_result_attribute[1]

resolved_object = self._object_resolution(resolved_object)

result_list[a_result_item[0]][a_result_attribute[0]] = resolved_object

return result_list

Expand Down Expand Up @@ -1079,10 +1092,23 @@ def build_class_registry(cls):
possible_label_combinations.append(base_label_set)

for label_set in possible_label_combinations:
if label_set not in db._NODE_CLASS_REGISTRY:
db._NODE_CLASS_REGISTRY[label_set] = cls
if not hasattr(cls, "__target_databases__"):
if label_set not in db._NODE_CLASS_REGISTRY:
db._NODE_CLASS_REGISTRY[label_set] = cls
else:
raise NodeClassAlreadyDefined(
cls, db._NODE_CLASS_REGISTRY, db._DB_SPECIFIC_CLASS_REGISTRY
)
else:
raise NodeClassAlreadyDefined(cls, db._NODE_CLASS_REGISTRY)
for database in cls.__target_databases__:
if database not in db._DB_SPECIFIC_CLASS_REGISTRY:
db._DB_SPECIFIC_CLASS_REGISTRY[database] = {}
if label_set not in db._DB_SPECIFIC_CLASS_REGISTRY[database]:
db._DB_SPECIFIC_CLASS_REGISTRY[database][label_set] = cls
else:
raise NodeClassAlreadyDefined(
cls, db._NODE_CLASS_REGISTRY, db._DB_SPECIFIC_CLASS_REGISTRY
)


NodeBase = NodeMeta("NodeBase", (PropertyManager,), {"__abstract_node__": True})
Expand Down
5 changes: 4 additions & 1 deletion neomodel/sync_/relationship_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -440,7 +440,10 @@ def __init__(
is_parent = issubclass(model_from_registry, model)
if is_direct_subclass(model, StructuredRel) and not is_parent:
raise RelationshipClassRedefined(
relation_type, db._NODE_CLASS_REGISTRY, model
relation_type,
db._NODE_CLASS_REGISTRY,
db._DB_SPECIFIC_CLASS_REGISTRY,
model,
)
else:
db._NODE_CLASS_REGISTRY[label_set] = model
Expand Down
Loading

0 comments on commit 2f4054c

Please sign in to comment.