added `returning` token for insert query

include/matador/query/intermediates/query_into_intermediate.hpp

@@ -2,24 +2,32 @@
 #define QUERY_INTO_INTERMEDIATE_HPP
 
 #include "matador/query/intermediates/executable_query.hpp"
+#include "matador/query/intermediates/fetchable_query.hpp"
 
 #include "matador/query/value_extractor.hpp"
 
 #include "matador/utils/placeholder.hpp"
 
 namespace matador::query {
+class table_column;
 
-class query_into_intermediate : public query_intermediate
-{
+class query_values_intermediate : public executable_query {
+public:
+  using executable_query::executable_query;
+
+  fetchable_query returning(const table_column &col);
+};
+
+class query_into_intermediate : public query_intermediate {
 public:
   using query_intermediate::query_intermediate;
 
-  executable_query values(std::initializer_list<std::variant<utils::placeholder, utils::database_type>> values);
-  executable_query values(std::vector<std::variant<utils::placeholder, utils::database_type>> &&values);
-  executable_query values(std::vector<utils::placeholder> &&values);
-  executable_query values(std::vector<utils::database_type> &&values);
+  query_values_intermediate values(std::initializer_list<std::variant<utils::placeholder, utils::database_type>> values);
+  query_values_intermediate values(std::vector<std::variant<utils::placeholder, utils::database_type>> &&values);
+  query_values_intermediate values(std::vector<utils::placeholder> &&values);
+  query_values_intermediate values(std::vector<utils::database_type> &&values);
   template<class Type>
-  executable_query values(const Type &obj) {
+  query_values_intermediate values(const Type &obj) {
     return values(value_extractor::extract(obj));
   }
 };

include/matador/query/internal/query_parts.hpp

@@ -317,8 +317,7 @@ private:
 /**
  * Represents the SQL VALUES part
  */
-class query_values_part final : public query_part
-{
+class query_values_part final : public query_part {
 public:
   explicit query_values_part(std::vector<std::variant<utils::placeholder, utils::database_type>> &&values);
 
@@ -331,6 +330,19 @@ private:
   std::vector<std::variant<utils::placeholder, utils::database_type>> values_;
 };
 
+class query_returning_part final : public query_part {
+public:
+  explicit query_returning_part(std::vector<table_column> columns);
+
+  [[nodiscard]] const std::vector<table_column>& columns() const;
+
+private:
+  void accept(query_part_visitor &visitor) override;
+
+private:
+  std::vector<table_column> columns_;
+};
+
 class query_update_part final : public query_part
 {
 public:

include/matador/query/query_builder.hpp

@@ -44,7 +44,6 @@ protected:
   void visit(internal::query_drop_key_constraint_part_by_name &part) override;
   void visit(internal::query_drop_key_constraint_part_by_constraint &part) override;
 
-protected:
   void visit(internal::query_select_part &part) override;
   void visit(internal::query_from_part &part) override;
   void visit(internal::query_join_table_part &part) override;
@@ -57,9 +56,11 @@ protected:
   void visit(internal::query_order_by_desc_part &part) override;
   void visit(internal::query_offset_part &part) override;
   void visit(internal::query_limit_part &part) override;
+
   void visit(internal::query_insert_part &part) override;
   void visit(internal::query_into_part &part) override;
   void visit(internal::query_values_part &part) override;
+  void visit(internal::query_returning_part &part) override;
 
   void visit(internal::query_update_part &part) override;
   void visit(internal::query_set_part &part) override;

include/matador/query/query_part_visitor.hpp

@@ -4,7 +4,9 @@
 namespace matador::query {
 
 namespace internal {
+// Alter common
 class query_alter_part;
+// Alter table
 class query_alter_table_part;
 class query_add_key_constraint_part;
 class query_drop_key_constraint_part_by_name;
@@ -13,6 +15,7 @@ class query_add_foreign_key_constraint_part;
 class query_add_constraint_part_by_constraint;
 class query_add_foreign_key_reference_part;
 class query_add_primary_key_constraint_part;
+// Select
 class query_select_part;
 class query_from_part;
 class query_join_table_part;
@@ -25,20 +28,30 @@ class query_order_by_asc_part;
 class query_order_by_desc_part;
 class query_offset_part;
 class query_limit_part;
+// Insert
 class query_insert_part;
 class query_into_part;
 class query_values_part;
+class query_returning_part;
+// Update
 class query_update_part;
 class query_set_part;
+// Delete
 class query_delete_part;
 class query_delete_from_part;
+// Create common
 class query_create_part;
+// Create table
 class query_create_table_part;
 class query_create_table_columns_part;
 class query_create_table_constraints_part;
+// Create schema
 class query_create_schema_part;
+// Drop common
 class query_drop_part;
+// Drop table
 class query_drop_table_part;
+// Drop schema
 class query_drop_schema_part;
 }
 
@@ -72,6 +85,7 @@ public:
   virtual void visit(internal::query_insert_part &part) = 0;
   virtual void visit(internal::query_into_part &part) = 0;
   virtual void visit(internal::query_values_part &part) = 0;
+  virtual void visit(internal::query_returning_part &part) = 0;
 
   virtual void visit(internal::query_update_part &part) = 0;
   virtual void visit(internal::query_set_part &part) = 0;

include/matador/query/query_utils.hpp

@@ -12,6 +12,7 @@ class dialect;
 namespace matador::query {
 class table_column;
 
+void prepare_column(std::string &out, const sql::dialect& d, const table_column &col);
 [[nodiscard]] std::string prepare_identifier(const sql::dialect& d, const table_column &col);
 [[nodiscard]] std::string prepare_criteria(const sql::dialect& d, const table_column &col);
 }

include/matador/sql/dialect.hpp

@@ -166,6 +166,7 @@ public:
   [[nodiscard]] const std::string& primary_key() const;
   [[nodiscard]] const std::string& references() const;
   [[nodiscard]] const std::string& remove() const;
+  [[nodiscard]] const std::string& returning() const;
   [[nodiscard]] const std::string& rollback() const;
   [[nodiscard]] const std::string& schema() const;
   [[nodiscard]] const std::string& select() const;
@@ -236,6 +237,7 @@ private:
     {dialect_token::PrimaryKey,  "PRIMARY KEY"},
     {dialect_token::References,  "REFERENCES"},
     {dialect_token::Remove,      "DELETE"},
+    {dialect_token::Returning,   "RETURNING"},
     {dialect_token::Rollback,    "ROLLBACK TRANSACTION"},
     {dialect_token::Schema,      "SCHEMA"},
     {dialect_token::Select,      "SELECT"},

include/matador/sql/dialect_token.hpp

@@ -51,6 +51,7 @@ enum class dialect_token : uint8_t {
   PrimaryKey,
   References,
   Remove,
+  Returning,
   Rollback,
   Schema,
   Select,

include/matador/sql/execute_result.hpp

@@ -1,11 +1,14 @@
 #ifndef MATADOR_EXECUTE_RESULT_HPP
 #define MATADOR_EXECUTE_RESULT_HPP
+
+#include "matador/utils/identifier.hpp"
+
 #include <vector>
 
 namespace matador::sql {
 struct execute_result {
   size_t affected_rows{};
-  std::vector<long long> insert_ids{};
+  std::vector<utils::identifier> generated_ids{};
 };
 }
 #endif  // MATADOR_EXECUTE_RESULT_HPP

include/matador/sql/query_context.hpp

@@ -27,9 +27,10 @@ enum class sql_command {
   AlterSchema
 };
 
-struct sql_command_info {
-  std::string sql;
-  sql_command command{};
+enum class return_mode {
+  None,
+  GeneratedKeys,
+  Rows
 };
 
 struct query_context {
@@ -44,6 +45,14 @@ struct query_context {
   // Data for resolving query result
   std::shared_ptr<resolver_service> resolver{};
   std::type_index result_type = typeid(void);
+  return_mode mode = return_mode::None;
+
+  [[nodiscard]] bool is_ddl() const { return command == sql_command::Create || command == sql_command::Alter; }
+  [[nodiscard]] bool is_dml() const { return command == sql_command::Insert || command == sql_command::Update || command == sql_command::Delete; }
+  [[nodiscard]] bool is_query() const { return command == sql_command::Select; }
+
+  [[nodiscard]] bool expects_rows() const { return mode == return_mode::Rows; }
+  [[nodiscard]] bool expects_generated_keys() const { return mode == return_mode::GeneratedKeys; }
 };
 }
 

source/orm/query/intermediates/query_into_intermediate.cpp

@@ -4,18 +4,21 @@
 #include "matador/query/query_data.hpp"
 
 namespace matador::query {
+fetchable_query query_values_intermediate::returning(const table_column &col) {
+  context_->parts.push_back(std::make_unique<internal::query_returning_part>(std::vector{col}));
+  return {context_};
+}
 
-executable_query query_into_intermediate::values(const std::initializer_list<std::variant<utils::placeholder, utils::database_type>> values)
-{
+query_values_intermediate query_into_intermediate::values(const std::initializer_list<std::variant<utils::placeholder, utils::database_type>> values) {
   return this->values(std::vector(values));
 }
 
-executable_query query_into_intermediate::values(std::vector<std::variant<utils::placeholder, utils::database_type>> &&values) {
+query_values_intermediate query_into_intermediate::values(std::vector<std::variant<utils::placeholder, utils::database_type>> &&values) {
   context_->parts.push_back(std::make_unique<internal::query_values_part>(std::move(values)));
   return {context_};
 }
 
-executable_query query_into_intermediate::values(std::vector<utils::placeholder>&& values) {
+query_values_intermediate query_into_intermediate::values(std::vector<utils::placeholder>&& values) {
   std::vector<std::variant<utils::placeholder, utils::database_type>> transformed_values;
   transformed_values.reserve(values.size());
   for (auto&& val : values) {
@@ -24,7 +27,7 @@ executable_query query_into_intermediate::values(std::vector<utils::placeholder>
   return this->values(std::move(transformed_values));
 }
 
-executable_query query_into_intermediate::values(std::vector<utils::database_type>&& values) {
+query_values_intermediate query_into_intermediate::values(std::vector<utils::database_type>&& values) {
   std::vector<std::variant<utils::placeholder, utils::database_type>> transformed_values;
   transformed_values.reserve(values.size());
   for (auto&& val : values) {

source/orm/query/internal/query_parts.cpp

@@ -311,6 +311,19 @@ void query_values_part::accept(query_part_visitor &visitor) {
   visitor.visit(*this);
 }
 
+query_returning_part::query_returning_part(std::vector<table_column> columns)
+: query_part(sql::dialect_token::Returning)
+, columns_(std::move(columns)) {
+}
+
+const std::vector<table_column> & query_returning_part::columns() const {
+  return columns_;
+}
+
+void query_returning_part::accept(query_part_visitor &visitor) {
+  visitor.visit(*this);
+}
+
 query_update_part::query_update_part(class table tab)
 : query_part(sql::dialect_token::Update)
 , table_(std::move(tab)) {

source/orm/query/query_builder.cpp

@@ -36,16 +36,10 @@ sql::query_context query_builder::compile(const query_data &data,
   return {query_};
 }
 
-std::string handle_column(sql::query_context &ctx, const sql::dialect *d, const table_column &col) {
-  if (col.is_function()) {
-    ctx.prototype.emplace_back(col.has_alias() ? col.alias() : col.canonical_name());
-    ctx.prototype.back().change_type(utils::basic_type::Int32);
-  } else {
-    ctx.prototype.emplace_back(col.has_alias() ? col.alias() : col.canonical_name());
-  }
-
-  return prepare_identifier(*d, col);
-}
+void build_columns_with_name_only(std::string &out, const std::vector<table_column> &cols, const sql::dialect &d);
+void build_columns(std::string &out, const std::vector<table_column> &cols, const sql::dialect &d);
+void build_fetchable_columns(sql::query_context &ctx, const std::vector<table_column> &cols, const sql::dialect &d);
+void prepare_prototype(std::vector<object::attribute> &prototype, const table_column &col);
 
 void query_builder::visit(internal::query_alter_part& part) {
   query_.sql = dialect_->token_at(part.token());
@@ -61,9 +55,6 @@ void query_builder::visit(internal::query_add_key_constraint_part& part) {
   query_.sql += " " + dialect_->add_constraint() + " " + part.name();
 }
 
-void build_columns_with_name_only(std::string &out, const std::vector<table_column> &cols, const sql::dialect &d);
-void build_columns(std::string &out, const std::vector<table_column> &cols, const sql::dialect &d);
-
 void query_builder::visit(internal::query_add_foreign_key_constraint_part& part) {
   query_.sql += " " + dialect_->token_at(part.token()) + " (";
   build_columns(query_.sql, part.columns(), *dialect_);
@@ -100,23 +91,7 @@ void query_builder::visit(internal::query_select_part &part) {
 
   query_.prototype.clear();
 
-  std::string result;
-  if (part.columns().empty()) {
-      result = dialect_->asterisk();
-  } else if (const auto &columns = part.columns(); columns.size() < 2) {
-    for (const auto &col: columns) {
-      result.append(handle_column(query_, dialect_, col ));
-    }
-  } else {
-    auto it = columns.begin();
-    result.append(handle_column(query_, dialect_, *it++));
-    for (; it != columns.end(); ++it) {
-      result.append(", ");
-      result.append(handle_column(query_, dialect_, *it));
-    }
-  }
-
-  query_.sql += result;
+  build_fetchable_columns(query_, part.columns(), *dialect_);
 }
 
 void query_builder::visit(internal::query_from_part &part) {
@@ -265,15 +240,20 @@ void query_builder::visit(internal::query_values_part &part) {
   query_.sql += " " + result;
 }
 
-void query_builder::visit(internal::query_update_part &part)
-{
-  query_.command = sql::sql_command::Update;
-  query_.table_name = part.table().name();
-  query_.sql += query_builder::build_table_name(part.token(), *dialect_, query_.table_name);
+void query_builder::visit(internal::query_returning_part &part) {
+  query_.mode = sql::return_mode::Rows;
+  query_.sql += " " + dialect_->returning() + " ";
+
+  build_fetchable_columns(query_, part.columns(), *dialect_);
 }
 
-void query_builder::visit(internal::query_delete_part &/*delete_part*/)
-{
+void query_builder::visit(internal::query_update_part &part) {
+  query_.command = sql::sql_command::Update;
+  query_.table_name = part.table().name();
+  query_.sql += build_table_name(part.token(), *dialect_, query_.table_name);
+}
+
+void query_builder::visit(internal::query_delete_part &/*delete_part*/) {
   query_.command = sql::sql_command::Delete;
   query_.sql = dialect_->remove();
 }
@@ -408,6 +388,27 @@ void build_columns(std::string &out, const std::vector<table_column> &cols,  con
   }
 }
 
+void build_fetchable_columns(sql::query_context &ctx, const std::vector<table_column> &cols, const sql::dialect &d) {
+  bool first = true;
+  for (const auto& col : cols) {
+    if (!first) {
+      ctx.sql.append(", ");
+    }
+    prepare_prototype(ctx.prototype, col);
+    prepare_column(ctx.sql, d, col);
+    first = false;
+  }
+}
+
+void prepare_prototype(std::vector<object::attribute> &prototype, const table_column &col) {
+  if (col.is_function()) {
+    prototype.emplace_back(col.has_alias() ? col.alias() : col.canonical_name());
+    prototype.back().change_type(utils::basic_type::Int32);
+  } else {
+    prototype.emplace_back(col.has_alias() ? col.alias() : col.canonical_name());
+  }
+}
+
 std::string build_constraint(const table_constraint& cons, const sql::dialect& d) {
     std::string result;
     if (!cons.name().empty()) {

source/orm/query/query_utils.cpp

@@ -4,6 +4,18 @@
 #include "matador/query/internal/string_builder_utils.hpp"
 
 namespace matador::query {
+void prepare_column(std::string &out, const sql::dialect &d, const table_column &col) {
+  if (!col.is_function()) {
+    prepare_identifier_string_append(out, col.name(), d);
+  } else {
+    if (col.column_name() == d.asterisk()) {
+      out += d.sql_function_at(col.function()) + "(" + col.column_name() + ")";
+    } else {
+      out += d.sql_function_at(col.function()) + "(" + col.column_name() + ") " + d.as() + " " + col.alias();
+    }
+  }
+}
+
 std::string prepare_identifier(const sql::dialect& d, const table_column& col) {
   std::string result;
   if (!col.is_function()) {

source/orm/sql/dialect.cpp

@@ -271,6 +271,10 @@ const std::string &dialect::remove() const {
   return token_at(dialect_token::Remove);
 }
 
+const std::string & dialect::returning() const {
+  return token_at(dialect_token::Returning);
+}
+
 const std::string &dialect::rollback() const {
   return token_at(dialect_token::Rollback);
 }

test/backends/QueryTest.cpp

@@ -87,6 +87,30 @@ TEST_CASE_METHOD(QueryFixture, "Execute select statement with where clause", "[q
   }
 }
 
+TEST_CASE_METHOD(QueryFixture, "Execute insert with returning", "[query][insert][returning]") {
+  REQUIRE(repo.attach<person>("persons"));
+  auto result = repo.create(db);
+  REQUIRE(result.is_ok());
+
+  check_table_exists(PERSON.table_name());
+
+  person george{7, "george", 45};
+  george.image.emplace_back(37);
+
+  auto res = query::insert()
+    .into(PERSON, {PERSON.id, PERSON.name, PERSON.age, PERSON.image})
+    .values(george)
+    .returning(PERSON.id)
+    .fetch_one(db);
+  REQUIRE(res.is_ok());
+
+  auto rec = res.release();
+  REQUIRE(rec->size() == 1);
+  REQUIRE(rec->at(0).name() == "persons.id");
+  REQUIRE(rec->at(0).is_integer());
+  REQUIRE(rec->at(0).as<uint32_t>() == 7);
+}
+
 TEST_CASE_METHOD(QueryFixture, "Execute insert statement", "[query][insert]") {
   auto res = query::create()
     .table("person")

todo.md

@@ -1,13 +1,25 @@
 # Todo
 
 - move `prepare_*` methods from `dialect` to `query_compiler`
-- add `session_insert_builder` and `session_update_builder` (returning multiple statements)
+- add `insert_query_builder` and `update_update_builder` (returning multiple statements)
 - finish fetch eager many-to-many relations
 - implement lazy loading
 - implement polymorphic class hierarchies
-- finish `schema_repository` classes (move add/drop from `session` to `schema`)
+- finish `database` (schema_repository) class (move add/drop from `session` to `schema`)
 - implement a flag class for enumerations
 
+- PK generator
+1. Introduce execute_context
+2. Introduce execute_result
+3. Add `abstract_pk_generator` class
+4. Add `identity_pk_generator` class
+5. Add `sequence_pk_generator` class
+6. Add `table_pk_generator` class
+7. Add `manual_pk_generator` class 
+8. Extend `session::insert` logic to use pk generator
+9. Extend `postgres_connection::execute` to handle `returning` clause
+10. Add generator to `schema_node` or `table` class when attaching type
+11. Extend `query_builder` and `intermediates` with `returning` intermediate
 
 __Proposal for polymorphic classes:__