structs.hpp

#pragma once
 
#include <ydb-cpp-sdk/client/result/result.h>
#include <ydb-cpp-sdk/client/value/value.h>
 
#include <cstddef>
#include <memory>
#include <optional>
#include <string_view>
#include <tuple>
#include <type_traits>
 
#include <boost/pfr/core.hpp>
#include <boost/pfr/core_name.hpp>
 
#include <userver/utils/assert.hpp>
#include <userver/utils/constexpr_indices.hpp>
#include <userver/utils/enumerate.hpp>
#include <userver/utils/trivial_map.hpp>
 
#include <userver/ydb/exceptions.hpp>
#include <userver/ydb/impl/cast.hpp>
#include <userver/ydb/io/traits.hpp>
 
USERVER_NAMESPACE_BEGIN
 
namespace ydb {
 
namespace impl {
 
struct NotStruct final {};
 
template <typename T>
constexpr decltype(T::kYdbMemberNames) DetectStructMemberNames() noexcept {
  return T::kYdbMemberNames;
}
 
template <typename T, typename... Args>
constexpr NotStruct DetectStructMemberNames(Args&&...) noexcept {
  return NotStruct{};
}
 
}  // namespace impl
 
/// @see ydb::CustomMemberNames
struct CustomMemberName final {
  std::string_view cpp_name;
  std::string_view ydb_name;
};
 
/// @brief Specifies C++ to YDB struct member names mapping. It's enough to only
/// specify the names that are different between C++ and YDB.
/// @see ydb::kStructMemberNames
template <std::size_t N>
struct StructMemberNames final {
  CustomMemberName custom_names[N];
};
 
// clang-format off
StructMemberNames() -> StructMemberNames<0>;
// clang-format on
 
template <std::size_t N>
StructMemberNames(CustomMemberName (&&)[N]) -> StructMemberNames<N>;
 
/// @brief Customization point for YDB serialization of structs.
///
/// In order to get serialization for a struct, you need to define
/// `kYdbMemberNames` inside it:
///
/// @snippet ydb/small_table.hpp  struct sample
///
/// Field names can be overridden:
///
/// @snippet ydb/tests/struct_io_test.cpp  custom names
///
/// To enable YDB serialization for an external struct, specialize
/// ydb::kStructMemberNames for it:
///
/// @snippet ydb/tests/struct_io_test.cpp  external specialization
///
/// @warning The struct must not reside in an anonymous namespace, otherwise
/// struct member names detection will break.
///
/// For extra fields on C++ side, parsing throws ydb::ParseError.
/// For extra fields on YDB side, parsing throws ydb::ParseError.
template <typename T>
constexpr auto kStructMemberNames = impl::DetectStructMemberNames<T>();
 
namespace impl {
 
template <typename T, std::size_t N>
constexpr auto WithDeducedNames(const StructMemberNames<N>& given_names) {
  auto names = boost::pfr::names_as_array<T>();
  for (const CustomMemberName& entry : given_names.custom_names) {
    std::size_t count = 0;
    for (auto& name : names) {
      if (name == entry.cpp_name) {
        name = entry.ydb_name;
        ++count;
        break;
      }
    }
    if (count != 1) {
      throw BaseError(
          "In a StructMemberNames, each cpp_name must match the C++ name of "
          "exactly 1 member of struct T");
    }
  }
  return names;
}
 
template <typename T, std::size_t... Indices>
constexpr auto MakeTupleOfOptionals(std::index_sequence<Indices...>) {
  return std::tuple<
      std::optional<boost::pfr::tuple_element_t<Indices, T>>...>();
}
 
template <typename T, typename Tuple, std::size_t... Indices>
constexpr auto MakeFromTupleOfOptionals(Tuple&& tuple,
                                        std::index_sequence<Indices...>) {
  return T{*std::get<Indices>(std::forward<Tuple>(tuple))...};
}
 
}  // namespace impl
 
template <typename T>
struct ValueTraits<
    T, std::enable_if_t<!std::is_same_v<decltype(kStructMemberNames<T>),
                                        const impl::NotStruct>>> {
  static_assert(std::is_aggregate_v<T>);
  static constexpr auto kFieldNames =
      impl::WithDeducedNames<T>(kStructMemberNames<T>);
  static constexpr auto kFieldNamesSet = utils::MakeTrivialSet<kFieldNames>();
  static constexpr auto kFieldsCount = kFieldNames.size();
 
  static T Parse(NYdb::TValueParser& parser, const ParseContext& context) {
    auto parsed_fields =
        impl::MakeTupleOfOptionals<T>(std::make_index_sequence<kFieldsCount>{});
    parser.OpenStruct();
 
    while (parser.TryNextMember()) {
      const auto& field_name = parser.GetMemberName();
      const auto index = kFieldNamesSet.GetIndex(field_name);
      if (!index.has_value()) {
        throw ColumnParseError(
            context.column_name,
            fmt::format("Unexpected field name '{}' for '{}' struct type, "
                        "expected one of: {}",
                        field_name, compiler::GetTypeName<T>(),
                        fmt::join(kFieldNames, ", ")));
      }
      utils::WithConstexprIndex<kFieldsCount>(*index, [&](auto index_c) {
        auto& field = std::get<decltype(index_c)::value>(parsed_fields);
        using FieldType = typename std::decay_t<decltype(field)>::value_type;
        field.emplace(ydb::Parse<FieldType>(parser, context));
      });
    }
 
    parser.CloseStruct();
 
    std::string_view missing_field;
    utils::ForEachIndex<kFieldsCount>([&](auto index_c) {
      if (!std::get<decltype(index_c)::value>(parsed_fields).has_value()) {
        missing_field = kFieldNames[index_c];
      }
    });
    if (!missing_field.empty()) {
      throw ColumnParseError(
          context.column_name,
          fmt::format("Missing field '{}' for '{}' struct type", missing_field,
                      compiler::GetTypeName<T>()));
    }
 
    return impl::MakeFromTupleOfOptionals<T>(
        std::move(parsed_fields), std::make_index_sequence<kFieldsCount>{});
  }
 
  template <typename Builder>
  static void Write(NYdb::TValueBuilderBase<Builder>& builder, const T& value) {
    builder.BeginStruct();
    boost::pfr::for_each_field(value, [&](const auto& field, std::size_t i) {
      builder.AddMember(impl::ToString(kFieldNames[i]));
      ydb::Write(builder, field);
    });
    builder.EndStruct();
  }
 
  static NYdb::TType MakeType() {
    NYdb::TTypeBuilder builder;
    builder.BeginStruct();
    utils::ForEachIndex<kFieldsCount>([&](auto index_c) {
      builder.AddMember(
          impl::ToString(kFieldNames[index_c]),
          ValueTraits<boost::pfr::tuple_element_t<decltype(index_c)::value,
                                                  T>>::MakeType());
    });
    builder.EndStruct();
    return builder.Build();
  }
};
 
namespace impl {
 
template <typename T>
struct StructRowParser final {
  static_assert(std::is_aggregate_v<T>);
  static constexpr auto kFieldNames =
      impl::WithDeducedNames<T>(kStructMemberNames<T>);
  static constexpr auto kFieldsCount = kFieldNames.size();
 
