query/include/matador/orm/session_query_builder.hpp

#ifndef QUERY_ENTITY_QUERY_BUILDER_HPP
#define QUERY_ENTITY_QUERY_BUILDER_HPP

#include "matador/orm/query_builder_exception.hpp"

#include "matador/query/criteria.hpp"
#include "matador/query/query.hpp"

#include "matador/sql/executor.hpp"
#include "matador/sql/statement.hpp"

#include "matador/object/join_columns_collector.hpp"
#include "matador/object/repository.hpp"

#include "matador/utils/primary_key_attribute.hpp"
#include "matador/utils/result.hpp"
#include "matador/utils/value.hpp"

#include <stack>
#include <unordered_map>

namespace matador::orm {

struct entity_query_data {
  std::shared_ptr<sql::table> root_table;
  std::string pk_column_name{};
  std::vector<sql::column> columns{};
  std::unordered_map<std::string, sql::statement> lazy_loading_statements{};
  std::vector<query::join_data> joins{};
  query::criteria_ptr where_clause;
};

class session_query_builder final {
public:
  session_query_builder(const object::repository &scm, sql::executor &exec)
  : schema_(scm)
  , executor_(exec){}

  template<class EntityType, typename PrimaryKeyType>
  utils::result<entity_query_data, query_build_error> build(const PrimaryKeyType &pk) {
    auto info = schema_.info<EntityType>();
    if (!info) {
      return utils::failure(query_build_error::UnknownType);
    }
    pk_ = pk;
    table_info_stack_.push({info.value(), std::make_shared<sql::table>(info.value().get().name(), build_alias('t', ++table_index))});
    entity_query_data_ = { table_info_stack_.top().table };
    processed_tables_.insert({info->get().name(), entity_query_data_.root_table});
    try {
      access::process(*this, info->get().prototype());

      return {utils::ok(std::move(entity_query_data_))};
    } catch (const query_builder_exception &ex) {
      return {utils::failure(ex.error_type())};
    } catch (...) {
      return {utils::failure(query_build_error::UnexpectedError)};
    }
  }

  template<class EntityType>
  utils::result<entity_query_data, query_build_error> build() {
    const auto info = schema_.info<EntityType>();
    if (!info) {
      return utils::failure(query_build_error::UnknownType);
    }
    pk_ = nullptr;
    table_info_stack_.push({info.value(), std::make_shared<sql::table>(info.value().get().name(), build_alias('t', ++table_index))});
    entity_query_data_ = { table_info_stack_.top().table };
    processed_tables_.insert({info->get().name(), entity_query_data_.root_table});
    try {
      access::process(*this, info->get().prototype());

      return {utils::ok(std::move(entity_query_data_))};
    } catch (const query_builder_exception &ex) {
      return {utils::failure(ex.error_type())};
    } catch (...) {
      return {utils::failure(query_build_error::UnexpectedError)};
    }
  }

  template < class V >
  void on_primary_key(const char *id, V &, const utils::primary_key_attribute& /*attr*/ = utils::default_pk_attributes)
  {
    push(id);
    if (!is_root_entity()) {
      return;
    }
    entity_query_data_.pk_column_name = id;
    if (!pk_.is_null()) {
      const auto c = sql::column{table_info_stack_.top().table, id, ""};
      using namespace matador::query;
      auto co = c == utils::_;
      entity_query_data_.where_clause = std::move(co);
    }
  }

  void on_revision(const char *id, uint64_t &/*rev*/);

  template<typename Type>
  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 &obj, const utils::foreign_attributes &attr)
  {
    on_foreign_object(id, obj, attr);
  }

  template<class Pointer>
  void on_has_one(const char *id, Pointer &obj, const utils::foreign_attributes &attr)
  {
    on_foreign_object(id, obj, attr);
  }

  template<class ContainerType>
  void on_has_many(const char * /*id*/, ContainerType &, const char *join_column, const utils::foreign_attributes &attr) {
    if (attr.fetch() == utils::fetch_type::EAGER) {
      const auto info = schema_.info<typename ContainerType::value_type::value_type>();
      if (!info) {
        throw query_builder_exception{query_build_error::UnknownType};
      }

      auto next = processed_tables_.find(info->get().name());
      if (next != processed_tables_.end()) {
        return;
      }
      table_info_stack_.push({info.value(), std::make_shared<sql::table>(info->get().name(), build_alias('t', ++table_index))});
      next = processed_tables_.insert({info->get().name(), table_info_stack_.top().table}).first;
      typename ContainerType::value_type::value_type obj;
      access::process(*this , obj);
      table_info_stack_.pop();

      auto pk = info->get().definition().primary_key();
      if (!pk) {
        throw query_builder_exception{query_build_error::MissingPrimaryKey};
      }

      append_join(
        sql::column{table_info_stack_.top().table, table_info_stack_.top().info.get().definition().primary_key()->name()},
        sql::column{next->second, join_column}
      );
    }
  }

