diff --git a/rel2graph/core/schema_compiler.py b/rel2graph/core/schema_compiler.py
index 4de0752..70caacb 100644
--- 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
index 49d999d..144d954 100644
--- 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))
+  
\ No newline at end of file
diff --git a/tests/unit/core/test_schema_compiler.py b/tests/unit/core/test_schema_compiler.py
index fbb7a4a..3256110 100644
--- 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():