table.hpp

#pragma once
 
#include <ydb-cpp-sdk/client/query/client.h>
#include <ydb-cpp-sdk/client/query/query.h>
#include <ydb-cpp-sdk/client/table/table.h>
 
#include <userver/dynamic_config/source.hpp>
#include <userver/utils/statistics/fwd.hpp>
 
#include <userver/ydb/builder.hpp>
#include <userver/ydb/query.hpp>
#include <userver/ydb/response.hpp>
#include <userver/ydb/settings.hpp>
#include <userver/ydb/transaction.hpp>
 
namespace NMonitoring {
class TMetricRegistry;
}  // namespace NMonitoring
 
USERVER_NAMESPACE_BEGIN
 
namespace tracing {
class Span;
}  // namespace tracing
 
namespace utils {
class RetryBudget;
}  // namespace utils
 
namespace ydb {
 
namespace impl {
struct Stats;
struct TableSettings;
class Driver;
template <typename Settings>
class RequestContext;
enum class IsStreaming : bool {};
}  // namespace impl
 
using DescribePathSettings = NYdb::NScheme::TDescribePathSettings;
using ListDirectorySettings = NYdb::NScheme::TListDirectorySettings;
using MakeDirectorySettings = NYdb::NScheme::TMakeDirectorySettings;
using RemoveDirectorySettings = NYdb::NScheme::TRemoveDirectorySettings;
 
using BulkUpsertSettings = NYdb::NTable::TBulkUpsertSettings;
using CreateTableSettings = NYdb::NTable::TCreateTableSettings;
using DescribeTableSettings = NYdb::NTable::TDescribeTableSettings;
using DropTableSettings = NYdb::NTable::TDropTableSettings;
using ScanQuerySettings = NYdb::NTable::TStreamExecScanQuerySettings;
 
/// @brief A dynamic transaction name for @see TableClient::Begin or
/// @see TableClient::RetryTx.
///
/// @warning Make sure that transaction name has low cardinality.
/// If transaction name is unique for every call, per-transaction metrics will overflow metrics quota,
/// and metrics will become unusable.
using DynamicTransactionName = utils::StrongTypedef<struct DynamicTransactionNameTag, std::string>;
 
class TableClient final {
public:
    /// @cond
    // For internal use only.
    TableClient(
        impl::TableSettings settings,
        OperationSettings operation_settings,
        dynamic_config::Source config_source,
        std::shared_ptr<impl::Driver> driver
    );
 
    ~TableClient();
    /// @endcond
 
    /// Query for creating/deleting tables
    void ExecuteSchemeQuery(const std::string& query);
 
    void MakeDirectory(const std::string& path, MakeDirectorySettings query_settings = {});
    void RemoveDirectory(const std::string& path, RemoveDirectorySettings query_settings = {});
 
    NYdb::NScheme::TDescribePathResult DescribePath(std::string_view path, DescribePathSettings query_settings = {});
    NYdb::NScheme::TListDirectoryResult ListDirectory(std::string_view path, ListDirectorySettings query_settings = {});
 
    NYdb::NTable::TDescribeTableResult DescribeTable(std::string_view path, DescribeTableSettings query_settings = {});
    void CreateTable(
        std::string_view path,
        NYdb::NTable::TTableDescription&& table_desc,
        CreateTableSettings query_settings = {}
    );
    void DropTable(std::string_view path, DropTableSettings query_settings = {});
 
    /// @name Data queries execution
    /// Execute a single data query outside of transactions. Query parameters are
    /// passed in `Args` as "string key - value" pairs:
    ///
    /// @code
    /// client.ExecuteDataQuery(query, "name1", value1, "name2", value2, ...);
    /// @endcode
    ///
    /// @warning ExecuteDataQuery returns no more than 1000 rows. Consider @ref ExecuteQuery() instead, but make sure
    ///          that it works for your case.
    ///
    /// Use ydb::PreparedArgsBuilder for storing a generic buffer of query params if needed.
    ///
    /// It is convenient to keep YQL queries in separate files, see @ref scripts/docs/en/userver/sql_files.md
    /// for more info.
    ///
    /// @{
    template <typename... Args>
    ExecuteResponse ExecuteDataQuery(const Query& query, Args&&... args);
 
    template <typename... Args>
    ExecuteResponse ExecuteDataQuery(OperationSettings settings, const Query& query, Args&&... args);
 
    ExecuteResponse ExecuteDataQuery(OperationSettings settings, const Query& query, PreparedArgsBuilder&& builder);
 
