feat(ydb): add load_balance_configuration and grpc_compression configuration

ydb/functional_tests/basic/static_config.yaml

@@ -78,6 +78,8 @@ components_manager:
                     max_pool_size: 10
                     min_pool_size: 5
                     get_session_retry_limit: 1
+                    grpc-compression-algorithm: gzip
+                    grpc-load-balancing-policy: round_robin
                 # Second logical database used by the YDB_DATABASE_ROUTING test.
                 # The testsuite ydb recipe points every database to the same
                 # physical instance, so routing here is observed via per-database

ydb/functional_tests/basic/tests/test_compression.py

@@ -0,0 +1,22 @@
+async def test_grpc_compression_algorithm_metric(service_client, monitor_client):
+    response = await service_client.post(
+        'ydb/upsert-row',
+        json={
+            'id': 'id-compression',
+            'name': 'name-compression',
+            'service': 'srv',
+            'channel': 123,
+        },
+    )
+    assert response.status_code == 200
+    assert response.json() == {}
+
+    metrics = await monitor_client.metrics(prefix='ydb')
+    assert (
+        metrics.value_at(
+            path='ydb.grpc-compression-algorithm',
+            labels={'ydb_database': 'sampledb', 'algorithm': 'gzip'},
+            default=0,
+        )
+        == 1
+    )

ydb/functional_tests/basic/tests/test_load_balancing.py

@@ -0,0 +1,20 @@
+async def test_grpc_load_balancing_policy_metric(service_client, monitor_client):
+    response = await service_client.post(
+        'ydb/select-rows',
+        json={
+            'service': 'srv',
+            'channels': [1, 2, 3],
+            'created': '2019-10-30T11:20:00+00:00',
+        },
+    )
+    assert response.status_code == 200
+
+    metrics = await monitor_client.metrics(prefix='ydb')
+    assert (
+        metrics.value_at(
+            path='ydb.grpc-load-balancing-policy',
+            labels={'ydb_database': 'sampledb', 'policy': 'round_robin'},
+            default=0,
+        )
+        == 1
+    )

ydb/include/userver/ydb/exceptions.hpp

@@ -23,6 +23,8 @@ class YdbResponseError : public BaseError {
 
     const NYdb::TStatus& GetStatus() const noexcept;
 
+    bool IsConstraintViolation() const noexcept;
+
 private:
     NYdb::TStatus status_;
 };

ydb/src/ydb/component.yaml

@@ -111,6 +111,19 @@ properties:
                     type: boolean
                     description: |
                         NYdb::TDriverConfig::SetGRpcKeepAlivePermitWithoutCalls. If omitted, not set (SDK default).
+                grpc-compression-algorithm:
+                    type: string
+                    enum: [none, gzip, deflate]
+                    default: none
+                    description: |
+                        NYdb::TDriverConfig::SetGRpcCompressionAlgorithm. If omitted, not set (SDK default).
+                        Valid values: "gzip", "deflate", "none".
+                grpc-load-balancing-policy:
+                    type: string
+                    default: round_robin
+                    description: |
+                        NYdb::TDriverConfig::SetGRpcLoadBalancingPolicy. If omitted, not set (SDK default).
+                        Valid values: "round_robin", "pick_first"
                 prefer_local_dc:
                     type: boolean
                     default: true

ydb/src/ydb/exceptions.cpp

