query/include/matador/orm/session.hpp

#ifndef QUERY_SESSION_HPP
#define QUERY_SESSION_HPP

#include "matador/orm/error_code.hpp"

#include "matador/query/select_query_builder.hpp"
#include "matador/query/criteria.hpp"
#include "matador/query/insert_query_builder.hpp"
#include "matador/query/query.hpp"
#include "matador/query/generator.hpp"
#include "matador/query/schema.hpp"

#include "matador/sql/connection.hpp"
#include "matador/sql/connection_pool.hpp"
#include "matador/sql/executor.hpp"
#include "matador/sql/resolver_service.hpp"
#include "matador/sql/statement.hpp"
#include "matador/sql/statement_cache.hpp"

#include "matador/object/object_ptr.hpp"

#include <unordered_map>
#include <utility>

namespace matador::orm {
utils::error make_error(error_code ec, const std::string &msg);

struct session_context {
  session_context(utils::message_bus &bus, std::string dns, const size_t count, const size_t cache_size = 500)
  : bus(bus)
  , dns(std::move(dns))
  , connection_count(count)
  , cache_size(cache_size) {
  }
  utils::message_bus &bus;
  std::string dns;
  size_t connection_count{};
  size_t cache_size{500};
  std::shared_ptr<sql::resolver_service> resolver_service = std::make_shared<sql::resolver_service>();
};

class session final : public sql::executor {
public:
  session(session_context &&ctx, const query::schema &scm);

  /**
   * Insert the given object into the session.
   *
   * @tparam Type Type of object to insert
   * @param obj Object to insert
   * @return Inserted object
   */
  template<typename Type>
  utils::result<object::object_ptr<Type>, utils::error> insert(object::object_ptr<Type> obj);

  template<typename Type>
  utils::result<object::object_ptr<Type>, utils::error> update(const object::object_ptr<Type> &obj);
  template<typename Type>
  utils::result<void, utils::error> remove(const object::object_ptr<Type> &obj);

  template<typename Type, typename PrimaryKeyType>
  utils::result<object::object_ptr<Type>, utils::error> find(const PrimaryKeyType &pk);
  template<typename Type>
  utils::result<sql::query_result<Type>, utils::error> find(query::criteria_ptr clause = {});

  template<typename Type>
  utils::result<void, utils::error> drop_table() const;
  utils::result<void, utils::error> drop_table(const std::string &table_name) const;

  [[nodiscard]] utils::result<sql::query_result<sql::record>, utils::error> fetch_all(const sql::query_context &q) const;
  [[nodiscard]] utils::result<sql::execute_result, utils::error> execute(const std::string &sql) const;

  [[nodiscard]] std::vector<object::attribute> describe_table(const std::string &table_name) const;
  [[nodiscard]] bool table_exists(const std::string &table_name) const;

  [[nodiscard]] utils::result<std::unique_ptr<sql::query_result_impl>, utils::error> fetch(const sql::query_context &ctx) const override;
  [[nodiscard]] utils::result<sql::execute_result, utils::error> execute(const sql::query_context &ctx) const override;
  [[nodiscard]] utils::result<sql::statement, utils::error> prepare(const sql::query_context &ctx) override;
  [[nodiscard]] std::string str(const sql::query_context &ctx) const override;
  [[nodiscard]] const sql::dialect &dialect() const override;
  [[nodiscard]] std::shared_ptr<sql::resolver_service> resolver() const override;
  [[nodiscard]] const query::basic_schema &schema() const;

private:
  sql::connection_pool pool_;
  mutable sql::statement_cache cache_;
  const sql::dialect &dialect_;