    ExecuteResponse ExecuteDataQuery(
        QuerySettings query_settings,
        OperationSettings settings,
        const Query& query,
        PreparedArgsBuilder&& builder
    );
    /// @}
 
    /// @name Transactions with retry
    /// @brief Execute a transactional function with automatic retries.
    ///
    /// The user-provided function receives a TxActor for executing queries
    /// and returns TxAction::kCommit or TxAction::kRollback. On transient
    /// errors the whole function is retried automatically.
    ///
    /// @code
    /// client.RetryTx("my_tx", {.retries = 3},
    ///     [](ydb::TxActor& tx) {
    ///         tx.Execute(query, "$id", 1);
    ///         return ydb::TxAction::kCommit;
    ///     });
    /// @endcode
    ///
    /// @{
    void RetryTx(utils::StringLiteral transaction_name, RetryTxSettings retry_settings, RetryTxFunction fn);
 
    /// @warning Make sure that `transaction_name` has low cardinality.
    void RetryTx(DynamicTransactionName transaction_name, RetryTxSettings retry_settings, RetryTxFunction fn);
    /// @}
 
    /// @name Transactions (deprecated)
    /// @brief Begin a transaction with the specified name. The settings are used
    /// for the `BEGIN` statement.
    /// @deprecated Use RetryTx instead for automatic retry support.
    /// @see ydb::Transaction
    ///
    /// @{
    Transaction Begin(utils::StringLiteral transaction_name, OperationSettings settings = {});
 
    /// @warning Make sure that `transaction_name` has low cardinality.
    /// If `transaction_name` is unique for every call, per-transaction metrics will overflow metrics quota,
    /// and metrics will become unusable.
    Transaction Begin(DynamicTransactionName transaction_name, OperationSettings settings = {});
 
    Transaction Begin(utils::StringLiteral transaction_name, TransactionMode tx_mode);
    /// @}
 
    /// Builder for storing dynamic query params.
    PreparedArgsBuilder GetBuilder() const;
 
    /// Efficiently write large ranges of table data.
    void BulkUpsert(
        std::string_view table,
        NYdb::TValue&& rows,
        OperationSettings settings = {},
        BulkUpsertSettings query_settings = {}
    );
 
    /// Efficiently write large ranges of table data.
    /// The passed range of structs is serialized to TValue.
    template <typename RangeOfStructs>
    void BulkUpsert(std::string_view table, const RangeOfStructs& rows, OperationSettings settings = {});
 
    /// Efficiently read large ranges of table data.
    ReadTableResults ReadTable(
        std::string_view table,
        NYdb::NTable::TReadTableSettings&& read_settings = {},
        OperationSettings settings = {}
    );
 
    /// @name Scan queries execution
    /// A separate data access interface designed primarily for performing
    /// analytical ad-hoc queries.
    /// @{
    template <typename... Args>
    ScanQueryResults ExecuteScanQuery(const Query& query, Args&&... args);
 
    template <typename... Args>
    ScanQueryResults ExecuteScanQuery(
        ScanQuerySettings&& scan_settings,
        OperationSettings settings,
        const Query& query,
        Args&&... args
    );
 
    ScanQueryResults ExecuteScanQuery(
        ScanQuerySettings&& scan_settings,
        OperationSettings settings,
        const Query& query,
        PreparedArgsBuilder&& builder
    );
    /// @}
 
    /// @name Queries execution (using YDB Query SDK)
    /// Execute a single query outside of transactions. Query parameters are
    /// passed in `Args` as "string key - value" pairs:
    ///
    /// @code
    /// client.ExecuteQuery(query, "name1", value1, "name2", value2, ...);
    /// @endcode
    ///
    /// Use ydb::PreparedArgsBuilder for storing a generic buffer of query params if needed.
    ///
    /// If both exec_settings and settings args are passed,
    /// exec_settings.client_timeout_ms and exec_settings.trace_id are ignored
    /// and are overwritten by settings.client_timeout_ms and settings.trace_id.
    ///
    /// It is convenient to keep YQL queries in separate files, see @ref scripts/docs/en/userver/sql_files.md
    /// for more info.
    /// @{
    template <typename... Args>
    ExecuteResponse ExecuteQuery(const Query& query, Args&&... args);
 
    template <typename... Args>
    ExecuteResponse ExecuteQuery(OperationSettings settings, const Query& query, Args&&... args);
 
    ExecuteResponse ExecuteQuery(OperationSettings settings, const Query& query, PreparedArgsBuilder&& builder);
 
