Extract client.check_connection()

Author: fantix

gel/__init__.py

@@ -30,6 +30,7 @@
 from .abstract import (
     Executor, AsyncIOExecutor, ReadOnlyExecutor, AsyncIOReadOnlyExecutor,
 )
+from .base_client import ConnectionInfo
 
 from .asyncio_client import (
     create_async_client,
@@ -52,6 +53,7 @@
     "Cardinality",
     "Client",
     "ConfigMemory",
+    "ConnectionInfo",
     "DateDuration",
     "EdgeDBError",
     "EdgeDBMessage",

gel/ai/core.py

@@ -27,15 +27,15 @@
 
 
 def create_rag_client(client: gel.Client, **kwargs) -> RAGClient:
-    client.ensure_connected()
-    return RAGClient(client, types.RAGOptions(**kwargs))
+    info = client.check_connection()
+    return RAGClient(info, types.RAGOptions(**kwargs))
 
 
 async def create_async_rag_client(
     client: gel.AsyncIOClient, **kwargs
 ) -> AsyncRAGClient:
-    await client.ensure_connected()
-    return AsyncRAGClient(client, types.RAGOptions(**kwargs))
+    info = await client.check_connection()
+    return AsyncRAGClient(info, types.RAGOptions(**kwargs))
 
 
 class BaseRAGClient:
@@ -45,25 +45,26 @@ class BaseRAGClient:
 
     def __init__(
         self,
-        client: typing.Union[gel.Client, gel.AsyncIOClient],
+        info: gel.ConnectionInfo,
         options: types.RAGOptions,
         **kwargs,
     ):
-        pool = client._impl
-        host, port = pool._working_addr
-        params = pool._working_params
-        proto = "http" if params.tls_security == "insecure" else "https"
-        branch = params.branch
+        proto = "http" if info.params.tls_security == "insecure" else "https"
+        branch = info.params.branch
         self.options = options
         self.context = types.QueryContext(**kwargs)
         args = dict(
-            base_url=f"{proto}://{host}:{port}/branch/{branch}/ext/ai",
-            verify=params.ssl_ctx,
+            base_url=(
+                f"{proto}://{info.host}:{info.port}/branch/{branch}/ext/ai"
+            ),
+            verify=info.params.ssl_ctx,
         )
-        if params.password is not None:
-            args["auth"] = (params.user, params.password)
-        elif params.secret_key is not None:
-            args["headers"] = {"Authorization": f"Bearer {params.secret_key}"}
+        if info.params.password is not None:
+            args["auth"] = (info.params.user, info.params.password)
+        elif info.params.secret_key is not None:
+            args["headers"] = {
+                "Authorization": f"Bearer {info.params.secret_key}"
+            }
         self._init_client(**args)
 
     def _init_client(self, **kwargs):

gel/asyncio_client.py

@@ -361,8 +361,11 @@ class AsyncIOClient(base_client.BaseClient, abstract.AsyncIOExecutor):
     __slots__ = ()
     _impl_class = _AsyncIOPoolImpl
 
+    async def check_connection(self) -> base_client.ConnectionInfo:
+        return await self._impl.ensure_connected()
+
     async def ensure_connected(self):
-        await self._impl.ensure_connected()
+        await self.check_connection()
         return self
 
     async def aclose(self):

gel/auth/email_password.py

@@ -170,13 +170,10 @@ async def make(
     verify_url: str,
     reset_url: str,
 ) -> EmailPassword:
-    await client.ensure_connected()
-    pool = client._impl
-    host, port = pool._working_addr
-    params = pool._working_params
-    proto = "http" if params.tls_security == "insecure" else "https"
-    branch = params.branch
-    auth_ext_url = f"{proto}://{host}:{port}/branch/{branch}/ext/auth/"
+    info = await client.check_connection()
+    proto = "http" if info.params.tls_security == "insecure" else "https"
+    branch = info.params.branch
+    auth_ext_url = f"{proto}://{info.host}:{info.port}/branch/{branch}/ext/auth/"
     return EmailPassword(
         auth_ext_url=auth_ext_url, verify_url=verify_url, reset_url=reset_url
     )

gel/base_client.py

@@ -18,6 +18,7 @@
 
 
 import abc
+import dataclasses
 import random
 import time
 import typing
@@ -411,6 +412,14 @@ def _release(self):
         self._pool._queue.put_nowait(self)
 
 
+@dataclasses.dataclass
+class ConnectionInfo:
+    host: str
+    port: int
+    params: con_utils.ResolvedConnectConfig
+    config: con_utils.ClientConfiguration
+
+
 class BasePoolImpl(abc.ABC):
     __slots__ = (
         "_connect_args",
@@ -624,7 +633,7 @@ def expire_connections(self):
         """
         self._generation += 1
 
-    async def ensure_connected(self):
+    async def ensure_connected(self) -> ConnectionInfo:
         self._ensure_initialized()
 
         for ch in self._holders:
@@ -635,6 +644,16 @@ async def ensure_connected(self):
         ch._con = None
         await ch.connect()
 
+        assert self._working_addr is not None
+        assert self._working_config is not None
+        assert self._working_params is not None
+        return ConnectionInfo(
+            host=self._working_addr[0],
+            port=self._working_addr[1],
+            config=self._working_config,
+            params=self._working_params,
+        )
+
 
 class BaseClient(abstract.BaseReadOnlyExecutor, _options._OptionsMixin):
     __slots__ = ("_impl", "_options")

gel/blocking_client.py

@@ -39,6 +39,18 @@
 MINIMUM_PING_WAIT_TIME = datetime.timedelta(seconds=1)
 
 
+T = typing.TypeVar("T")
+
+
+def iter_coroutine(coro: typing.Coroutine[None, None, T]) -> T:
+    try:
+        coro.send(None)
+    except StopIteration as ex:
+        return ex.value
+    finally:
+        coro.close()
+
+
 class BlockingIOConnection(base_client.BaseConnection):
     __slots__ = ("_ping_wait_time",)
 
@@ -328,7 +340,7 @@ def __enter__(self):
     def __exit__(self, extype, ex, tb):
         with self._exclusive():
             self._managed = False
-            return self._client._iter_coroutine(self._exit(extype, ex))
+            return iter_coroutine(self._exit(extype, ex))
 
     async def _ensure_transaction(self):
         if not self._managed:
@@ -340,11 +352,11 @@ async def _ensure_transaction(self):
 
     def _query(self, query_context: abstract.QueryContext):
         with self._exclusive():
-            return self._client._iter_coroutine(super()._query(query_context))
+            return iter_coroutine(super()._query(query_context))
 
     def _execute(self, execute_context: abstract.ExecuteContext) -> None:
         with self._exclusive():
-            self._client._iter_coroutine(super()._execute(execute_context))
+            iter_coroutine(super()._execute(execute_context))
 
     @contextlib.contextmanager
     def _exclusive(self):
@@ -392,22 +404,17 @@ class Client(base_client.BaseClient, abstract.Executor):
     __slots__ = ()
     _impl_class = _PoolImpl
 
-    def _iter_coroutine(self, coro):
-        try:
-            coro.send(None)
-        except StopIteration as ex:
-            return ex.value
-        finally:
-            coro.close()
-
     def _query(self, query_context: abstract.QueryContext):
-        return self._iter_coroutine(super()._query(query_context))
+        return iter_coroutine(super()._query(query_context))
 
     def _execute(self, execute_context: abstract.ExecuteContext) -> None:
-        self._iter_coroutine(super()._execute(execute_context))
+        iter_coroutine(super()._execute(execute_context))
+
+    def check_connection(self) -> base_client.ConnectionInfo:
+        return iter_coroutine(self._impl.ensure_connected())
 
     def ensure_connected(self):
-        self._iter_coroutine(self._impl.ensure_connected())
+        self.check_connection()
         return self
 
     def transaction(self) -> Retry:
@@ -421,7 +428,7 @@ def close(self, timeout=None):
         in ``close()`` the pool will terminate by calling
         Client.terminate() .
         """
-        self._iter_coroutine(self._impl.close(timeout))
+        iter_coroutine(self._impl.close(timeout))
 
     def __enter__(self):
         return self.ensure_connected()
@@ -438,7 +445,7 @@ def _describe_query(
         output_format: OutputFormat = OutputFormat.BINARY,
         expect_one: bool = False,
     ) -> abstract.DescribeResult:
-        return self._iter_coroutine(self._describe(abstract.DescribeContext(
+        return iter_coroutine(self._describe(abstract.DescribeContext(
             query=query,
             state=self._get_state(),
             inject_type_names=inject_type_names,

tests/test_sync_query.py

@@ -28,6 +28,7 @@
 
 from edgedb import abstract
 from gel import _testbase as tb
+from gel import blocking_client
 from edgedb.protocol import protocol
 
 
@@ -778,7 +779,7 @@ def test_json(self):
 
     def test_json_elements(self):
         self.client.ensure_connected()
-        result = self.client._iter_coroutine(
+        result = blocking_client.iter_coroutine(
             self.client.connection.raw_query(
                 abstract.QueryContext(
                     query=abstract.QueryWithArgs(