Add support for single row mode and chunked mode

taocpp · Nov 21, 2024 · 34bdab6 · 34bdab6
1 parent 9cafd37
commit 34bdab6
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 0 deletions.
diff --git a/include/tao/pq/result.hpp b/include/tao/pq/result.hpp
@@ -54,6 +54,7 @@ namespace tao::pq
    public:
       [[nodiscard]] auto has_rows_affected() const noexcept -> bool;
       [[nodiscard]] auto rows_affected() const -> std::size_t;
+      [[nodiscard]] auto is_final() const -> bool;
 
       [[nodiscard]] auto columns() const noexcept -> std::size_t
       {

diff --git a/include/tao/pq/transaction.hpp b/include/tao/pq/transaction.hpp
@@ -141,6 +141,11 @@ namespace tao::pq
          }
       }
 
+      void set_single_row_mode();
+#if defined( LIBPQ_HAS_CHUNK_MODE )
+      void set_chunk_mode( const std::size_t rows );
+#endif
+
       [[nodiscard]] auto get_result( const std::chrono::steady_clock::time_point start = std::chrono::steady_clock::now() ) -> result;
 
       template< typename... As >

diff --git a/src/lib/pq/result.cpp b/src/lib/pq/result.cpp
@@ -45,6 +45,10 @@ namespace tao::pq
       switch( PQresultStatus( pgresult ) ) {
          case PGRES_COMMAND_OK:
          case PGRES_TUPLES_OK:
+         case PGRES_SINGLE_TUPLE:
+#if defined( LIBPQ_HAS_CHUNK_MODE )
+         case PGRES_TUPLES_CHUNK:
+#endif
             return;
 
          case PGRES_EMPTY_QUERY:
@@ -74,6 +78,24 @@ namespace tao::pq
       return internal::from_chars< std::size_t >( str );
    }
 
+   auto result::is_final() const -> bool
+   {
+      switch( PQresultStatus( m_pgresult.get() ) ) {
+         case PGRES_COMMAND_OK:
+         case PGRES_TUPLES_OK:
+            return true;
+
+         case PGRES_SINGLE_TUPLE:
+#if defined( LIBPQ_HAS_CHUNK_MODE )
+         case PGRES_TUPLES_CHUNK:
+#endif
+            return false;
+
+         default:
+            TAO_PQ_UNREACHABLE;  // LCOV_EXCL_LINE
+      }
+   }
+
    auto result::name( const std::size_t column ) const -> std::string
    {
       if( column >= m_columns ) {

diff --git a/src/lib/pq/transaction.cpp b/src/lib/pq/transaction.cpp
@@ -141,6 +141,24 @@ namespace tao::pq
       m_connection->send_params( statement, n_params, types, values, lengths, formats );
    }
 
+   void transaction::set_single_row_mode()
+   {
+      check_current_transaction();
+      if( PQsetSingleRowMode( m_connection->underlying_raw_ptr() ) == 0 ) {
+         throw std::runtime_error( "unable to switch to single row mode" );
+      }
+   }
+
+#if defined( LIBPQ_HAS_CHUNK_MODE )
+   void transaction::set_chunk_mode( const std::size_t rows )
+   {
+      check_current_transaction();
+      if( PQsetChunkedRowsMode( m_connection->underlying_raw_ptr(), rows ) == 0 ) {
+         throw std::runtime_error( "unable to switch to chunk mode" );
+      }
+   }
+#endif
+
    auto transaction::get_result( const std::chrono::steady_clock::time_point start ) -> result
    {
       check_current_transaction();
@@ -159,6 +177,12 @@ namespace tao::pq
                m_connection->clear_results( end );
                throw std::runtime_error( "unexpected COPY TO statement" );
 
+            case PGRES_SINGLE_TUPLE:
+#if defined( LIBPQ_HAS_CHUNK_MODE )
+            case PGRES_TUPLES_CHUNK:
+#endif
+               return pq::result( result.release() );
+
             default:;
          }
          while( auto next = m_connection->get_result( end ) ) {