response.hpp

#pragma once
 
#include <ydb-cpp-sdk/client/query/client.h>
#include <ydb-cpp-sdk/client/result/result.h>
#include <ydb-cpp-sdk/client/table/table.h>
 
#include <cstddef>
#include <iterator>
#include <memory>
#include <optional>
#include <typeinfo>
 
#include <userver/utils/meta.hpp>
#include <userver/utils/not_null.hpp>
#include <userver/ydb/impl/cast.hpp>
#include <userver/ydb/io/insert_row.hpp>
#include <userver/ydb/io/list.hpp>
#include <userver/ydb/io/primitives.hpp>
#include <userver/ydb/io/traits.hpp>
#include <userver/ydb/types.hpp>
 
USERVER_NAMESPACE_BEGIN
 
namespace ydb {
 
namespace impl {
 
template <typename T>
struct StructRowParser;
 
struct ParseState final {
    explicit ParseState(const NYdb::TResultSet& result_set);
 
    NYdb::TResultSetParser parser;
    const std::type_info* row_type_id{nullptr};
    std::unique_ptr<std::size_t[]> cpp_to_ydb_field_mapping{};
};
 
}  // namespace impl
 
using ValueType = NYdb::EPrimitiveType;
 
class Cursor;
 
class Row final {
public:
    /// @cond
    // For internal use only.
    explicit Row(impl::ParseState& parse_state);
    /// @endcond
 
    Row(const Row&) = delete;
    Row(Row&&) noexcept = default;
    Row& operator=(const Row&) = delete;
    Row& operator=(Row&&) = delete;
 
    /// @brief Parses the whole row to `T`, which must be a struct type.
    /// ydb::kStructMemberNames must be specialized for `T`.
    ///
    /// @throws ydb::ColumnParseError on parsing error
    /// @throws ydb::ParseError on extra fields on C++ side
    /// @throws ydb::ParseError on extra fields on YDB side
    template <typename T>
    T As() &&;
 
    /// @brief Parses the specified column to `T`.
    /// `Get` can only be called once for each column.
    /// @throws ydb::BaseError on parsing error
    template <typename T>
    T Get(std::string_view column_name);
 
    /// @brief Parses the specified column to `T`.
    /// `Get` can only be called once for each column.
    /// @throws ydb::BaseError on parsing error
    template <typename T>
    T Get(std::size_t column_index);
 
private:
    NYdb::TValueParser& GetColumn(std::size_t index);
    NYdb::TValueParser& GetColumn(std::string_view name);
 
    void ConsumedColumnsCheck(std::size_t column_index);
 
    impl::ParseState& parse_state_;
    std::vector<bool> consumed_columns_;
};
 
class CursorIterator final {
public:
    using difference_type = std::ptrdiff_t;
    using value_type = Row;
    using reference = Row;
    using iterator_category = std::input_iterator_tag;
 
    CursorIterator() = default;
 
    CursorIterator(const CursorIterator&) = delete;
    CursorIterator(CursorIterator&&) noexcept = default;
    CursorIterator& operator=(const CursorIterator&) = delete;
    CursorIterator& operator=(CursorIterator&&) = default;
 
    Row operator*() const;
 
    CursorIterator& operator++();
 
    void operator++(int);
 
    bool operator==(const std::default_sentinel_t& other) const noexcept;
 
private:
    friend class Cursor;
 
    explicit CursorIterator(Cursor& cursor);
 
    impl::ParseState* parse_state_{nullptr};
};
 
class Cursor final {
public:
    /// @cond
    explicit Cursor(const NYdb::TResultSet& result_set);
    /// @endcond
 
    Cursor(const Cursor&) = delete;
    Cursor(Cursor&&) noexcept = default;
    Cursor& operator=(const Cursor&) = delete;
    Cursor& operator=(Cursor&&) noexcept = default;
 
    size_t ColumnsCount() const;
    size_t RowsCount() const;
    /// @throws EmptyResponseError if GetFirstRow() or begin() called before or
    /// cursor is empty
    Row GetFirstRow();
 
    /// @brief Extract first row
    /// @throws EmptyResponseError if @ref empty().
    /// @throws IgnoreResultsError if @ref size() > 1.
    Row GetSingleRow() &&;
 
    /// @brief Extract data into a container. Each row is parsed using @ref Row::As.
    template <typename Container>
    Container AsContainer() &&;
 
    /// @brief Extract first row into user type using @ref Row::As.
    /// @throws EmptyResponseError if @ref empty().
    template <typename T>
    T AsSingleRow() &&;
 
    /// @brief Extract first row into user type using @ref Row::As.
    /// @returns A single row result set if non empty result was returned, empty `std::optional` otherwise
    template <typename T>
    std::optional<T> AsOptionalSingleRow() &&;
 
    /// Returns true if response has been truncated to the database limit
    /// (currently 1000 rows)
    bool IsTruncated() const;
 
