fix: Modify AlloyDBClient to use sync transport for sync connector (#442)

rhatgadkar-goog · web-flow · commit 984dd800bbaa · 2025-04-25T09:58:25.000-07:00
diff --git a/google/cloud/alloydb/connector/client.py b/google/cloud/alloydb/connector/client.py
@@ -79,14 +79,28 @@ def __init__(
                 requests to AlloyDB APIs.
                 Optional, defaults to None and creates new client.
             driver (str): Database driver to be used by the client.
+            user_agent (str): The custom user-agent string to use in the HTTP
+                header when making requests to AlloyDB APIs.
+                Optional, defaults to None and uses a pre-defined one.
         """
         user_agent = _format_user_agent(driver, user_agent)
 
-        self._client = (
-            client
-            if client
-            else v1beta.AlloyDBAdminAsyncClient(
+        # TODO(rhatgadkar-goog): Rollback the PR of deciding between creating
+        # AlloyDBAdminClient or AlloyDBAdminAsyncClient when either
+        # https://github.com/grpc/grpc/issues/25364 is resolved or an async REST
+        # transport for AlloyDBAdminAsyncClient gets introduced.
+        # The issue is that the async gRPC transport does not work with multiple
+        # event loops in the same process. So all calls to the AlloyDB Admin
+        # API, even from multiple threads, need to be made to a single-event
+        # loop. See https://github.com/GoogleCloudPlatform/alloydb-python-connector/issues/435
+        # for more details.
+        self._is_sync = False
+        if client:
+            self._client = client
+        elif driver == "pg8000":
+            self._client = v1beta.AlloyDBAdminClient(
                 credentials=credentials,
+                transport="grpc",
                 client_options=ClientOptions(
                     api_endpoint=alloydb_api_endpoint,
                     quota_project_id=quota_project,
@@ -95,7 +109,20 @@ def __init__(
                     user_agent=user_agent,
                 ),
             )
-        )
+            self._is_sync = True
+        else:
+            self._client = v1beta.AlloyDBAdminAsyncClient(
+                credentials=credentials,
+                transport="grpc_asyncio",
+                client_options=ClientOptions(
+                    api_endpoint=alloydb_api_endpoint,
+                    quota_project_id=quota_project,
+                ),
+                client_info=ClientInfo(
+                    user_agent=user_agent,
+                ),
+            )
+
         self._credentials = credentials
         # asyncpg does not currently support using metadata exchange
         # only use metadata exchange for pg8000 driver
@@ -131,7 +158,10 @@ async def _get_metadata(
         )
 
         req = v1beta.GetConnectionInfoRequest(parent=parent)
-        resp = await self._client.get_connection_info(request=req)
+        if self._is_sync:
+            resp = self._client.get_connection_info(request=req)
+        else:
+            resp = await self._client.get_connection_info(request=req)
 
         # Remove trailing period from PSC DNS name.
         psc_dns = resp.psc_dns_name
@@ -178,7 +208,10 @@ async def _get_client_certificate(
             public_key=pub_key,
             use_metadata_exchange=self._use_metadata,
         )
-        resp = await self._client.generate_client_certificate(request=req)
+        if self._is_sync:
+            resp = self._client.generate_client_certificate(request=req)
+        else:
+            resp = await self._client.generate_client_certificate(request=req)
         return (resp.ca_cert, resp.pem_certificate_chain)
 
     async def get_connection_info(
diff --git a/tests/unit/mocks.py b/tests/unit/mocks.py
@@ -478,3 +478,36 @@ async def generate_client_certificate(
         ccr.pem_certificate_chain.append("This is the intermediate cert")
         ccr.pem_certificate_chain.append("This is the root cert")
         return ccr
+
+
+class FakeAlloyDBAdminClient:
+    def get_connection_info(
+        self, request: alloydb_v1beta.GetConnectionInfoRequest
+    ) -> alloydb_v1beta.types.resources.ConnectionInfo:
+        ci = alloydb_v1beta.types.resources.ConnectionInfo()
+        ci.ip_address = "10.0.0.1"
+        ci.public_ip_address = "127.0.0.1"
+        ci.instance_uid = "123456789"
+        ci.psc_dns_name = "x.y.alloydb.goog"
+
+        parent = request.parent
+        instance = parent.split("/")[-1]
+        if instance == "test-instance":
+            ci.public_ip_address = ""
+            ci.psc_dns_name = ""
+        elif instance == "public-instance":
+            ci.psc_dns_name = ""
+        else:
+            ci.ip_address = ""
+            ci.public_ip_address = ""
+        return ci
+
+    def generate_client_certificate(
+        self, request: alloydb_v1beta.GenerateClientCertificateRequest
+    ) -> alloydb_v1beta.types.service.GenerateClientCertificateResponse:
+        ccr = alloydb_v1beta.types.service.GenerateClientCertificateResponse()
+        ccr.ca_cert = "This is the CA cert"
+        ccr.pem_certificate_chain.append("This is the client cert")
+        ccr.pem_certificate_chain.append("This is the intermediate cert")
+        ccr.pem_certificate_chain.append("This is the root cert")
+        return ccr
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -14,7 +14,9 @@
 
 from typing import Optional
 
+import google.cloud.alloydb_v1beta as v1beta
 from mocks import FakeAlloyDBAdminAsyncClient
+from mocks import FakeAlloyDBAdminClient
 from mocks import FakeCredentials
 import pytest
 
@@ -80,6 +82,40 @@ async def test__get_metadata_with_psc(credentials: FakeCredentials) -> None:
     }
 
 
+async def test__get_metadata_with_async_client(credentials: FakeCredentials) -> None:
+    """
+    Test _get_metadata returns successfully for an async client.
+    """
+    test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminAsyncClient())
+    test_client._is_sync = False
+    assert (
+        await test_client._get_metadata(
+            "test-project",
+            "test-region",
+            "test-cluster",
+            "psc-instance",
+        )
+        is not None
+    )
+
+
+async def test__get_metadata_with_sync_client(credentials: FakeCredentials) -> None:
+    """
+    Test _get_metadata returns successfully for a sync client.
+    """
+    test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminClient())
+    test_client._is_sync = True
+    assert (
+        await test_client._get_metadata(
+            "test-project",
+            "test-region",
+            "test-cluster",
+            "psc-instance",
+        )
+        is not None
+    )
+
+
 @pytest.mark.asyncio
 async def test__get_client_certificate(credentials: FakeCredentials) -> None:
     """
@@ -97,6 +133,40 @@ async def test__get_client_certificate(credentials: FakeCredentials) -> None:
     assert cert_chain[2] == "This is the root cert"
 
 
+async def test__get_client_certificate_with_async_client(
+    credentials: FakeCredentials,
+) -> None:
+    """
+    Test _get_client_certificate returns successfully for an async client.
+    """
+    test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminAsyncClient())
+    test_client._is_sync = False
+    keys = await generate_keys()
+    assert (
+        await test_client._get_client_certificate(
+            "test-project", "test-region", "test-cluster", keys[1]
+        )
+        is not None
+    )
+
+
+async def test__get_client_certificate_with_sync_client(
+    credentials: FakeCredentials,
+) -> None:
+    """
+    Test _get_client_certificate returns successfully for a sync client.
+    """
+    test_client = AlloyDBClient("", "", credentials, FakeAlloyDBAdminClient())
+    test_client._is_sync = True
+    keys = await generate_keys()
+    assert (
+        await test_client._get_client_certificate(
+            "test-project", "test-region", "test-cluster", keys[1]
+        )
+        is not None
+    )
+
+
 @pytest.mark.asyncio
 async def test_AlloyDBClient_init_(credentials: FakeCredentials) -> None:
     """
@@ -129,6 +199,48 @@ async def test_AlloyDBClient_init_custom_user_agent(
     )
 
 
+async def test_AlloyDBClient_init_specified_client(
+    credentials: FakeCredentials,
+) -> None:
+    """
+    Test to check that __init__ method of AlloyDBClient uses specified client.
+    """
+    client = AlloyDBClient(
+        "www.test-endpoint.com",
+        "my-quota-project",
+        credentials,
+        FakeAlloyDBAdminAsyncClient(),
+    )
+    assert client._is_sync is False
+    assert type(client._client) is FakeAlloyDBAdminAsyncClient
+
+
+async def test_AlloyDBClient_init_sync_client(credentials: FakeCredentials) -> None:
+    """
+    Test to check that __init__ method of AlloyDBClient creates a sync client
+    when client is not specified and driver is pg8000.
+    """
+    client = AlloyDBClient(
+        "www.test-endpoint.com", "my-quota-project", credentials, driver="pg8000"
+    )
+    assert client._is_sync is True
+    assert type(client._client) is v1beta.AlloyDBAdminClient
+    assert client._client.transport.kind == "grpc"
+
+
+async def test_AlloyDBClient_init_async_client(credentials: FakeCredentials) -> None:
+    """
+    Test to check that __init__ method of AlloyDBClient creates an async client
+    when client is not specified and driver is not pg8000.
+    """
+    client = AlloyDBClient(
+        "www.test-endpoint.com", "my-quota-project", credentials, driver=""
+    )
+    assert client._is_sync is False
+    assert type(client._client) is v1beta.AlloyDBAdminAsyncClient
+    assert client._client.transport.kind == "grpc_asyncio"
+
+
 @pytest.mark.parametrize(
     "driver",
     [None, "pg8000", "asyncpg"],
