From 1d83abcb1bec92a5047330b86e5cc96edf6746a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sascha=20K=C3=BChl?= <s.kuehl@tms-bonn.de>
Date: Tue, 11 Feb 2025 16:34:06 +0100
Subject: [PATCH] schema node analyze progress

---
 .../matador/object/attribute_definition.hpp   |  1 +
 .../object/attribute_definition_generator.hpp | 25 +++++----
 include/matador/object/basic_object_info.hpp  |  8 +--
 include/matador/object/object_info.hpp        | 12 +++--
 .../matador/object/primary_key_resolver.hpp   | 20 +++----
 include/matador/object/schema.hpp             | 52 ++++++++++---------
 include/matador/object/schema_node.hpp        | 25 ++++++++-
 include/matador/sql/column_generator.hpp      |  9 ++--
 source/core/object/attribute_definition.cpp   |  4 ++
 .../object/attribute_definition_generator.cpp |  5 +-
 source/core/object/basic_object_info.cpp      | 28 ++++++----
 source/core/object/primary_key_resolver.cpp   |  4 +-
 source/core/object/schema.cpp                 | 27 +++++-----
 source/core/object/schema_node.cpp            | 18 ++++---
 test/core/CMakeLists.txt                      |  3 +-
 ...p => AttributeDefinitionGeneratorTest.cpp} |  5 ++
 test/core/object/PrimaryKeyResolverTest.cpp   | 22 ++++++++
 test/orm/orm/SessionQueryBuilderTest.cpp      |  2 +
 18 files changed, 174 insertions(+), 96 deletions(-)
 rename test/core/object/{ColumnDefinitionGeneratorTest.cpp => AttributeDefinitionGeneratorTest.cpp} (91%)
 create mode 100644 test/core/object/PrimaryKeyResolverTest.cpp

diff --git a/include/matador/object/attribute_definition.hpp b/include/matador/object/attribute_definition.hpp
index 3ee8b86..c9925eb 100644
--- a/include/matador/object/attribute_definition.hpp
+++ b/include/matador/object/attribute_definition.hpp
@@ -59,6 +59,7 @@ public:
   attribute_definition(std::string name, utils::basic_type type, size_t index, const std::shared_ptr<attribute_definition> &ref_column, const utils::field_attributes& attr, null_option null_opt);
 
   [[nodiscard]] const std::string& name() const;
+  void name(const std::string& n);
   [[nodiscard]] std::string full_name() const;
   [[nodiscard]] std::string table_name() const;
   [[nodiscard]] int index() const;
diff --git a/include/matador/object/attribute_definition_generator.hpp b/include/matador/object/attribute_definition_generator.hpp
index 215c0de..a8bd051 100644
--- a/include/matador/object/attribute_definition_generator.hpp
+++ b/include/matador/object/attribute_definition_generator.hpp
@@ -5,6 +5,7 @@
 
 #include "matador/utils/access.hpp"
 #include "matador/utils/data_type_traits.hpp"
+#include "matador/utils/error.hpp"
 #include "matador/utils/field_attributes.hpp"
 #include "matador/utils/foreign_attributes.hpp"
 
@@ -21,15 +22,13 @@ public:
   fk_attribute_generator() = default;
 
   template<class Type>
