postgresql_async_streaming_source.hpp

Go to the documentation of this file.

#pragma once
 
#include "../results/streaming_result.hpp"
#include "postgresql_async_connection.hpp"
 
#include <memory>
#include <optional>
#include <string>
#include <vector>
 
#include <boost/asio.hpp>
#include <boost/asio/awaitable.hpp>
 
// Forward declarations
struct pg_conn;
using PGconn = pg_conn;
 
namespace relx::connection {
 
template <typename T>
struct is_awaitable : std::false_type {};
 
template <typename T>
struct is_awaitable<boost::asio::awaitable<T>> : std::true_type {};
 
template <typename T>
inline constexpr bool is_awaitable_v = is_awaitable<T>::value;
 
class PostgreSQLAsyncStreamingSource {
public:
  PostgreSQLAsyncStreamingSource(PostgreSQLAsyncConnection& connection, std::string sql,
                                 std::vector<std::string> params = {});
 
  ~PostgreSQLAsyncStreamingSource();
 
  // Disable copy operations
  PostgreSQLAsyncStreamingSource(const PostgreSQLAsyncStreamingSource&) = delete;
  PostgreSQLAsyncStreamingSource& operator=(const PostgreSQLAsyncStreamingSource&) = delete;
 
  // Allow move operations
  PostgreSQLAsyncStreamingSource(PostgreSQLAsyncStreamingSource&& other) noexcept;
  PostgreSQLAsyncStreamingSource& operator=(PostgreSQLAsyncStreamingSource&& other) noexcept;
 
  boost::asio::awaitable<ConnectionResult<void>> initialize();
 
  boost::asio::awaitable<std::optional<std::string>> get_next_row();
 
  const std::vector<std::string>& get_column_names() const;
 
  bool is_initialized() const { return initialized_; }
 
  bool has_more_rows() const { return !finished_; }
 
  boost::asio::awaitable<void> async_cleanup();
 
  PostgreSQLAsyncConnection& get_connection() { return connection_; }
 
private:
  PostgreSQLAsyncConnection& connection_;
  std::string sql_;
  std::vector<std::string> params_;
 
  std::vector<std::string> column_names_;
  std::vector<bool> is_bytea_column_;
  bool initialized_;
  bool finished_;
  bool convert_bytea_;
 
  // Track query state
  bool query_active_;
 
  // Cache for the first row (since we consume it during metadata processing)
  std::optional<std::string> first_row_cached_;
 
  // Streaming state for proper result management
  std::unique_ptr<struct pg_result, void (*)(struct pg_result*)> current_result_;
  size_t current_row_index_;
  bool has_pending_results_;
 
  boost::asio::awaitable<ConnectionResult<void>> start_query();
 
  void process_column_metadata_from_pg_result(struct pg_result* pg_result);
 
  std::string format_single_row(struct pg_result* pg_result);
 
  std::string convert_pg_bytea_to_binary(const std::string& hex_value) const;
 
  void cleanup();
};
 
template <typename DataSource>
class AsyncStreamingResultSet {
public:
  class async_streaming_iterator {
  public:
    async_streaming_iterator(DataSource& source, AsyncStreamingResultSet& result_set,
                             bool at_end = false)
        : source_(source), result_set_(result_set), current_row_(), at_end_(at_end) {}
 
    const auto& operator*() const { return current_row_; }
 
    boost::asio::awaitable<void> advance() {
      if (at_end_) {
        co_return;
      }
 
      auto next_row_data = co_await source_.get_next_row();
      if (next_row_data) {
        current_row_ = result::LazyRow(std::move(*next_row_data), source_.get_column_names());
      } else {
        at_end_ = true;
        // Automatically reset connection state when streaming completes
        co_await result_set_.auto_reset_connection_state();
      }
    }
 
    bool is_at_end() const { return at_end_; }
 
  private:
    DataSource& source_;
    AsyncStreamingResultSet& result_set_;
    result::LazyRow current_row_;
    bool at_end_;
  };
 
  AsyncStreamingResultSet(DataSource source) : source_(std::move(source)), reset_called_(false) {}
 
  ~AsyncStreamingResultSet() {
    // Call synchronous cleanup for RAII when going out of scope
    if (!reset_called_) {
      reset_called_ = true;
      // Only call reset for PostgreSQL async streaming sources
      if constexpr (std::is_same_v<DataSource, PostgreSQLAsyncStreamingSource>) {
        source_.get_connection().reset_connection_state_sync();
      }
    }
  }
 
  // Disable copy operations to ensure single ownership
  AsyncStreamingResultSet(const AsyncStreamingResultSet&) = delete;
  AsyncStreamingResultSet& operator=(const AsyncStreamingResultSet&) = delete;
 
