Implement returning of auto id using RETURNING clause in PostgreSQL

Before we used a separate SELECT lastval() query which was both inefficient
and error-prone in cases where INSERT may cause triggers to override the
last value.

4 files changed, 53 insertions, 28 deletions

diff --git a/odb/pgsql/container-statements.hxx b/odb/pgsql/container-statements.hxx
index 252eb77..a8d15d8 100644
--- a/odb/pgsql/container-statements.hxx
+++ b/odb/pgsql/container-statements.hxx
@@ -187,7 +187,8 @@ namespace odb
               insert_one_types_,
               insert_one_count_,
               data_image_binding_,
-              data_image_native_binding_));
+              data_image_native_binding_,
+              false));
         }
 
         return *insert_one_;
diff --git a/odb/pgsql/object-statements.hxx b/odb/pgsql/object-statements.hxx
index 6942361..2a26cce 100644
--- a/odb/pgsql/object-statements.hxx
+++ b/odb/pgsql/object-statements.hxx
@@ -265,7 +265,8 @@ namespace odb
               object_traits::persist_statement_types,
               insert_column_count,
               insert_image_binding_,
-              insert_image_native_binding_));
+              insert_image_native_binding_,
+              object_traits::auto_id));
         }
 
         return *persist_;
@@ -522,7 +523,8 @@ namespace odb
               object_traits::persist_statement_types,
               insert_column_count,
               insert_image_binding_,
-              insert_image_native_binding_));
+              insert_image_native_binding_,
+              false));
         }
 
         return *persist_;
diff --git a/odb/pgsql/statement.cxx b/odb/pgsql/statement.cxx
index 671bf61..e76e081 100644
--- a/odb/pgsql/statement.cxx
+++ b/odb/pgsql/statement.cxx
@@ -11,6 +11,7 @@
 
 #include <odb/exceptions.hxx> // object_not_persistent
 
+#include <odb/pgsql/pgsql-oid.hxx>
 #include <odb/pgsql/statement.hxx>
 #include <odb/pgsql/connection.hxx>
 #include <odb/pgsql/transaction.hxx>
@@ -442,19 +443,18 @@ namespace odb
                       const Oid* types,
                       size_t types_count,
                       binding& data,
-                      native_binding& native_data)
+                      native_binding& native_data,
+                      bool returning)
         : statement (conn, name, stmt, types, types_count),
           data_ (data),
           native_data_ (native_data),
-          id_cached_ (false)
+          returning_ (returning)
     {
     }
 
     bool insert_statement::
     execute ()
     {
-      id_cached_ = false;
-
       bind_param (native_data_, data_);
 
       auto_handle<PGresult> h (
@@ -481,26 +481,44 @@ namespace odb
         translate_error (conn_, h);
       }
 
-      return true;
-    }
-
-    unsigned long long insert_statement::
-    id ()
-    {
-      if (id_cached_)
-        return id_;
-
-      auto_handle<PGresult> h (
-        PQexecParams (conn_.handle (), "select lastval ()", 0, 0, 0, 0, 0, 1));
-
-      if (!is_good_result (h))
-        translate_error (conn_, h);
+      if (returning_)
+      {
+        // Get the id value that was returned using the RETURNING clause.
+        //
+        const char* v (PQgetvalue (h, 0, 0));
 
-      id_ = endian_traits::ntoh (*reinterpret_cast<unsigned long long*> (
-                                   PQgetvalue (h, 0, 0)));
-      id_cached_ = true;
+        // While the ODB auto id type can only be INT or BIGINT, handle the
+        // SMALLINT integer in case we are dealing with a custom schema.
+        //
+        switch (PQftype (h, 0))
+        {
+        case int2_oid:
+          {
+            id_ = endian_traits::ntoh (
+              *reinterpret_cast<const unsigned short*> (v));
+            break;
+          }
+        case int4_oid:
+          {
+            id_ = endian_traits::ntoh (
+              *reinterpret_cast<const unsigned int*> (v));
+            break;
+          }
+        case int8_oid:
+          {
+            id_ = endian_traits::ntoh (
+              *reinterpret_cast<const unsigned long long*> (v));
+            break;
+          }
+        default:
+          {
+            assert (false);
+            break;
+          }
+        }
+      }
 
-      return id_;
+      return true;
     }
 
     //
diff --git a/odb/pgsql/statement.hxx b/odb/pgsql/statement.hxx
index d7cf6dc..8488e12 100644
--- a/odb/pgsql/statement.hxx
+++ b/odb/pgsql/statement.hxx
@@ -180,7 +180,8 @@ namespace odb
                         const Oid* types,
                         std::size_t types_count,
                         binding& data,
-                        native_binding& native_data);
+                        native_binding& native_data,
+                        bool returning);
 
       // Return true if successful and false if the row is a duplicate.
       // All other errors are reported by throwing exceptions.
@@ -189,7 +190,10 @@ namespace odb
       execute ();
 
       unsigned long long
-      id ();
+      id ()
+      {
+        return id_;
+      }
 
     private:
       insert_statement (const insert_statement&);
@@ -199,7 +203,7 @@ namespace odb
       binding& data_;
       native_binding& native_data_;
 
-      bool id_cached_;
+      bool returning_;
       unsigned long long id_;
     };