Support preferable location for connections

ydb-platform · Feb 4, 2025 · 264d712 · 264d712
1 parent 6aaa1a3
commit 264d712
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbConfig.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbConfig.java
@@ -26,6 +26,7 @@ public static YdbConfig createForTesting(String host, int port, String tablespac
                 database,
                 null,
                 HostAndPort.fromParts(host, port),
+                null,
                 SESSION_CREATE_TIMEOUT_DEFAULT,
                 SESSION_CREATE_RETRY_COUNT_DEFAULT,
                 SESSION_KEEP_ALIVE_TIME_DEFAULT,
@@ -61,6 +62,13 @@ public static YdbConfig createForTesting(String host, int port, String tablespac
     @With
     HostAndPort hostAndPort;
 
+    /**
+     * Name of preferable location (data center) for connections.
+     * If it is not specified, all available cluster nodes will be used.
+     */
+    @With
+    String preferableLocation;
+
     @With
     Duration sessionCreationTimeout;
     @With

diff --git a/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java b/repository-ydb-v2/src/main/java/tech/ydb/yoj/repository/ydb/YdbRepository.java
@@ -10,6 +10,7 @@
 import org.slf4j.LoggerFactory;
 import tech.ydb.auth.AuthProvider;
 import tech.ydb.auth.NopAuthProvider;
+import tech.ydb.core.grpc.BalancingSettings;
 import tech.ydb.core.grpc.GrpcTransport;
 import tech.ydb.core.grpc.GrpcTransportBuilder;
 import tech.ydb.core.impl.SingleChannelTransport;
@@ -117,6 +118,10 @@ private static GrpcTransportBuilder makeGrpcTransportBuilder(@NonNull YdbConfig
             throw new IllegalArgumentException("one of [discoveryEndpoint, hostAndPort] must be set");
         }
 
+        if (config.getPreferableLocation() != null) {
+            transportBuilder.withBalancingSettings(BalancingSettings.fromLocation(config.getPreferableLocation()));
+        }
+
         if (config.isUseTLS()) {
             if (config.isUseTrustStore()) {
                 transportBuilder.withSecureConnection();