    ExecuteResponse ExecuteQuery(
        NYdb::NQuery::TExecuteQuerySettings&& exec_settings,
        OperationSettings settings,
        const Query& query,
        PreparedArgsBuilder&& builder
    );
    /// @}
 
    /// @cond
    // For internal use only.
    friend void DumpMetric(utils::statistics::Writer& writer, const TableClient& table_client);
    /// @endcond
 
    /// Get native table or query client
    /// @warning Use with care! Facilities from
    /// `<core/include/userver/drivers/subscribable_futures.hpp>` can help with
    /// non-blocking wait operations.
    /// @{
    NYdb::NTable::TTableClient& GetNativeTableClient();
 
    NYdb::NQuery::TQueryClient& GetNativeQueryClient();
    /// @}
 
    utils::RetryBudget& GetRetryBudget();
 
private:
    friend class Transaction;
    friend class TxActor;
    template <typename Settings>
    friend class impl::RequestContext;
 
    std::string JoinDbPath(std::string_view path) const;
 
    void Select1();
 
    NYdb::NQuery::TExecuteQuerySettings ToExecuteQuerySettings(const QuerySettings& query_settings) const;
    NYdb::NTable::TExecDataQuerySettings ToExecDataQuerySettings(const QuerySettings& query_settings) const;
 
    template <typename... Args>
    PreparedArgsBuilder MakeBuilder(Args&&... args);
 
    // Func: (TSession, const std::string& full_path, const Settings&)
    //       -> NThreading::TFuture<T>
    //       OR
    //       (TTableClient&, const std::string& full_path, const Settings&)
    //       -> NThreading::TFuture<T>
    // ExecuteSchemeQueryImpl -> T
    template <typename QuerySettings, typename Func>
    auto ExecuteWithPathImpl(
        std::string_view path,
        std::string_view operation_name,
        OperationSettings settings,
        QuerySettings&& query_settings,
        Func&& func
    );
 
    dynamic_config::Source config_source_;
    const OperationSettings default_settings_;
    const bool keep_in_query_cache_;
    const bool use_query_client_;
    std::unique_ptr<impl::Stats> stats_;
    std::shared_ptr<impl::Driver> driver_;
    std::unique_ptr<NYdb::NScheme::TSchemeClient> scheme_client_;
    std::unique_ptr<NYdb::NTable::TTableClient> table_client_;
    std::unique_ptr<NYdb::NQuery::TQueryClient> query_client_;
};
 
template <typename... Args>
PreparedArgsBuilder TableClient::MakeBuilder(Args&&... args) {
    auto builder = GetBuilder();
    builder.AddParams(std::forward<Args>(args)...);
    return builder;
}
 
template <typename... Args>
ExecuteResponse TableClient::ExecuteDataQuery(const Query& query, Args&&... args) {
    return ExecuteDataQuery(OperationSettings{}, query, MakeBuilder(std::forward<Args>(args)...));
}
 
template <typename... Args>
ExecuteResponse TableClient::ExecuteDataQuery(OperationSettings settings, const Query& query, Args&&... args) {
    return ExecuteDataQuery(settings, query, MakeBuilder(std::forward<Args>(args)...));
}
 
template <typename RangeOfStructs>
void TableClient::BulkUpsert(std::string_view table, const RangeOfStructs& rows, OperationSettings settings) {
    NYdb::TValueBuilder builder;
    ydb::Write(builder, rows);
    BulkUpsert(table, builder.Build(), std::move(settings));
}
 
template <typename... Args>
ScanQueryResults TableClient::ExecuteScanQuery(const Query& query, Args&&... args) {
    return ExecuteScanQuery(ScanQuerySettings{}, OperationSettings{}, query, MakeBuilder(std::forward<Args>(args)...));
}
 
template <typename... Args>
ScanQueryResults TableClient::ExecuteScanQuery(
    ScanQuerySettings&& scan_settings,
    OperationSettings settings,
    const Query& query,
    Args&&... args
) {
    return ExecuteScanQuery(
        std::move(scan_settings),
        std::move(settings),
        query,
        MakeBuilder(std::forward<Args>(args)...)
    );
}
 
template <typename... Args>
ExecuteResponse TableClient::ExecuteQuery(const Query& query, Args&&... args) {
    return ExecuteQuery(OperationSettings{}, query, MakeBuilder(std::forward<Args>(args)...));
}
 
template <typename... Args>
ExecuteResponse TableClient::ExecuteQuery(OperationSettings settings, const Query& query, Args&&... args) {
    return ExecuteQuery(settings, query, MakeBuilder(std::forward<Args>(args)...));
}
 
}  // namespace ydb
 
USERVER_NAMESPACE_END