Fixed bug where primary keys were not correctly parsed if the primary…

… attribute was wrapped
jkminder · Jun 10, 2022 · eb633ef · eb633ef
1 parent a382824
commit eb633ef
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 1 deletion.
diff --git a/rel2graph/core/schema_compiler.py b/rel2graph/core/schema_compiler.py
@@ -15,6 +15,7 @@
 from collections import Counter
 import re
 import logging
+from numpy import extract
 from ply import lex, yacc
 
 from .factories.registrar import get_factory
@@ -136,6 +137,13 @@ def p_entity(self, p):
         # clear identifier list
         self._identifiers = []
 
+    @staticmethod
+    def _extract_key_from_attribute(attribute):
+        if attribute[0] == "AttributeFactory":
+            return attribute[1][0]
+        else:
+            return SchemaConfigParser._extract_key_from_attribute(attribute[1][0])
+
     @staticmethod
     def _inject_graphelement_args(instructions, attributes, identifier):
         """
@@ -157,7 +165,7 @@ def _inject_graphelement_args(instructions, attributes, identifier):
                     if primary_key is not None:
                         raise SchemaConfigException(
                             f"Setting two or more primary keys for one graphelement is not allowed. Conflict: '{primary_key}' <-> '{attribute[1][0]}'")
-                    primary_key = attribute[1][0]
+                    primary_key = SchemaConfigParser._extract_key_from_attribute(attribute)
             instructions[1].insert(0, raw_attributes)
             instructions[1].extend([primary_key, identifier])
         else:

diff --git a/tests/unit/core/resources/primary_keys.yaml b/tests/unit/core/resources/primary_keys.yaml
@@ -5,9 +5,18 @@ ENTITY("entity"):
   NODE("pk") c:
     + pk = entity.pk
     - attr = "someattr"
+  NODE("pk") e:
+      + pk = WRAPPER(WRAPPER(entity.pk))
+  NODE("pk") d:
+    + pk = WRAPPER(entity.pk)
   RELATION(a, "noattr", c):
   RELATION(a, "nopk", c):
     - attr = "someattr"
   RELATION(a, "pk", c):
     + pk = entity.pk
     - attr = "someattr"
+  RELATION(a, "pk", c):
+    + pk = WRAPPER(entity.pk)
+  RELATION(a, "pk", c):
+    + pk = WRAPPER(WRAPPER(entity.pk))
+
diff --git a/tests/unit/core/test_schema_compiler.py b/tests/unit/core/test_schema_compiler.py
@@ -14,6 +14,8 @@
 from rel2graph.core.schema_compiler import compile_schema, SchemaConfigParser, _precompile
 from rel2graph import register_attribute_preprocessor, SchemaConfigException
 
+
+
 ######## TESTS PRECOMPILER ##########
 
 def test_precompile_commentremoval():