  const query::basic_schema &schema_;
  mutable std::unordered_map<std::string, std::vector<object::attribute> > prototypes_;
  std::shared_ptr<sql::resolver_service> resolver_service_;
};

template<typename Type>
utils::result<object::object_ptr<Type>, utils::error> session::insert(object::object_ptr<Type> obj) {
  if (obj.is_persistent()) {
    return utils::ok(obj);
  }

  const auto it = schema_.find(typeid(Type));
  if (it == schema_.end()) {
    return utils::failure(make_error(error_code::UnknownType, "Failed to determine requested type."));
  }

  // Build dependency-ordered insert steps (deps first, root last)
  query::insert_query_builder iqb(schema_);
  auto steps = iqb.build(obj);
  if (!steps.is_ok()) {
    return utils::failure(make_error(error_code::FailedToBuildQuery, "Failed to build insert dependency queries."));
  }

  // Execute all steps; for Identity steps read RETURNING and write pk back into the object
  for (auto &step : *steps) {
    if (step.pk_is_unset && step.set_pk) {
      if (step.pk_generator == utils::generator_type::Manual) {
        if (step.pk_is_unset()) {
          return utils::failure(make_error(error_code::NoPrimaryKey, "Manual primary key is required but unset."));
        }
      } else if (step.pk_generator == utils::generator_type::Sequence) {
        if (step.pk_is_unset()) {
          // hard-coded naming as you specified earlier
          // <table_name>_seq (table name known in schema meta; if you prefer, store it in step too)
          // For now, we derive from the schema type used for the root insert: simplistic but works if table name == entity name.
          const std::string seq_name = it->second.name() + "_seq";

          auto id_res = query::query::select().nextval(seq_name).fetch_value<std::uint64_t>(*this);
          if (!id_res.is_ok() || !id_res.value().has_value()) {
            return utils::failure(make_error(error_code::FailedToBuildQuery, "Failed to obtain next sequence value."));
          }
          step.set_pk(*id_res.value());
        }
      } else if (step.pk_generator == utils::generator_type::Table) {
        if (step.pk_is_unset()) {
          // TODO: implement table id generation; same idea as above:
          // UPDATE <table>_tbl_seq ... RETURNING ...
          return utils::failure(make_error(error_code::Failed, "Table primary key generator not implemented yet."));
        }
      }
    }

    // --- Execute step ---
    if (std::holds_alternative<query::executable_query>(step.query)) {
      const auto &q = std::get<query::executable_query>(step.query);
      const auto exec_res = q.execute(*this);
      if (!exec_res.is_ok()) {
        return utils::failure(exec_res.err());
      }
      continue;
    }

    const auto &q = std::get<query::fetchable_query>(step.query);
    auto rec_res = q.fetch_one(*this);
    if (!rec_res.is_ok()) {
      return utils::failure(rec_res.err());
    }
    if (!rec_res.value().has_value()) {
      return utils::failure(make_error(error_code::FailedToFindObject, "INSERT ... RETURNING did not return a row."));
    }
    if (step.apply_returning) {
      step.apply_returning(*rec_res.value());
    }
  }

  obj.change_state(object::object_state::Persistent);
  return utils::ok(obj);
}

class pk_object_binder final {
public:
  explicit pk_object_binder(sql::statement &stmt, const size_t position)
  : stmt_(stmt)
  , binding_position_(position) {
  }

  template<class Type>
  sql::statement &bind(Type &obj) {
    access::process(*this, obj);
    return stmt_;
  }

  template<class Type>
  void on_primary_key(const char * /*id*/, Type &x, const utils::primary_key_attribute & /*attr*/ = utils::default_pk_attributes) {
    stmt_.bind(binding_position_, x);
  }

  static void on_revision(const char * /*id*/, uint64_t &/*rev*/) {
  }

  template<class Type>
  static void on_attribute(const char * /*id*/, Type &/*x*/, const utils::field_attributes &/*attr*/ = utils::null_attributes) {}
  template<class Pointer>
  static void on_belongs_to(const char * /*id*/, Pointer &/*x*/, const utils::foreign_attributes &/*attr*/ = utils::CascadeNoneFetchLazy) {}
  template<class Pointer>
  static void on_has_one(const char * /*id*/, Pointer &/*x*/, const utils::foreign_attributes &/*attr*/ = utils::CascadeNoneFetchLazy) {}

  template<class ContainerType>
  static void on_has_many(const char * /*id*/,
                          ContainerType &/*c*/,
                          const char * /*join_column*/,
                          const utils::foreign_attributes &/*attr*/ = utils::CascadeNoneFetchLazy) {}

  template<class ContainerType>
  static void on_has_many_to_many(const char * /*id*/,
                                  ContainerType &/*c*/,
                                  const char * /*join_column*/,
                                  const char * /*inverse_join_column*/,
                                  const utils::foreign_attributes &/*attr*/) {}

