Support for ip addresses discovery in C++ and Python sdk, KIKIMR-13572

ref:9032969
sync: https://proxy.sandbox.yandex-team.ru/2700957527

Author: snaury

ydb/aio/connection.py

@@ -13,6 +13,7 @@
     _get_request_timeout,
     _set_server_timeouts,
     _RpcState as RpcState,
+    EndpointOptions,
     channel_factory,
     YDB_DATABASE_HEADER,
     YDB_TRACE_ID_HEADER,
@@ -106,10 +107,17 @@ class Connection:
         "closing",
     )
 
-    def __init__(self, endpoint: str, driver_config: DriverConfig = None):
+    def __init__(
+        self,
+        endpoint: str,
+        driver_config: DriverConfig = None,
+        endpoint_options: EndpointOptions = None,
+    ):
         global _stubs_list
         self.endpoint = endpoint
-        self._channel = channel_factory(self.endpoint, driver_config, grpc.aio)
+        self._channel = channel_factory(
+            self.endpoint, driver_config, grpc.aio, endpoint_options=endpoint_options
+        )
         self._driver_config = driver_config
 
         self._stub_instances = {}

ydb/aio/pool.py

@@ -123,31 +123,40 @@ async def execute_discovery(self):
             return False
 
         resolved_endpoints = set(
-            details.endpoint for details in resolve_details.endpoints
+            endpoint
+            for resolved_endpoint in resolve_details.endpoints
+            for endpoint, endpoint_options in resolved_endpoint.endpoints_with_options()
         )
         for cached_endpoint in self._cache.values():
             if cached_endpoint.endpoint not in resolved_endpoints:
                 self._cache.make_outdated(cached_endpoint)
 
         for resolved_endpoint in resolve_details.endpoints:
-            if self._cache.size >= self._max_size or self._cache.already_exists(
-                resolved_endpoint.endpoint
-            ):
-                continue
-
             if self._ssl_required and not resolved_endpoint.ssl:
                 continue
 
             if not self._ssl_required and resolved_endpoint.ssl:
                 continue
 
-            endpoint = resolved_endpoint.endpoint
             preferred = resolve_details.self_location == resolved_endpoint.location
 
-            ready_connection = Connection(endpoint, self._driver_config)
-            await ready_connection.connection_ready(ready_timeout=self._ready_timeout)
+            for (
+                endpoint,
+                endpoint_options,
+            ) in resolved_endpoint.endpoints_with_options():
+                if self._cache.size >= self._max_size or self._cache.already_exists(
+                    endpoint
+                ):
+                    continue
+
+                ready_connection = Connection(
+                    endpoint, self._driver_config, endpoint_options=endpoint_options
+                )
+                await ready_connection.connection_ready(
+                    ready_timeout=self._ready_timeout
+                )
 
-            self._cache.add(ready_connection, preferred)
+                self._cache.add(ready_connection, preferred)
 
         await self._cache.cleanup_outdated()
         return self._cache.size > 0

ydb/connection.py

@@ -162,10 +162,18 @@ def _get_request_timeout(settings):
     return settings.timeout
 
 
-def _construct_channel_options(driver_config):
+class EndpointOptions(object):
+    __slots__ = ("ssl_target_name_override",)
+
+    def __init__(self, ssl_target_name_override=None):
+        self.ssl_target_name_override = ssl_target_name_override
+
+
+def _construct_channel_options(driver_config, endpoint_options=None):
     """
     Constructs gRPC channel initialization options
     :param driver_config: A driver config instance
+    :param endpoint_options: Endpoint options
     :return: A channel initialization options
     """
     _max_message_size = 64 * 10 ** 6
@@ -187,6 +195,14 @@ def _construct_channel_options(driver_config):
                 ("grpc.keepalive_permit_without_calls", 0),
             ]
         )
+    if endpoint_options is not None:
+        if endpoint_options.ssl_target_name_override:
+            _default_connect_options.append(
+                (
+                    "grpc.ssl_target_name_override",
+                    endpoint_options.ssl_target_name_override,
+                )
+            )
     if driver_config.channel_options is None:
         return _default_connect_options
     channel_options = copy.deepcopy(driver_config.channel_options)
@@ -276,9 +292,11 @@ def _set_server_timeouts(request, settings, default_value):
     _set_duration(request.operation_params.cancel_after, cancel_after)
 
 
-def channel_factory(endpoint, driver_config, channel_provider=None):
+def channel_factory(
+    endpoint, driver_config, channel_provider=None, endpoint_options=None
+):
     channel_provider = channel_provider if channel_provider is not None else grpc
-    options = _construct_channel_options(driver_config)
+    options = _construct_channel_options(driver_config, endpoint_options)
     logger.debug("Channel options: {}".format(options))
 
     if driver_config.root_certificates is None and not driver_config.secure_channel:
@@ -306,7 +324,7 @@ class Connection(object):
         "closing",
     )
 
