stream_query.hxx

/* Definition of the pqxx::internal::stream_query class.
 *
 * Enables optimized batch reads from a database table.
 *
 * DO NOT INCLUDE THIS FILE DIRECTLY; include pqxx/stream_query instead.
 *
 * Copyright (c) 2000-2024, Jeroen T. Vermeulen.
 *
 * See COPYING for copyright license.  If you did not receive a file called
 * COPYING with this source code, please notify the distributor of this
 * mistake, or contact the author.
 */
#ifndef PQXX_H_STREAM_QUERY
#define PQXX_H_STREAM_QUERY
 
#if !defined(PQXX_HEADER_PRE)
#  error "Include libpqxx headers as <pqxx/header>, not <pqxx/header.hxx>."
#endif
 
#include <cassert>
#include <functional>
#include <variant>
 
#include "pqxx/connection.hxx"
#include "pqxx/except.hxx"
#include "pqxx/internal/concat.hxx"
#include "pqxx/internal/encoding_group.hxx"
#include "pqxx/internal/encodings.hxx"
#include "pqxx/internal/gates/connection-stream_from.hxx"
#include "pqxx/internal/stream_iterator.hxx"
#include "pqxx/separated_list.hxx"
#include "pqxx/transaction_base.hxx"
#include "pqxx/transaction_focus.hxx"
#include "pqxx/util.hxx"
 
 
namespace pqxx
{
class transaction_base;
} // namespace pqxx
 
 
namespace pqxx::internal
{
class stream_query_end_iterator
{};
 
 
// C++20: Can we use generators, and maybe get speedup from HALO?
 
template<typename... TYPE> class stream_query : transaction_focus
{
public:
  using line_handle = std::unique_ptr<char, void (*)(void const *)>;
 
  inline stream_query(transaction_base &tx, std::string_view query);
 
  stream_query(stream_query &&) = delete;
  stream_query &operator=(stream_query &&) = delete;
 
  ~stream_query() noexcept
  {
    try
    {
      close();
    }
    catch (std::exception const &e)
    {
      reg_pending_error(e.what());
    }
  }
 
  bool done() const & noexcept { return m_char_finder == nullptr; }
 
  inline auto begin() &;
 
  auto end() const & { return stream_query_end_iterator{}; }
 
  std::tuple<TYPE...> parse_line(zview line) &
  {
    assert(not done());
 
    auto const line_size{std::size(line)};
 
    // This function uses m_row as a buffer, across calls.  The only reason for
    // it to carry over across calls is to avoid reallocation.
 
    // Make room for unescaping the line.  It's a pessimistic size.
    // Unusually, we're storing terminating zeroes *inside* the string.
    // This is the only place where we modify m_row.  MAKE SURE THE BUFFER DOES
    // NOT GET RESIZED while we're working, because we're working with views
    // into its buffer.
    m_row.resize(line_size + 1);
 
    std::size_t offset{0u};
    char *write{m_row.data()};
 
    // DO NOT shrink m_row to fit.  We're carrying views pointing into the
    // buffer.  (Also, how useful would shrinking really be?)
 
    // Folding expression: scan and unescape each field, and convert it to its
    // requested type.
    std::tuple<TYPE...> data{parse_field<TYPE>(line, offset, write)...};
 
    assert(offset == line_size + 1u);
    return data;
  }
 
  std::pair<line_handle, std::size_t> read_line() &;
 
private:
 
  static inline char_finder_func *get_finder(transaction_base const &tx);
 
 
  std::tuple<std::size_t, char *, zview>
  read_field(zview line, std::size_t offset, char *write)
  {
#if !defined(NDEBUG)
    auto const line_size{std::size(line)};
#endif
 
    assert(offset <= line_size);
 
    char const *lp{std::data(line)};
 
    // The COPY line now ends in a tab.  (We replace the trailing newline with
    // that to simplify the loop here.)
    assert(lp[line_size] == '\t');
    assert(lp[line_size + 1] == '\0');
 
    if ((lp[offset] == '\\') and (lp[offset + 1] == 'N'))
    {
      // Null field.  Consume the "\N" and the field separator.
      offset += 3;
      assert(offset <= (line_size + 1));
      assert(lp[offset - 1] == '\t');
      // Return a null value.  There's nothing to write into m_row.
      return {offset, write, {}};
    }
 
    // Beginning of the field text in the row buffer.
    char const *const field_begin{write};
 
    // We're relying on several assumptions just for making the main loop
    // condition work:
    // * The COPY line ends in a newline.
    // * Multibyte characters never start with an ASCII-range byte.
    // * We can index a view beyond its bounds (but within its address space).
    //
    // Effectively, the newline acts as a final field separator.
    while (lp[offset] != '\t')
    {
      assert(lp[offset] != '\0');
 
      // Beginning of the next character of interest (or the end of the line).
      auto const stop_char{m_char_finder(line, offset)};
      PQXX_ASSUME(stop_char > offset);
      assert(stop_char < (line_size + 1));
 
      // Copy the text we have so far.  It's got no special characters in it.
      std::memcpy(write, &lp[offset], stop_char - offset);
      write += (stop_char - offset);
      offset = stop_char;
 
      // We're still within the line.
      char const special{lp[offset]};
      if (special == '\\')
      {
        // Escape sequence.
        // Consume the backslash.
        ++offset;
        assert(offset < line_size);
 
        // The database will only escape ASCII characters, so we assume that
        // we're dealing with a single-byte character.
        char const escaped{lp[offset]};
        assert((escaped >> 7) == 0);
        ++offset;
        *write++ = unescape_char(escaped);
      }
      else
      {
        // Field separator.  Fall out of the loop.
        assert(special == '\t');
      }
    }
 
    // Hit the end of the field.
    assert(lp[offset] == '\t');
    *write = '\0';
    ++write;
    ++offset;
    return {offset, write, {field_begin, write - field_begin - 1}};
  }
 
 
  template<typename TARGET>
  TARGET parse_field(zview line, std::size_t &offset, char *&write)
  {
    using field_type = strip_t<TARGET>;
    using nullity = nullness<field_type>;
 
    assert(offset <= std::size(line));
 
    auto [new_offset, new_write, text]{read_field(line, offset, write)};
    PQXX_ASSUME(new_offset > offset);
    PQXX_ASSUME(new_write >= write);
    offset = new_offset;
    write = new_write;
    if constexpr (nullity::always_null)
    {
      if (std::data(text) != nullptr)
        throw conversion_error{concat(
          "Streaming a non-null value into a ", type_name<field_type>,
          ", which must always be null.")};
    }
    else if (std::data(text) == nullptr)
    {
      if constexpr (nullity::has_null)
        return nullity::null();
      else
        internal::throw_null_conversion(type_name<field_type>);
    }
    else
    {
      // Don't ever try to convert a non-null value to nullptr_t!
      return from_string<field_type>(text);
    }
  }
 
  void close() noexcept
  {
    if (not done())
    {
      m_char_finder = nullptr;
      unregister_me();
    }
  }
 
 
  char_finder_func *m_char_finder;
 
 
  std::string m_row;
};
} // namespace pqxx::internal
#endif