@@ -15,6 +15,10 @@ YdbResponseError::YdbResponseError(std::string_view operation_name, NYdb::TStatu
 
 const NYdb::TStatus& YdbResponseError::GetStatus() const noexcept { return status_; }
 
+bool YdbResponseError::IsConstraintViolation() const noexcept {
+    return NYdb::NStatusHelpers::StatusContainsIssueWithCode(status_, NYdb::NIssue::CONSTRAINT_VIOLATION);
+}
+
 UndefinedDatabaseError::UndefinedDatabaseError(std::string error)
     : BaseError(std::move(error))
 {}

ydb/src/ydb/impl/config.cpp

@@ -3,6 +3,7 @@
 #include <userver/formats/json/serialize.hpp>
 #include <userver/formats/json/value.hpp>
 #include <userver/formats/parse/common_containers.hpp>
+#include <userver/utils/assert.hpp>
 #include <userver/utils/retry_budget.hpp>
 #include <userver/yaml_config/yaml_config.hpp>
 
@@ -56,6 +57,31 @@ TableSettings ParseTableSettings(const yaml_config::YamlConfig& dbconfig, const
     return result;
 }
 
+NYdb::EGrpcCompressionAlgorithm ToCompressionAlgorithm(std::string_view alg) {
+    if (alg == "gzip") {
+        return NYdb::EGrpcCompressionAlgorithm::Gzip;
+    } else if (alg == "deflate") {
+        return NYdb::EGrpcCompressionAlgorithm::Deflate;
+    } else if (alg == "none") {
+        return NYdb::EGrpcCompressionAlgorithm::None;
+    } else {
+        throw yaml_config::Exception(fmt::format("Unknown grpc-compression-algorithm: {}", alg));
+    }
+
+}
+
+std::string_view ToString(NYdb::EGrpcCompressionAlgorithm algorithm) {
+    switch (algorithm) {
+        case NYdb::EGrpcCompressionAlgorithm::None:
+            return "none";
+        case NYdb::EGrpcCompressionAlgorithm::Deflate:
+            return "deflate";
+        case NYdb::EGrpcCompressionAlgorithm::Gzip:
+            return "gzip";
+    }
+    UINVARIANT(false, "Unhandled EGrpcCompressionAlgorithm value");
+}
+
 DriverSettings ParseDriverSettings(
     const yaml_config::YamlConfig& dbconfig,
     const secdist::DatabaseSettings& dbsecdist,
@@ -83,6 +109,11 @@ DriverSettings ParseDriverSettings(
     result.grpc_keepalive_permit_without_calls =
         dbconfig["grpc-keepalive-permit-without-calls"].As<std::optional<bool>>();
 
+    if (auto alg = dbconfig["grpc-compression-algorithm"].As<std::optional<std::string>>()) {                                                                                                                                       
+        result.grpc_compression_algorithm = ToCompressionAlgorithm(*alg);                                                                                                                                                           
+    }
+    result.grpc_load_balancing_policy = dbconfig["grpc-load-balancing-policy"].As<std::optional<std::string>>();
+
     result.endpoint = MergeWithSecdist(dbsecdist.endpoint, std::move(config_endpoint), dbconfig, "endpoint");
     result.database = MergeWithSecdist(dbsecdist.database, std::move(config_database), dbconfig, "database");
     result.oauth_token = dbsecdist.oauth_token;

ydb/src/ydb/impl/config.hpp

@@ -15,6 +15,7 @@
 #include <userver/yaml_config/fwd.hpp>
 
 #include <userver/ydb/settings.hpp>
+#include <ydb-cpp-sdk/client/types/ydb.h>
 
 USERVER_NAMESPACE_BEGIN
 
@@ -58,6 +59,9 @@ struct DriverSettings {
     std::optional<std::chrono::milliseconds> grpc_keepalive_timeout{};
     std::optional<bool> grpc_keepalive_permit_without_calls{};
 
+    std::optional<NYdb::EGrpcCompressionAlgorithm> grpc_compression_algorithm{};
+    std::optional<std::string> grpc_load_balancing_policy{};
+
     bool prefer_local_dc{false};
     std::optional<std::string> oauth_token;
     std::optional<std::string> iam_jwt_params;
@@ -67,6 +71,8 @@ struct DriverSettings {
     std::shared_ptr<NYdb::ICredentialsProviderFactory> credentials_provider_factory;
 };
 
+std::string_view ToString(NYdb::EGrpcCompressionAlgorithm algorithm);
+
 TableSettings ParseTableSettings(const yaml_config::YamlConfig& dbconfig, const secdist::DatabaseSettings& dbsecdist);
 
 DriverSettings ParseDriverSettings(

ydb/src/ydb/impl/driver.cpp

@@ -18,6 +18,8 @@ namespace ydb::impl {
 Driver::Driver(std::string dbname, impl::DriverSettings settings)
     : dbname_(std::move(dbname)),
       dbpath_(settings.database),
+      grpc_compression_algorithm_(settings.grpc_compression_algorithm),
+      grpc_load_balancing_policy_(settings.grpc_load_balancing_policy),
       native_metrics_(std::make_unique<NMonitoring::TMetricRegistry>(NMonitoring::TLabels{})),
       retry_budget_(utils::RetryBudgetSettings{})
 {
@@ -70,6 +72,13 @@ Driver::Driver(std::string dbname, impl::DriverSettings settings)
         driver_config.SetGRpcKeepAlivePermitWithoutCalls(*settings.grpc_keepalive_permit_without_calls);
     }
 
+    if (settings.grpc_compression_algorithm.has_value()) {
+        driver_config.SetGRpcCompressionAlgorithm(*settings.grpc_compression_algorithm);
+    }
+    if (settings.grpc_load_balancing_policy.has_value()) {
+        driver_config.SetGRpcLoadBalancingPolicy(*settings.grpc_load_balancing_policy);
+    }
+
     AppendUserverYdbBuildInfo(driver_config);
 
     driver_ = std::make_unique<NYdb::TDriver>(driver_config);
@@ -86,7 +95,17 @@ const std::string& Driver::GetDbPath() const { return dbpath_; }
 
 utils::RetryBudget& Driver::GetRetryBudget() { return retry_budget_; }
 
-void DumpMetric(utils::statistics::Writer& writer, const Driver& driver) { writer["native"] = *driver.native_metrics_; }
+void DumpMetric(utils::statistics::Writer& writer, const Driver& driver) {
+    writer["native"] = *driver.native_metrics_;
+    if (driver.grpc_compression_algorithm_.has_value()) {
+        writer["grpc-compression-algorithm"].ValueWithLabels(
+            1, {{"algorithm", ToString(*driver.grpc_compression_algorithm_)}}
+        );
+    }
+    if (driver.grpc_load_balancing_policy_.has_value()) {
+        writer["grpc-load-balancing-policy"].ValueWithLabels(1, {{"policy", *driver.grpc_load_balancing_policy_}});
+    }
+}
 
 std::string JoinPath(std::string_view database_path, std::string_view path) {
     UASSERT(!database_path.ends_with("/"));

ydb/src/ydb/impl/driver.hpp

@@ -1,12 +1,14 @@
 #pragma once
 
 #include <memory>
+#include <optional>
 #include <string>
 
 #include <userver/utils/retry_budget.hpp>
 #include <userver/utils/statistics/fwd.hpp>
 
 #include <ydb-cpp-sdk/client/driver/fwd.h>
+#include <ydb-cpp-sdk/client/types/ydb.h>
 
 namespace NMonitoring {
 class TMetricRegistry;
@@ -44,6 +46,9 @@ class Driver final {
     const std::string dbname_;
     const std::string dbpath_;
 
+    const std::optional<NYdb::EGrpcCompressionAlgorithm> grpc_compression_algorithm_;
+    const std::optional<std::string> grpc_load_balancing_policy_;
+
     std::unique_ptr<NMonitoring::TMetricRegistry> native_metrics_;
     // The retry_budget_ is used in driver_ threads, so it must be before the
     // driver_

ydb/tests/driver_config_test.cpp

@@ -0,0 +1,62 @@
+#include <userver/formats/yaml.hpp>
+#include <userver/utest/utest.hpp>
+#include <userver/yaml_config/yaml_config.hpp>
+
+#include <ydb/impl/config.hpp>
+#include <ydb/impl/secdist.hpp>
+
+USERVER_NAMESPACE_BEGIN
+
+namespace {
+
+ydb::impl::DriverSettings DriverSettingsFromYaml(std::string_view yaml) {
+    const yaml_config::YamlConfig config{formats::yaml::FromString(std::string{yaml}), {}};
+    ydb::impl::secdist::DatabaseSettings secdist;
+    secdist.endpoint = "localhost:2136";
+    secdist.database = "local";
+    return ydb::impl::ParseDriverSettings(config, secdist, nullptr);
+}
+
+}  // namespace
+
+UTEST(YdbDriverConfig, GrpcCompressionAlgorithmGzip) {
+    const auto settings = DriverSettingsFromYaml("grpc-compression-algorithm: gzip");
+    ASSERT_TRUE(settings.grpc_compression_algorithm.has_value());
+    EXPECT_EQ(*settings.grpc_compression_algorithm, NYdb::EGrpcCompressionAlgorithm::Gzip);
+}
+
+UTEST(YdbDriverConfig, GrpcCompressionAlgorithmDeflate) {
+    const auto settings = DriverSettingsFromYaml("grpc-compression-algorithm: deflate");
+    ASSERT_TRUE(settings.grpc_compression_algorithm.has_value());
+    EXPECT_EQ(*settings.grpc_compression_algorithm, NYdb::EGrpcCompressionAlgorithm::Deflate);
+}
+
+UTEST(YdbDriverConfig, GrpcCompressionAlgorithmNone) {
+    const auto settings = DriverSettingsFromYaml("grpc-compression-algorithm: none");
+    ASSERT_TRUE(settings.grpc_compression_algorithm.has_value());
+    EXPECT_EQ(*settings.grpc_compression_algorithm, NYdb::EGrpcCompressionAlgorithm::None);
+}
+
+UTEST(YdbDriverConfig, GrpcCompressionAlgorithmUnknownThrows) {
+    EXPECT_THROW(DriverSettingsFromYaml("grpc-compression-algorithm: brotli"), yaml_config::Exception);
+}
+
+UTEST(YdbDriverConfig, GrpcLoadBalancingPolicyRoundRobin) {
+    const auto settings = DriverSettingsFromYaml("grpc-load-balancing-policy: round_robin");
+    ASSERT_TRUE(settings.grpc_load_balancing_policy.has_value());
+    EXPECT_EQ(*settings.grpc_load_balancing_policy, "round_robin");
+}
+
+UTEST(YdbDriverConfig, GrpcLoadBalancingPolicyPickFirst) {
+    const auto settings = DriverSettingsFromYaml("grpc-load-balancing-policy: pick_first");
+    ASSERT_TRUE(settings.grpc_load_balancing_policy.has_value());
+    EXPECT_EQ(*settings.grpc_load_balancing_policy, "pick_first");
+}
+
+UTEST(YdbDriverConfig, MissingKeysLeaveSettingsUnset) {
+    const auto settings = DriverSettingsFromYaml("max_pool_size: 10");
+    EXPECT_FALSE(settings.grpc_compression_algorithm.has_value());
+    EXPECT_FALSE(settings.grpc_load_balancing_policy.has_value());
+}
+
+USERVER_NAMESPACE_END