  template<class ContainerType>
  void on_has_many_to_many(const char *id, ContainerType &/*cont*/, const char *join_column, const char *inverse_join_column, const utils::foreign_attributes &attr) {
    if (attr.fetch() != utils::fetch_type::EAGER) {
      return;
    }
    const auto info = schema_.info<typename ContainerType::value_type::value_type>();
    if (!info) {
      throw query_builder_exception{query_build_error::UnknownType};
    }

    auto next = processed_tables_.find(info->get().name());
    if (next != processed_tables_.end()) {
      return;
    }
    auto relation = processed_tables_.find(id);
    if (relation == processed_tables_.end()) {
      relation = processed_tables_.insert({id, std::make_shared<sql::table>(id, build_alias('t', ++table_index))}).first;
    }

    table_info_stack_.push({info.value(), std::make_shared<sql::table>(info->get().name(), build_alias('t', ++table_index))});
    next = processed_tables_.insert({info->get().name(), table_info_stack_.top().table}).first;
    typename ContainerType::value_type::value_type obj;
    access::process(*this , obj);
    table_info_stack_.pop();

    auto pk = info->get().definition().primary_key();
    if (!pk) {
      throw query_builder_exception{query_build_error::MissingPrimaryKey};
    }

    append_join(
      sql::column{table_info_stack_.top().table, table_info_stack_.top().info.get().definition().primary_key()->name()},
      sql::column{relation->second, join_column}
    );
    append_join(
      sql::column{relation->second, inverse_join_column},
      sql::column{next->second, pk->name()}
    );
  }

  template<class ContainerType>
  void on_has_many_to_many(const char *id, ContainerType &/*cont*/, const utils::foreign_attributes &attr) {
    if (attr.fetch() != utils::fetch_type::EAGER) {
      return;
    }
    const auto info = schema_.info<typename ContainerType::value_type::value_type>();
    if (!info) {
      throw query_builder_exception{query_build_error::UnknownType};
    }

    auto next = processed_tables_.find(info->get().name());
    if (next != processed_tables_.end()) {
      return;
    }

    auto relation = processed_tables_.find(id);
    if (relation == processed_tables_.end()) {
      relation = processed_tables_.insert({id, std::make_shared<sql::table>(id, build_alias('t', ++table_index))}).first;
    }
    table_info_stack_.push({info.value(), std::make_shared<sql::table>(info->get().name(), build_alias('t', ++table_index))});
    next = processed_tables_.insert({info->get().name(), table_info_stack_.top().table}).first;
    typename ContainerType::value_type::value_type obj;
    access::process(*this , obj);
    table_info_stack_.pop();

    auto pk = info->get().definition().primary_key();
    if (!pk) {
      throw query_builder_exception{query_build_error::MissingPrimaryKey};
    }

    const auto join_columns = join_columns_collector_.collect<typename ContainerType::value_type::value_type>();

    append_join(
      sql::column{table_info_stack_.top().table, table_info_stack_.top().info.get().definition().primary_key()->name()},
      sql::column{relation->second, join_columns.inverse_join_column}
    );
    append_join(
      sql::column{relation->second, join_columns.join_column},
      sql::column{next->second, pk->name()}
    );
  }

private:
  template<class Pointer>
  void on_foreign_object(const char *id, Pointer &, const utils::foreign_attributes &attr);
  void push(const std::string &column_name);
  static std::string build_alias(char prefix, unsigned int count);
  [[nodiscard]] bool is_root_entity() const;
  void append_join(const sql::column &left, const sql::column &right);

private:
  utils::value pk_;
  struct table_info {
    std::reference_wrapper<const object::basic_object_info> info;
    std::shared_ptr<sql::table> table;
  };

  std::stack<table_info> table_info_stack_{};
  std::unordered_map<std::string, std::shared_ptr<sql::table>> processed_tables_{};
  const object::repository &schema_;
  entity_query_data entity_query_data_{};
  unsigned int column_index{0};
  unsigned int table_index{0};
  object::join_columns_collector join_columns_collector_{};
  sql::executor &executor_;
};

template<class Pointer>
void session_query_builder::on_foreign_object(const char *id, Pointer &, const utils::foreign_attributes &attr) {
  const auto info = schema_.info<typename Pointer::value_type>();
  if (!info) {
    throw query_builder_exception{query_build_error::UnknownType};
  }
  auto pk = info->get().definition().primary_key();
  if (!pk) {
    throw query_builder_exception{query_build_error::MissingPrimaryKey};
  }

  const auto foreign_table = std::make_shared<sql::table>(info->get().name(), build_alias('t', ++table_index));
  if (attr.fetch() == utils::fetch_type::EAGER) {

    auto next = processed_tables_.find(info->get().name());
    if (next != processed_tables_.end()) {
      return;
    }
    table_info_stack_.push({info.value(), foreign_table});
    next = processed_tables_.insert({info->get().name(), table_info_stack_.top().table}).first;
    typename Pointer::value_type obj;
    access::process(*this, obj);
    table_info_stack_.pop();

    append_join(
      sql::column{table_info_stack_.top().table, id},
      sql::column{next->second, pk->name()}
    );
  } else {
    push(id);
    using namespace matador::utils;
    using namespace matador::query;
    // create select query
    auto result = matador::query::query::select(sql::column_generator::generate<typename Pointer::value_type>(schema_, true))
      .from(*foreign_table)
      .where(sql::column(foreign_table, pk->name(), "") == _)
      .prepare(executor_);
    if (!result) {
      throw query_builder_exception(query_build_error::QueryError, result.release_error());
    }

    entity_query_data_.lazy_loading_statements.emplace(id, std::move(result.release()));
  }
}

}
#endif //QUERY_ENTITY_QUERY_BUILDER_HPP