    bool empty() const;
    std::size_t size() const;
 
    CursorIterator begin();
    std::default_sentinel_t end();
 
private:
    friend class Row;
    friend class CursorIterator;
 
    bool truncated_;
    bool is_consumed_{false};
    // Row should not be invalidated after moving Cursor, hence the indirection.
    utils::UniqueRef<impl::ParseState> parse_state_;
};
 
class ExecuteResponse final {
public:
    /// @cond
    explicit ExecuteResponse(std::variant<NYdb::NQuery::TExecuteQueryResult, NYdb::NTable::TDataQueryResult>&&
                                 query_result);
    /// @endcond
 
    ExecuteResponse(const ExecuteResponse&) = delete;
    ExecuteResponse(ExecuteResponse&&) noexcept = default;
    ExecuteResponse& operator=(const ExecuteResponse&) = delete;
    ExecuteResponse& operator=(ExecuteResponse&&) = delete;
 
    std::size_t GetCursorCount() const;
    Cursor GetCursor(std::size_t index) const;
    Cursor GetSingleCursor() const;
 
    /// Query stats are only available if initially requested
    const std::optional<NYdb::NTable::TQueryStats>&  //
    GetQueryStats() const noexcept;
 
    /// Returns true if Execute used the server query cache
    bool IsFromServerQueryCache() const noexcept;
 
private:
    void EnsureResultSetsNotEmpty() const;
 
    std::optional<NYdb::NTable::TQueryStats> query_stats_;
    std::vector<NYdb::TResultSet> result_sets_;
};
 
class ReadTableResults final {
public:
    /// @cond
    explicit ReadTableResults(NYdb::NTable::TTablePartIterator iterator);
    /// @endcond
 
    std::optional<Cursor> GetNextResult();
 
    ReadTableResults(const ReadTableResults&) = delete;
    ReadTableResults(ReadTableResults&&) noexcept = default;
    ReadTableResults& operator=(const ReadTableResults&) = delete;
    ReadTableResults& operator=(ReadTableResults&&) = delete;
 
private:
    NYdb::NTable::TTablePartIterator iterator_;
};
 
class ScanQueryResults final {
    using TScanQueryPartIterator = NYdb::NTable::TScanQueryPartIterator;
 
public:
    using TScanQueryPart = NYdb::NTable::TScanQueryPart;
 
    /// @cond
    explicit ScanQueryResults(TScanQueryPartIterator iterator);
    /// @endcond
 
    std::optional<TScanQueryPart> GetNextResult();
 
    std::optional<Cursor> GetNextCursor();
 
    ScanQueryResults(const ScanQueryResults&) = delete;
    ScanQueryResults(ScanQueryResults&&) noexcept = default;
    ScanQueryResults& operator=(const ScanQueryResults&) = delete;
    ScanQueryResults& operator=(ScanQueryResults&&) = delete;
 
private:
    TScanQueryPartIterator iterator_;
};
 
template <typename T>
T Row::As() && {
    if (&typeid(T) != parse_state_.row_type_id) {
        parse_state_.cpp_to_ydb_field_mapping = impl::StructRowParser<T>::MakeCppToYdbFieldMapping(parse_state_.parser);
        parse_state_.row_type_id = &typeid(T);
    }
    return impl::StructRowParser<T>::ParseRow(parse_state_.parser, parse_state_.cpp_to_ydb_field_mapping);
}
 
template <typename T>
T Row::Get(std::string_view column_name) {
#ifndef NDEBUG
    ConsumedColumnsCheck(parse_state_.parser.ColumnIndex(impl::ToString(column_name)));
#endif
    auto& column = GetColumn(column_name);
    return Parse<T>(column, ParseContext{/*column_name=*/column_name});
}
 
template <typename T>
T Row::Get(std::size_t column_index) {
#ifndef NDEBUG
    ConsumedColumnsCheck(column_index);
#endif
    auto& column = GetColumn(column_index);
    const auto column_name = std::to_string(column_index);
    return Parse<T>(column, ParseContext{/*column_name=*/column_name});
}
 
template <typename Container>
Container Cursor::AsContainer() && {
    using ValueType = typename Container::value_type;
    Container c;
    if constexpr (meta::kIsReservable<Container>) {
        c.reserve(size());
    }
 
    auto inserter = meta::Inserter(c);
    for (Row row : *this) {
        *inserter = std::move(row).As<ValueType>();
        ++inserter;
    }
 
    return c;
}
 
template <typename T>
T Cursor::AsSingleRow() && {
    return std::move(*this).GetSingleRow().As<T>();
}
 
template <typename T>
std::optional<T> Cursor::AsOptionalSingleRow() && {
    if (empty()) {
        return std::nullopt;
    }
    return std::move(*this).AsSingleRow<T>();
}
 
}  // namespace ydb
 
USERVER_NAMESPACE_END