-  attribute_definition generate(const char *id, Type &x, const std::shared_ptr<attribute_definition> &ref_column)
-  {
+  attribute_definition generate(const char *id, Type &x, const std::shared_ptr<attribute_definition> &ref_column) {
     access::process(*this, x);
     return attribute_definition{id, type_, 0, ref_column, {utils::constraints::FOREIGN_KEY }, null_option::NOT_NULL};
   }
 
   template<typename ValueType>
-  void on_primary_key(const char *, ValueType &/*pk*/, std::enable_if_t<std::is_integral_v<ValueType> && !std::is_same_v<bool, ValueType>>* = nullptr)
-  {
+  void on_primary_key(const char *, ValueType &/*pk*/, std::enable_if_t<std::is_integral_v<ValueType> && !std::is_same_v<bool, ValueType>>* = nullptr) {
     type_ = utils::data_type_traits<ValueType>::type(0);
   }
   void on_primary_key(const char * /*id*/, std::string &/*pk*/, size_t size);
@@ -81,16 +80,16 @@ public:
   void on_attribute(const char *id, std::optional<Type> &x, const utils::field_attributes &attr = utils::null_attributes);
 
   template<class Pointer>
-  void on_belongs_to(const char *id, Pointer &x, const utils::foreign_attributes &/*attr*/)
-  {
-    const auto ref_column = determine_foreign_ref(std::type_index(typeid(typename Pointer::value_type)));
-    columns_.push_back(fk_column_generator_.generate(id, *x, ref_column));
+  void on_belongs_to(const char *id, Pointer &x, const utils::foreign_attributes &/*attr*/) {
+    if (const auto result = determine_foreign_ref(std::type_index(typeid(typename Pointer::value_type)))) {
+      columns_.push_back(fk_column_generator_.generate(id, *x, *result));
+    }
   }
   template<class Pointer>
-  void on_has_one(const char *id, Pointer &x, const utils::foreign_attributes &/*attr*/)
-  {
-    const auto ref_column = determine_foreign_ref(std::type_index(typeid(typename Pointer::value_type)));
-    columns_.push_back(fk_column_generator_.generate(id, *x, ref_column));
+  void on_has_one(const char *id, Pointer &x, const utils::foreign_attributes &/*attr*/) {
+    if (const auto result = determine_foreign_ref(std::type_index(typeid(typename Pointer::value_type)))) {
+      columns_.push_back(fk_column_generator_.generate(id, *x, *result));
+    }
   }
   template<class ContainerType>
   void on_has_many(const char *id, ContainerType &, const char *, const utils::foreign_attributes &/*attr*/) {}
@@ -100,7 +99,7 @@ public:
   void on_has_many_to_many(const char *id, ContainerType &c, const utils::foreign_attributes &/*attr*/) {}
 
 private:
-  std::shared_ptr<attribute_definition> determine_foreign_ref(const std::type_index &ti);
+  [[nodiscard]] utils::result<std::shared_ptr<attribute_definition>, utils::error> determine_foreign_ref(const std::type_index &ti) const;
 
 private:
   size_t index_ = 0;
diff --git a/include/matador/object/basic_object_info.hpp b/include/matador/object/basic_object_info.hpp
index 8e2dd7d..7853f1d 100644
--- a/include/matador/object/basic_object_info.hpp
+++ b/include/matador/object/basic_object_info.hpp
@@ -2,6 +2,7 @@
 #define BASIC_PROTOTYPE_INFO_HPP
 
 #include "matador/object/object_definition.hpp"
+#include "matador/utils/identifier.hpp"
 
 #include <string>
 #include <typeindex>
@@ -20,14 +21,15 @@ public:
   [[nodiscard]] std::shared_ptr<attribute_definition> reference_column() const;
 
 protected:
-  basic_object_info(std::type_index type_index, object_definition &&definition, std::shared_ptr<attribute_definition> &&reference_column);
-  basic_object_info(std::type_index type_index, std::shared_ptr<attribute_definition> &&reference_column);
+  basic_object_info(std::shared_ptr<schema_node> node, std::type_index type_index, utils::identifier &&pk, std::shared_ptr<attribute_definition> &&pk_column, object_definition &&definition);
+  basic_object_info(std::shared_ptr<schema_node> node, std::type_index type_index, utils::identifier &&pk, std::shared_ptr<attribute_definition> &&pk_column);
 
 protected:
   std::shared_ptr<schema_node> node_;            /**< prototype node of the represented object type */
   std::type_index type_index_;   /**< type index of the represented object type */
   object_definition definition_;
-  std::shared_ptr<attribute_definition> reference_column_;
+  std::optional<utils::identifier> identifier_;
+  std::shared_ptr<attribute_definition> pk_column_;
 };
 
 using basic_object_info_ref = std::reference_wrapper<const basic_object_info>;
diff --git a/include/matador/object/object_info.hpp b/include/matador/object/object_info.hpp
index cb32a03..ab6d2cc 100644
--- a/include/matador/object/object_info.hpp
+++ b/include/matador/object/object_info.hpp
@@ -10,13 +10,15 @@ class schema_node;
 template<typename Type>
 class object_info final : public basic_object_info {
 public:
-  explicit object_info(object_definition &&definition,
+  explicit object_info(const std::shared_ptr<schema_node>& node,
                        std::shared_ptr<attribute_definition> &&ref_column)
-  : basic_object_info(typeid(Type), std::move(definition), std::move(ref_column)) {
+  : basic_object_info(node, typeid(Type), {}, std::move(ref_column), {}) {
   }
-  explicit object_info(schema_node &node,
-                       std::shared_ptr<attribute_definition> &&ref_column)
-  : basic_object_info(typeid(Type), std::move(ref_column)) {
+  explicit object_info(const std::shared_ptr<schema_node>& node,
+                       utils::identifier &&pk,
+                       std::shared_ptr<attribute_definition> &&ref_column,
+                       object_definition &&definition)
+  : basic_object_info(node, typeid(Type), std::move(pk), std::move(ref_column), std::move(definition)) {
   }
 
   const Type &prototype() const { return prototype_; }
diff --git a/include/matador/object/primary_key_resolver.hpp b/include/matador/object/primary_key_resolver.hpp
index c50486b..f108ba4 100644
--- a/include/matador/object/primary_key_resolver.hpp
+++ b/include/matador/object/primary_key_resolver.hpp
@@ -14,7 +14,8 @@
 namespace matador::object {
 
 struct primary_key_info {
-  std::string column_name;
+  std::string pk_column_name;
+  utils::basic_type type;
   utils::identifier pk;
 };
 
@@ -36,25 +37,26 @@ public:
 
   template < class Type >
   void on_primary_key(const char *id, Type &pk, std::enable_if_t<std::is_integral_v<Type> && !std::is_same_v<bool, Type>>* = nullptr) {
-    primary_key_info_.column_name = id;
+    primary_key_info_.pk_column_name = id;
+    primary_key_info_.type = utils::data_type_traits<Type>::type();
     primary_key_info_.pk = pk;
   }
 
   void on_primary_key(const char *id, const std::string &pk, size_t size);
 
-  void on_revision(const char * /*id*/, uint64_t &/*rev*/) {}
+  static void on_revision(const char * /*id*/, uint64_t &/*rev*/) {}
   template<typename Type>
-  void on_attribute(const char * /*id*/, Type &/*val*/, const utils::field_attributes &/*attr*/ = utils::null_attributes) {}
+  static void on_attribute(const char * /*id*/, Type &/*val*/, const utils::field_attributes &/*attr*/ = utils::null_attributes) {}
   template<class Pointer>
-  void on_belongs_to(const char * /*id*/, Pointer &/*val*/, const utils::foreign_attributes &/*attr*/) {}
+  static void on_belongs_to(const char * /*id*/, Pointer &/*val*/, const utils::foreign_attributes &/*attr*/) {}
   template<class Pointer>
-  void on_has_one(const char * /*id*/, Pointer &/*val*/, const utils::foreign_attributes &/*attr*/) {}
+  static void on_has_one(const char * /*id*/, Pointer &/*val*/, const utils::foreign_attributes &/*attr*/) {}
   template<class ContainerType>
-  void on_has_many(const char * /*id*/, ContainerType &/*col*/, const char *, const utils::foreign_attributes &/*attr*/) {}
+  static void on_has_many(const char * /*id*/, ContainerType &/*col*/, const char *, const utils::foreign_attributes &/*attr*/) {}
   template<class ContainerType>
-  void on_has_many_to_many(const char * /*id*/, ContainerType &/*col*/, const char * /*join_column*/, const char * /*inverse_join_column*/, const utils::foreign_attributes &/*attr*/) {}
+  static void on_has_many_to_many(const char * /*id*/, ContainerType &/*col*/, const char * /*join_column*/, const char * /*inverse_join_column*/, const utils::foreign_attributes &/*attr*/) {}
   template<class ContainerType>
-  void on_has_many_to_many(const char * /*id*/, ContainerType &/*col*/, const utils::foreign_attributes &/*attr*/) {}
+  static void on_has_many_to_many(const char * /*id*/, ContainerType &/*col*/, const utils::foreign_attributes &/*attr*/) {}
 
 private:
   primary_key_info primary_key_info_{};
diff --git a/include/matador/object/schema.hpp b/include/matador/object/schema.hpp
index b47f817..3269e40 100644
--- a/include/matador/object/schema.hpp
+++ b/include/matador/object/schema.hpp
@@ -1,6 +1,7 @@
 #ifndef SCHEMA_HPP
 #define SCHEMA_HPP
 
+#include "primary_key_resolver.hpp"
 #include "matador/object/error_code.hpp"
 #include "matador/object/schema_node.hpp"
 #include "matador/object/schema_node_iterator.hpp"
@@ -24,14 +25,11 @@ class type_analyzer final {
 public:
   using value_type = Type;
 
-  static void analyze(schema &scm) {
+  static void prepare_expected_nodes(schema &scm) {
     type_analyzer analyzer(scm);
 
     Type obj;
-    const auto ti = std::type_index(typeid(Type));
-    analyzer.known_types_.emplace(ti);
     access::process(analyzer, obj);
-    analyzer.known_types_.erase(ti);
   }
 
   template < class PrimaryKeyType >
@@ -47,12 +45,12 @@ public:
 
   template<class ForeignPointerType>
   void on_belongs_to(const char * /*id*/, ForeignPointerType &/*obj*/, const utils::foreign_attributes &/*attr*/) {
-    // const auto [ref_table, ref_column] = determine_foreign_ref(std::type_index(typeid(typename Pointer::value_type)));
-    // columns_.push_back(fk_column_generator_.generate(id, *x, ref_table, ref_column));
+    on_foreign_key<ForeignPointerType>();
   }
-
   template<class ForeignPointerType>
-  void on_has_one(const char * /*id*/, ForeignPointerType &/*obj*/, const utils::foreign_attributes &/*attr*/);
+  void on_has_one(const char * /*id*/, ForeignPointerType &/*obj*/, const utils::foreign_attributes &/*attr*/) {
+    on_foreign_key<ForeignPointerType>();
+  }
 
   template<class CollectionType>
   void on_has_many(const char * /*id*/, CollectionType &, const char *, const utils::foreign_attributes &/*attr*/) {}
@@ -61,6 +59,10 @@ public:
   template<class ContainerType>
   void on_has_many_to_many(const char * /*id*/, ContainerType &/*col*/, const utils::foreign_attributes &/*attr*/) {}
 
+private:
+  template<class ForeignPointerType>
+  void on_foreign_key();
+
 private:
   explicit type_analyzer(schema& schema)
   : schema_(schema) {};
@@ -68,7 +70,6 @@ private:
 
 private:
   schema &schema_;
-  std::unordered_set<std::type_index> known_types_;
 };
 
 
@@ -79,23 +80,30 @@ public:
   /**
    * Creates an empty schema
    */
-  explicit schema( const std::string& name = "");
+  explicit schema( std::string  name = "");
 
   template <typename Type>
   [[nodiscard]] utils::result<void, utils::error> attach(const std::string& name, const std::string &parent = "") {
     if (has_node(name)) {
       return utils::failure(make_error(error_code::NodeAlreadyExists, "Node '" + name + "' already exists"));
     }
+    if (const auto it = expected_node_map_.find(typeid(Type)); it != expected_node_map_.end()) {
+      const auto node = it->second;
+      expected_node_map_.erase(it);
 
-    // analyze node (collect unknown types by type index)
-    type_analyzer<Type>::analyze(*this);
+      node->update_name(name);
 
-    const auto node = schema_node::make_node(*this, name, {});
+      node_map_.insert({node->name(), node})/*.first*/;
+      type_index_node_map_.insert({node->type_index(), node});
+    } else {
+      // analyze node (collect unknown types by type index)
+      type_analyzer<Type>::prepare_expected_nodes(*this);
 
-    if (auto result = attach_node(node, parent); !result) {
-      return utils::failure(result.err());
+      const auto node = schema_node::make_node<Type>(*this, name);
+      if (auto result = attach_node(node, parent); !result) {
+        return utils::failure(result.err());
+      }
     }
-
     return utils::ok<void>();
   }
 
@@ -195,22 +203,18 @@ private:
 
   t_node_map node_map_;
   t_type_index_node_map type_index_node_map_;
+  t_type_index_node_map expected_node_map_;
 };
 
 template<typename Type>
 template<class ForeignPointerType>
-void type_analyzer<Type>::on_has_one(const char *, ForeignPointerType &, const utils::foreign_attributes &) {
+void type_analyzer<Type>::on_foreign_key() {
   auto ti = std::type_index(typeid(typename ForeignPointerType::value_type));
-  if (schema_.has_node(ti)) {
+  if (schema_.has_node(ti) || schema_.expected_node_map_.count(ti) > 0) {
     return;
   }
 
-  auto info = std::make_unique<object_info<typename ForeignPointerType::value_type>>(object_definition{}, std::make_unique<attribute_definition>(std::string(""), std::string(ti.name()), utils::basic_type::type_null, utils::constraints::FOREIGN_KEY));
-
-  known_types_.insert(ti);
-
-  // const auto [ref_table, ref_column] = determine_foreign_ref(std::type_index(typeid(typename Pointer::value_type)));
-  // columns_.push_back(fk_column_generator_.generate(id, *x, ref_table, ref_column));
+  schema_.expected_node_map_.insert({ti, schema_node::make_node<typename ForeignPointerType::value_type>(schema_, ti.name())});
 }
 
 }
diff --git a/include/matador/object/schema_node.hpp b/include/matador/object/schema_node.hpp
index 48a286b..5c0bad6 100644
--- a/include/matador/object/schema_node.hpp
+++ b/include/matador/object/schema_node.hpp
@@ -3,6 +3,7 @@
 
 #include "matador/object/attribute_definition_generator.hpp"
 #include "matador/object/object_info.hpp"
+#include "matador/object/primary_key_resolver.hpp"
 
 #include <memory>
 
@@ -15,7 +16,24 @@ class schema_node final {
 public:
   using node_ptr = std::shared_ptr<schema_node>;
 
-  static std::shared_ptr<schema_node> make_node(schema& tree, const std::string& name, std::unique_ptr<basic_object_info> &&info);
+  template < typename Type >
+  static std::shared_ptr<schema_node> make_node(schema& tree, const std::string& name) {
+    auto node = std::shared_ptr<schema_node>(new schema_node(tree, name));
+
+    primary_key_resolver resolver;
+    auto pk_info = resolver.resolve<Type>();
+    auto info = std::make_unique<object_info<Type>>(
+      node,
+      std::move(pk_info.pk),
+      std::make_shared<attribute_definition>(pk_info.pk_column_name, name, pk_info.type, utils::constraints::FOREIGN_KEY),
+      object_definition{attribute_definition_generator::generate<Type>(tree)}
+    );
+    node->info_ = std::move(info);
+
+    return node;
+  }
+
+  static std::shared_ptr<schema_node> make_null_node(schema& tree);
 
   schema_node(const schema_node& other) = delete;
   schema_node(schema_node&& other) = default;
@@ -31,14 +49,17 @@ public:
 
   [[nodiscard]] const basic_object_info& basic_info() const;
 
+  void update_name(const std::string& name);
+
   template <typename Type>
   object_info_ref<Type> info() const {
+      auto ref = std::ref(static_cast<const object_info<Type>&>(*info_));
       return std::ref(static_cast<const object_info<Type>&>(*info_));
   }
 
 private:
   explicit schema_node(schema& tree);
-  schema_node(schema& tree, std::string name, std::unique_ptr<basic_object_info> &&info);
+  schema_node(schema& tree, std::string name);
 
 private:
   friend class schema;
diff --git a/include/matador/sql/column_generator.hpp b/include/matador/sql/column_generator.hpp
index 835f7e3..36f9e4f 100644
--- a/include/matador/sql/column_generator.hpp
+++ b/include/matador/sql/column_generator.hpp
@@ -41,22 +41,19 @@ public:
   }
 
   template < class V >
-  void on_primary_key(const char *id, V &, std::enable_if_t<std::is_integral_v<V> && !std::is_same_v<bool, V>>* = nullptr)
-  {
+  void on_primary_key(const char *id, V &, std::enable_if_t<std::is_integral_v<V> && !std::is_same_v<bool, V>>* = nullptr) {
     push(id);
   }
   void on_primary_key(const char *id, std::string &, size_t);
   void on_revision(const char *id, unsigned long long &/*rev*/);
 
   template<typename Type>
-  void on_attribute(const char *id, Type &, const utils::field_attributes &/*attr*/ = utils::null_attributes)
-  {
+  void on_attribute(const char *id, Type &, const utils::field_attributes &/*attr*/ = utils::null_attributes) {
     push(id);
   }
 
   template<class Pointer>
-  void on_belongs_to(const char *id, Pointer &, const utils::foreign_attributes &attr)
-  {
+  void on_belongs_to(const char *id, Pointer &, const utils::foreign_attributes &attr) {
     if (attr.fetch() == utils::fetch_type::LAZY || force_lazy_) {
       push(id);
     } else {
diff --git a/source/core/object/attribute_definition.cpp b/source/core/object/attribute_definition.cpp
index 05ecabe..1442cbd 100644
--- a/source/core/object/attribute_definition.cpp
+++ b/source/core/object/attribute_definition.cpp
@@ -65,6 +65,10 @@ const std::string &attribute_definition::name() const {
   return name_;
 }
 
+void attribute_definition::name( const std::string& n ) {
+  name_ = n;
+}
+
 std::string attribute_definition::full_name() const {
   return table_ + "." + name_;
 }
diff --git a/source/core/object/attribute_definition_generator.cpp b/source/core/object/attribute_definition_generator.cpp
index d700816..1a27b63 100644
--- a/source/core/object/attribute_definition_generator.cpp
+++ b/source/core/object/attribute_definition_generator.cpp
@@ -18,9 +18,8 @@ void attribute_definition_generator::on_revision(const char *id, uint64_t &rev)
     on_attribute(id, rev);
 }
 
-std::shared_ptr<attribute_definition> attribute_definition_generator::determine_foreign_ref(const std::type_index &ti)
-{
-    return repo_.reference(ti).value();
+utils::result<std::shared_ptr<attribute_definition>, utils::error> attribute_definition_generator::determine_foreign_ref(const std::type_index &ti) const {
+    return repo_.reference(ti);
 }
 
 void fk_attribute_generator::on_primary_key(const char *, std::string &, const size_t size)
diff --git a/source/core/object/basic_object_info.cpp b/source/core/object/basic_object_info.cpp
index 2422314..25b44dd 100644
--- a/source/core/object/basic_object_info.cpp
+++ b/source/core/object/basic_object_info.cpp
@@ -5,18 +5,26 @@
 #include <algorithm>
 
 namespace matador::object {
-basic_object_info::basic_object_info(const std::type_index type_index,
-                                     object_definition &&definition,
-                                     std::shared_ptr<attribute_definition> &&reference_column)
-: type_index_(type_index)
+basic_object_info::basic_object_info(std::shared_ptr<schema_node> node,
+                                     const std::type_index type_index,
+                                     utils::identifier &&pk,
+                                     std::shared_ptr<attribute_definition> &&pk_column,
+                                     object_definition &&definition)
+: node_(std::move(node))
+, type_index_(type_index)
 , definition_(std::move(definition))
-, reference_column_(std::move(reference_column)) {
+, identifier_(std::move(pk))
+, pk_column_(std::move(pk_column)) {
 }
 
-basic_object_info::basic_object_info(const std::type_index type_index,
-                                     std::shared_ptr<attribute_definition> &&reference_column)
-: type_index_(type_index)
-, reference_column_(std::move(reference_column)) {
+basic_object_info::basic_object_info(std::shared_ptr<schema_node> node,
+                                     const std::type_index type_index,
+                                     utils::identifier &&pk,
+                                     std::shared_ptr<attribute_definition> &&pk_column)
+: node_(std::move(node))
+, type_index_(type_index)
+, identifier_(std::move(pk))
+, pk_column_(std::move(pk_column)) {
 }
 
 std::type_index basic_object_info::type_index() const {
@@ -32,7 +40,7 @@ const object_definition &basic_object_info::definition() const {
 }
 
 std::shared_ptr<attribute_definition> basic_object_info::reference_column() const {
-  return reference_column_;
+  return pk_column_;
 }
 
 } // namespace matador::object
diff --git a/source/core/object/primary_key_resolver.cpp b/source/core/object/primary_key_resolver.cpp
index 11e0f31..9e3eb71 100644
--- a/source/core/object/primary_key_resolver.cpp
+++ b/source/core/object/primary_key_resolver.cpp
@@ -2,7 +2,9 @@
 
 namespace matador::object {
 void primary_key_resolver::on_primary_key(const char *id, const std::string &pk, size_t /*size*/) {
-  primary_key_info_.column_name = id;
+  primary_key_info_.pk_column_name = id;
+  primary_key_info_.type = utils::basic_type::type_varchar;
   primary_key_info_.pk = pk;
 }
+
 }
diff --git a/source/core/object/schema.cpp b/source/core/object/schema.cpp
index 1a41ef7..7bcab40 100644
--- a/source/core/object/schema.cpp
+++ b/source/core/object/schema.cpp
@@ -1,3 +1,5 @@
+#include <utility>
+
 #include "matador/object/schema.hpp"
 
 namespace matador::object {
@@ -5,13 +7,13 @@ utils::error make_error(const error_code ec, const std::string &msg) {
   return utils::error(ec, msg);
 }
 
-schema::schema(const std::string &name)
-  : name_(name)
-    , root_(std::shared_ptr<schema_node>(new schema_node(*this))) {
-  root_->first_child_ = std::shared_ptr<schema_node>(new schema_node(*this));
-  root_->last_child_ = std::shared_ptr<schema_node>(new schema_node(*this));
-  root_->first_child_->next_sibling_ = root_->last_child_;
-  root_->last_child_->previous_sibling_ = root_->first_child_;
+schema::schema(std::string name)
+: name_(std::move(name))
+, root_(schema_node::make_null_node(*this)) {
+root_->first_child_ = std::shared_ptr<schema_node>(new schema_node(*this));
+root_->last_child_ = std::shared_ptr<schema_node>(new schema_node(*this));
+root_->first_child_->next_sibling_ = root_->last_child_;
+root_->last_child_->previous_sibling_ = root_->first_child_;
 }
 
 schema::const_iterator schema::begin() const {
@@ -36,16 +38,15 @@ std::string schema::name() const {
 
 utils::result<std::shared_ptr<attribute_definition>, utils::error> schema::reference(const std::type_index &type_index) const {
   const auto result = find_node(type_index);
-  if (!result) {
-    return utils::failure(result.err());
+  if (result) {
+    return utils::ok((*result)->basic_info().reference_column());
   }
 
-  const auto &node = *result;
-  if (!node->basic_info().definition().has_primary_key()) {
-    return utils::failure(make_error(error_code::Failure, "Primary key not found"));
+  if (const auto it = expected_node_map_.find(type_index); it != expected_node_map_.end()) {
+    return utils::ok(it->second->basic_info().reference_column());
   }
 
-  return utils::ok(std::make_shared<attribute_definition>(node->name(), node->basic_info().definition().primary_key()->name(), utils::basic_type::type_null, utils::constraints::FOREIGN_KEY));
+  return utils::failure(result.err());
 }
 
 utils::result<std::shared_ptr<schema_node>, utils::error> schema::attach_node(const std::shared_ptr<schema_node> &node,
diff --git a/source/core/object/schema_node.cpp b/source/core/object/schema_node.cpp
index e330240..814bc57 100644
--- a/source/core/object/schema_node.cpp
+++ b/source/core/object/schema_node.cpp
@@ -4,13 +4,11 @@
 
 namespace matador::object {
 schema_node::schema_node(schema &tree)
-: schema_(tree)
-, info_(std::make_unique<null_info>(object_definition{}, std::shared_ptr<attribute_definition>{})) {
+: schema_(tree) {
 }
 
-schema_node::schema_node(schema &tree, std::string name, std::unique_ptr<basic_object_info> &&info)
+schema_node::schema_node(schema &tree, std::string name)
 : schema_(tree)
-, info_(std::move(info))
 , first_child_(std::shared_ptr<schema_node>(new schema_node(tree)))
 , last_child_(std::shared_ptr<schema_node>(new schema_node(tree)))
 , name_(std::move(name)) {
@@ -18,8 +16,11 @@ schema_node::schema_node(schema &tree, std::string name, std::unique_ptr<basic_o
   last_child_->previous_sibling_ = first_child_;
 }
 
-std::shared_ptr<schema_node> schema_node::make_node(schema &tree, const std::string &name, std::unique_ptr<basic_object_info> &&info) {
-  return std::shared_ptr<schema_node>(new schema_node(tree, name, std::move(info)));
+std::shared_ptr<schema_node> schema_node::make_null_node(schema &tree) {
+  auto node = std::shared_ptr<schema_node>(new schema_node(tree));
+  node->info_ = std::make_unique<null_info>(node, std::shared_ptr<attribute_definition>{});
+
+  return node;
 }
 
 std::string schema_node::name() const {
@@ -34,6 +35,11 @@ const basic_object_info &schema_node::basic_info() const {
   return *info_;
 }
 
+void schema_node::update_name(const std::string& name) {
+  name_ = name;
+  info_->reference_column()->name(name);
+}
+
 schema_node::node_ptr schema_node::next() const {
   // if we have a child, child is the next iterator to return
   // (if we don't do iterate over the siblings)
diff --git a/test/core/CMakeLists.txt b/test/core/CMakeLists.txt
index 2ec9470..3857c92 100644
--- a/test/core/CMakeLists.txt
+++ b/test/core/CMakeLists.txt
@@ -12,7 +12,8 @@ add_executable(CoreTests
   utils/FieldAttributeTest.cpp
   utils/VersionTest.cpp
   utils/StringTest.cpp
-  object/ColumnDefinitionGeneratorTest.cpp
+  object/AttributeDefinitionGeneratorTest.cpp
+  object/PrimaryKeyResolverTest.cpp
   object/SchemaTest.cpp
 )
 
diff --git a/test/core/object/ColumnDefinitionGeneratorTest.cpp b/test/core/object/AttributeDefinitionGeneratorTest.cpp
similarity index 91%
rename from test/core/object/ColumnDefinitionGeneratorTest.cpp
rename to test/core/object/AttributeDefinitionGeneratorTest.cpp
index 70c1b62..7202d84 100644
--- a/test/core/object/ColumnDefinitionGeneratorTest.cpp
+++ b/test/core/object/AttributeDefinitionGeneratorTest.cpp
@@ -12,6 +12,11 @@ using namespace matador::utils;
 TEST_CASE("Generate column definitions from object", "[column][definition][generator]") {
   schema repo("main");
 
+  auto result = repo.attach<matador::test::product>("products")
+    .and_then([&repo] { return repo.attach<matador::test::supplier>("supplier"); })
+    .and_then([&repo] { return repo.attach<matador::test::category>("categories"); });
+  REQUIRE(result);
+
   auto columns = attribute_definition_generator::generate<matador::test::product>(repo);
 
   const std::vector expected_columns = {
diff --git a/test/core/object/PrimaryKeyResolverTest.cpp b/test/core/object/PrimaryKeyResolverTest.cpp
new file mode 100644
index 0000000..cd1bc74
--- /dev/null
+++ b/test/core/object/PrimaryKeyResolverTest.cpp
@@ -0,0 +1,22 @@
+#include <catch2/catch_test_macros.hpp>
+
+#include "matador/object/primary_key_resolver.hpp"
+
+#include "../test/models/book.hpp"
+#include "../test/models/product.hpp"
+
+TEST_CASE("Test primary key resolver", "[PrimaryKeyResolver]") {
+    using namespace matador;
+
+    object::primary_key_resolver resolver;
+
+    auto pk_info = resolver.resolve<test::book>();
+
+    REQUIRE(pk_info.pk_column_name == "id");
+    REQUIRE(pk_info.pk.is_integer());
+
+    pk_info = resolver.resolve<test::product>();
+
+    REQUIRE(pk_info.pk_column_name == "product_name");
+    REQUIRE(pk_info.pk.is_varchar());
+}
\ No newline at end of file
diff --git a/test/orm/orm/SessionQueryBuilderTest.cpp b/test/orm/orm/SessionQueryBuilderTest.cpp
index e16fc2a..152233b 100644
--- a/test/orm/orm/SessionQueryBuilderTest.cpp
+++ b/test/orm/orm/SessionQueryBuilderTest.cpp
@@ -133,6 +133,8 @@ TEST_CASE("Create sql query data for entity with eager has many belongs to", "[q
   schema scm("noop");
   auto result = scm.attach<product>("products")
     .and_then( [&scm] { return scm.attach<order_details>("order_details"); } )
+    .and_then( [&scm] { return scm.attach<supplier>("suppliers"); } )
+    .and_then( [&scm] { return scm.attach<category>("categories"); } )
     .and_then( [&scm] { return scm.attach<order>("orders"); } );
   REQUIRE(result);