  static std::unique_ptr<std::size_t[]> MakeCppToYdbFieldMapping(
      NYdb::TResultSetParser& parser) {
    auto result = std::make_unique<std::size_t[]>(kFieldsCount);
    for (const auto [pos, field_name] : utils::enumerate(kFieldNames)) {
      const auto column_index = parser.ColumnIndex(impl::ToString(field_name));
      if (column_index == -1) {
        throw ParseError(fmt::format("Missing column '{}' for '{}' struct type",
                                     field_name, compiler::GetTypeName<T>()));
      }
      result[pos] = static_cast<std::size_t>(column_index);
    }
 
    const auto columns_count = parser.ColumnsCount();
    UASSERT(columns_count >= kFieldsCount);
    if (columns_count != kFieldsCount) {
      throw ParseError(fmt::format(
          "Unexpected extra columns while parsing row to '{}' struct type",
          compiler::GetTypeName<T>()));
    }
 
    return result;
  }
 
  static T ParseRow(NYdb::TResultSetParser& parser,
                    const std::unique_ptr<std::size_t[]>& cpp_to_ydb_mapping) {
    return ParseRowImpl(parser, cpp_to_ydb_mapping,
                        std::make_index_sequence<kFieldsCount>{});
  }
 
  template <std::size_t... Indices>
  static T ParseRowImpl(
      NYdb::TResultSetParser& parser,
      const std::unique_ptr<std::size_t[]>& cpp_to_ydb_mapping,
      std::index_sequence<Indices...>) {
    return T{
        Parse<boost::pfr::tuple_element_t<Indices, T>>(
            parser.ColumnParser(cpp_to_ydb_mapping[Indices]),
            ParseContext{/*column_name=*/kFieldNames[Indices]})...,
    };
  }
};
 
}  // namespace impl
 
}  // namespace ydb
 
USERVER_NAMESPACE_END