  // Allow move operations
  AsyncStreamingResultSet(AsyncStreamingResultSet&& other) noexcept
      : source_(std::move(other.source_)), reset_called_(other.reset_called_) {
    other.reset_called_ = true;  // Prevent cleanup in moved-from object
  }
 
  AsyncStreamingResultSet& operator=(AsyncStreamingResultSet&& other) noexcept {
    if (this != &other) {
      // Cleanup current object before move
      if (!reset_called_) {
        reset_called_ = true;
        if constexpr (std::is_same_v<DataSource, PostgreSQLAsyncStreamingSource>) {
          source_.get_connection().reset_connection_state_sync();
        }
      }
 
      source_ = std::move(other.source_);
      reset_called_ = other.reset_called_;
      other.reset_called_ = true;  // Prevent cleanup in moved-from object
    }
    return *this;
  }
 
  async_streaming_iterator begin() { return async_streaming_iterator(source_, *this); }
 
  async_streaming_iterator end() { return async_streaming_iterator(source_, *this, true); }
 
  template <typename Func>
  boost::asio::awaitable<void> for_each(Func&& func) {
    auto it = begin();
 
    // Advance to first row
    co_await it.advance();
 
    while (!it.is_at_end()) {
      // Handle different function signatures
      using ReturnType = std::invoke_result_t<Func, decltype(*it)>;
 
      if constexpr (is_awaitable_v<ReturnType>) {
        // Asynchronous function - return type is boost::asio::awaitable<T>
 
        // Extract the value type from the awaitable
        using AwaitableValueType = typename ReturnType::value_type;
 
        if constexpr (std::is_same_v<AwaitableValueType, bool>) {
          // Async function returning bool - check for early termination
          bool result = co_await func(*it);
          if (result) {
            // Early termination - need to manually reset connection state
            co_await auto_reset_connection_state();
            co_return;
          }
        } else {
          // Async function returning void - just await it
          co_await func(*it);
        }
      } else {
        // Synchronous function
        if constexpr (std::is_same_v<ReturnType, bool>) {
          // Sync function returning bool - check for early termination
          if (func(*it)) {
            // Early termination - need to manually reset connection state
            co_await auto_reset_connection_state();
            co_return;
          }
        } else {
          // Sync function returning void
          func(*it);
        }
      }
 
      co_await it.advance();
    }
 
    // Connection state reset is automatically called when iterator reaches end via advance()
  }
 
  boost::asio::awaitable<void> cleanup() {
    co_await source_.async_cleanup();
    co_await auto_reset_connection_state();
  }
 
  boost::asio::awaitable<void> auto_reset_connection_state() {
    if (!reset_called_) {
      reset_called_ = true;
      // Only call reset for PostgreSQL async streaming sources
      if constexpr (std::is_same_v<DataSource, PostgreSQLAsyncStreamingSource>) {
        auto reset_result = co_await source_.get_connection().reset_connection_state();
        // Note: We don't propagate reset errors since streaming operation was successful
        (void)reset_result;  // Silence unused variable warning
      }
    }
  }
 
private:
  DataSource source_;
  bool reset_called_;
};
 
template <typename... Args>
AsyncStreamingResultSet<PostgreSQLAsyncStreamingSource> create_async_streaming_result(
    PostgreSQLAsyncConnection& connection, const std::string& sql, Args&&... args) {
  // Convert parameters to strings (similar to execute_typed)
  std::vector<std::string> param_strings;
  if constexpr (sizeof...(Args) > 0) {
    param_strings.reserve(sizeof...(Args));
 
    auto add_param = [&param_strings](auto&& param) {
      using ParamType = std::remove_cvref_t<decltype(param)>;
 
      if constexpr (std::is_same_v<ParamType, std::nullptr_t>) {
        param_strings.push_back("NULL");
      } else if constexpr (std::is_same_v<ParamType, std::string> ||
                           std::is_same_v<ParamType, const char*> ||
                           std::is_same_v<ParamType, std::string_view>) {
        param_strings.push_back(std::string(param));
      } else if constexpr (std::is_arithmetic_v<ParamType>) {
        param_strings.push_back(std::to_string(param));
      } else if constexpr (std::is_same_v<ParamType, bool>) {
        param_strings.push_back(param ? "t" : "f");
      } else {
        std::ostringstream ss;
        ss << param;
        param_strings.push_back(ss.str());
      }
    };
 
    (add_param(std::forward<Args>(args)), ...);
  }
 
  return AsyncStreamingResultSet<PostgreSQLAsyncStreamingSource>(
      PostgreSQLAsyncStreamingSource(connection, sql, std::move(param_strings)));
}
 
}  // namespace relx::connection