-    def __init__(self, endpoint, driver_config=None):
+    def __init__(self, endpoint, driver_config=None, endpoint_options=None):
         """
         Object that wraps gRPC channel and encapsulates gRPC request execution logic
         :param endpoint: endpoint to connect (in pattern host:port), constructed by user or
@@ -315,7 +333,9 @@ def __init__(self, endpoint, driver_config=None):
         """
         global _stubs_list
         self.endpoint = endpoint
-        self._channel = channel_factory(self.endpoint, driver_config)
+        self._channel = channel_factory(
+            self.endpoint, driver_config, endpoint_options=endpoint_options
+        )
         self._driver_config = driver_config
         self._call_states = {}
         self._stub_instances = {}
@@ -436,8 +456,10 @@ def __call__(
             self._finish_call(rpc_state)
 
     @classmethod
-    def ready_factory(cls, endpoint, driver_config, ready_timeout=10):
-        candidate = cls(endpoint, driver_config)
+    def ready_factory(
+        cls, endpoint, driver_config, ready_timeout=10, endpoint_options=None
+    ):
+        candidate = cls(endpoint, driver_config, endpoint_options=endpoint_options)
         ready_future = candidate.ready_future()
         try:
             ready_future.result(timeout=ready_timeout)

ydb/pool.py

@@ -215,30 +215,40 @@ def execute_discovery(self):
                 return False
 
             resolved_endpoints = set(
-                details.endpoint for details in resolve_details.endpoints
+                endpoint
+                for resolved_endpoint in resolve_details.endpoints
+                for endpoint, endpoint_options in resolved_endpoint.endpoints_with_options()
             )
             for cached_endpoint in self._cache.values():
                 if cached_endpoint.endpoint not in resolved_endpoints:
                     self._cache.make_outdated(cached_endpoint)
 
             for resolved_endpoint in resolve_details.endpoints:
-                if self._cache.size >= self._max_size or self._cache.already_exists(
-                    resolved_endpoint.endpoint
-                ):
-                    continue
-
                 if self._ssl_required and not resolved_endpoint.ssl:
                     continue
 
                 if not self._ssl_required and resolved_endpoint.ssl:
                     continue
 
-                endpoint = resolved_endpoint.endpoint
                 preferred = resolve_details.self_location == resolved_endpoint.location
-                ready_connection = connection_impl.Connection.ready_factory(
-                    endpoint, self._driver_config, self._ready_timeout
-                )
-                self._cache.add(ready_connection, preferred)
+
+                for (
+                    endpoint,
+                    endpoint_options,
+                ) in resolved_endpoint.endpoints_with_options():
+                    if (
+                        self._cache.size >= self._max_size
+                        or self._cache.already_exists(endpoint)
+                    ):
+                        continue
+
+                    ready_connection = connection_impl.Connection.ready_factory(
+                        endpoint,
+                        self._driver_config,
+                        self._ready_timeout,
+                        endpoint_options=endpoint_options,
+                    )
+                    self._cache.add(ready_connection, preferred)
 
         self._cache.cleanup_outdated()
 

ydb/resolver.py

@@ -10,14 +10,46 @@
 
 
 class EndpointInfo(object):
-    __slots__ = ("address", "endpoint", "location", "port", "ssl")
+    __slots__ = (
+        "address",
+        "endpoint",
+        "location",
+        "port",
+        "ssl",
+        "ipv4_addrs",
+        "ipv6_addrs",
+        "ssl_target_name_override",
+    )
 
     def __init__(self, endpoint_info):
         self.address = endpoint_info.address
         self.endpoint = "%s:%s" % (endpoint_info.address, endpoint_info.port)
         self.location = endpoint_info.location
         self.port = endpoint_info.port
         self.ssl = endpoint_info.ssl
+        self.ipv4_addrs = tuple(endpoint_info.ip_v4)
+        self.ipv6_addrs = tuple(endpoint_info.ip_v6)
+        self.ssl_target_name_override = endpoint_info.ssl_target_name_override
+
+    def endpoints_with_options(self):
+        if self.ssl:
+            if self.ssl_target_name_override:
+                endpoint_options = conn_impl.EndpointOptions(
+                    self.ssl_target_name_override
+                )
+            elif self.ipv6_addrs or self.ipv4_addrs:
+                endpoint_options = conn_impl.EndpointOptions(self.address)
+            else:
+                endpoint_options = None
+        else:
+            endpoint_options = None
+        if self.ipv6_addrs or self.ipv4_addrs:
+            for ipv6addr in self.ipv6_addrs:
+                yield ("ipv6:[%s]:%s" % (ipv6addr, self.port), endpoint_options)
+            for ipv4addr in self.ipv4_addrs:
+                yield ("ipv4:%s:%s" % (ipv4addr, self.port), endpoint_options)
+        else:
+            yield (self.endpoint, endpoint_options)
 
     def __str__(self):
         return "<Endpoint %s, location %s, ssl: %s>" % (