  template<class ContainerType>
  static void on_has_many_to_many(const char * /*id*/,
                                  ContainerType &/*c*/,
                                  const utils::foreign_attributes &/*attr*/) {}

private:
  sql::statement &stmt_;
  size_t binding_position_{0};
  sql::object_pk_binder pk_binder_{};
};

template<typename Type>
utils::result<object::object_ptr<Type>, utils::error> session::update(const object::object_ptr<Type> &obj) {
  const auto it = schema_.find(typeid(Type));
  if (it == schema_.end()) {
    return utils::failure(make_error(error_code::UnknownType, "Failed to determine requested type."));
  }
  using namespace matador::utils;
  using namespace matador::query;

  const auto col = table_column(it->second.node().info().primary_key_attribute()->name());
  auto res = matador::query::query::update(it->second.name())
    .set<Type>()
    .where(col == _)
    .prepare(*this);
  if (!res) {
    return utils::failure(res.err());
  }

  res->bind(*obj);
  pk_object_binder binder(res.value(), res->bind_pos());
  if (const auto update_result = binder.bind(*obj).execute(); !update_result.is_ok()) {
    return utils::failure(update_result.err());
  }
  return utils::ok(object::object_ptr{obj});
}

template<typename Type>
utils::result<void, utils::error> session::remove(const object::object_ptr<Type> &obj) {
  const auto it = schema_.find(typeid(Type));
  if (it == schema_.end()) {
    return utils::failure(make_error(error_code::UnknownType, "Failed to determine requested type."));
  }
  using namespace matador::utils;
  using namespace matador::query;

  const auto col = table_column(it->second.node().info().primary_key_attribute()->name());
  auto res = matador::query::query::remove()
    .from(it->second.name())
    .where(col == _)
    .prepare(*this);
  if (!res) {
    return utils::failure(res.err());
  }

  pk_object_binder binder(res.value(), res->bind_pos());
  if (const auto update_result = binder.bind(*obj).execute(); !update_result.is_ok()) {
    return utils::failure(update_result.err());
  }
  return utils::ok<void>();
}

template<typename Type, typename PrimaryKeyType>
utils::result<object::object_ptr<Type>, utils::error> session::find(const PrimaryKeyType &pk) {
  const auto it = schema_.find(typeid(Type));
  if (it == schema_.end()) {
    return utils::failure(make_error(error_code::UnknownType, "Failed to determine requested type."));
  }
  const auto &info = it->second.node().info();
  if (!info.has_primary_key()) {
    return utils::failure(make_error(error_code::NoPrimaryKey, "Type hasn't primary key."));
  }

  query::select_query_builder eqb(schema_);
  const query::table_column c(&it->second.table(), info.primary_key_attribute()->name());
  using namespace matador::query;
  auto data = eqb.build<Type>(c == utils::_);
  if (!data.is_ok()) {
    return utils::failure(make_error(error_code::FailedToBuildQuery,
                                     "Failed to build query for type " + info.name() + "."));
  }

  auto res = data->prepare(*this);

  if (!res) {
    return utils::failure(res.err());
  }
  auto stmt_result = res->bind(0, const_cast<PrimaryKeyType &>(pk))
    .template fetch_one<Type>();
  if (!stmt_result) {
    return utils::failure(make_error(error_code::FailedToFindObject,
                                     "Failed to find object of type " + info.name() + " with primary key " +
                                     std::to_string(pk) + "."));
  }
  return utils::ok(*stmt_result);
}

template<typename Type>
utils::result<sql::query_result<Type>, utils::error> session::find(query::criteria_ptr clause) {
  const auto it = schema_.find(typeid(Type));
  if (it == schema_.end()) {
    return utils::failure(make_error(error_code::UnknownType, "Failed to determine requested type."));
  }

  query::select_query_builder eqb(schema_);
  auto data = eqb.build<Type>(std::move(clause));
  if (!data.is_ok()) {
    return utils::failure(make_error(error_code::FailedToBuildQuery,
                                     "Failed to build query for type " + it->second.name() + "."));
  }

  auto result = data->prepare(*this);
  if (!result.is_ok()) {
    return utils::failure(result.err());
  }

  return result->template fetch<Type>();
}

template<typename Type>
utils::result<void, utils::error> session::drop_table() const {
  const auto it = schema_.find(typeid(Type));
  if (it == schema_.end()) {
    return utils::failure(make_error(error_code::UnknownType, "Failed to determine requested type."));
  }
  return drop_table(it->second.name());
}

}
#endif //QUERY_SESSION_HPP