Address PR comments

Author: rhatgadkar-goog

google/cloud/alloydb/connector/async_connector.py

@@ -15,10 +15,9 @@
 from __future__ import annotations
 
 import asyncio
-import io
 import logging
 from types import TracebackType
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, Optional, TYPE_CHECKING
 
 import google.auth
 from google.auth.credentials import with_scopes_if_required
@@ -30,7 +29,7 @@
 from google.cloud.alloydb.connector.enums import RefreshStrategy
 from google.cloud.alloydb.connector.instance import RefreshAheadCache
 from google.cloud.alloydb.connector.lazy import LazyRefreshCache
-from google.cloud.alloydb.connector.static import StaticConnectionInfoCache
+from google.cloud.alloydb.connector.types import CacheTypes
 from google.cloud.alloydb.connector.utils import generate_keys
 
 if TYPE_CHECKING:
@@ -61,9 +60,6 @@ class AsyncConnector:
             of the following: RefreshStrategy.LAZY ("LAZY") or
             RefreshStrategy.BACKGROUND ("BACKGROUND").
             Default: RefreshStrategy.BACKGROUND
-        static_conn_info (io.TextIOBase): A file-like JSON object that contains
-            static connection info for the StaticConnectionInfoCache.
-            Defaults to None, which will not use the StaticConnectionInfoCache.
     """
 
     def __init__(
@@ -75,9 +71,8 @@ def __init__(
         ip_type: str | IPTypes = IPTypes.PRIVATE,
         user_agent: Optional[str] = None,
         refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
-        static_conn_info: Optional[io.TextIOBase] = None,
     ) -> None:
-        self._cache: dict[str, Union[RefreshAheadCache, LazyRefreshCache]] = {}
+        self._cache: dict[str, CacheTypes] = {}
         # initialize default params
         self._quota_project = quota_project
         self._alloydb_api_endpoint = alloydb_api_endpoint
@@ -106,7 +101,6 @@ def __init__(
         except RuntimeError:
             self._keys = None
         self._client: Optional[AlloyDBClient] = None
-        self._static_conn_info = static_conn_info
 
     async def connect(
         self,
@@ -147,11 +141,8 @@ async def connect(
         enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)
 
         # use existing connection info if possible
-        cache: Union[RefreshAheadCache, LazyRefreshCache, StaticConnectionInfoCache]
         if instance_uri in self._cache:
             cache = self._cache[instance_uri]
-        elif self._static_conn_info:
-            cache = StaticConnectionInfoCache(instance_uri, self._static_conn_info)
         else:
             if self._refresh_strategy == RefreshStrategy.LAZY:
                 logger.debug(

google/cloud/alloydb/connector/connector.py

@@ -22,7 +22,7 @@
 import struct
 from threading import Thread
 from types import TracebackType
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, Optional, TYPE_CHECKING
 
 from google.auth import default
 from google.auth.credentials import TokenState
@@ -35,7 +35,7 @@
 from google.cloud.alloydb.connector.instance import RefreshAheadCache
 from google.cloud.alloydb.connector.lazy import LazyRefreshCache
 import google.cloud.alloydb.connector.pg8000 as pg8000
-from google.cloud.alloydb.connector.static import StaticConnectionInfoCache
+from google.cloud.alloydb.connector.types import CacheTypes
 from google.cloud.alloydb.connector.utils import generate_keys
 import google.cloud.alloydb_connectors_v1.proto.resources_pb2 as connectorspb
 
@@ -93,7 +93,7 @@ def __init__(
         self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
         self._thread = Thread(target=self._loop.run_forever, daemon=True)
         self._thread.start()
-        self._cache: dict[str, Union[RefreshAheadCache, LazyRefreshCache]] = {}
+        self._cache: dict[str, CacheTypes] = {}
         # initialize default params
         self._quota_project = quota_project
         self._alloydb_api_endpoint = alloydb_api_endpoint
@@ -176,11 +176,8 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
             )
         enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)
         # use existing connection info if possible
-        cache: Union[RefreshAheadCache, LazyRefreshCache, StaticConnectionInfoCache]
         if instance_uri in self._cache:
             cache = self._cache[instance_uri]
-        elif self._static_conn_info:
-            cache = StaticConnectionInfoCache(instance_uri, self._static_conn_info)
         else:
             if self._refresh_strategy == RefreshStrategy.LAZY:
                 logger.debug(

google/cloud/alloydb/connector/types.py

@@ -0,0 +1,23 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+
+from google.cloud.alloydb.connector.instance import RefreshAheadCache
+from google.cloud.alloydb.connector.lazy import LazyRefreshCache
+from google.cloud.alloydb.connector.static import StaticConnectionInfoCache
+
+CacheTypes = typing.Union[
+    RefreshAheadCache, LazyRefreshCache, StaticConnectionInfoCache
+]

tests/unit/conftest.py

@@ -16,6 +16,7 @@
 import socket
 import ssl
 from threading import Thread
+from typing import Final
 
 from aiofiles.tempfile import TemporaryDirectory
 from mocks import FakeAlloyDBClient
@@ -26,6 +27,8 @@
 
 from google.cloud.alloydb.connector.utils import _write_to_file
 
+DELAY: Final[float] = 1.0
+
 
 @pytest.fixture
 def credentials() -> FakeCredentials:
@@ -86,4 +89,4 @@ def proxy_server(fake_instance: FakeInstance) -> None:
         daemon=True,
     )
     thread.start()
-    thread.join(0.1)  # wait 100ms to allow the proxy server to start
+    thread.join(DELAY)  # add a delay to allow the proxy server to start

tests/unit/test_async_connector.py

@@ -20,7 +20,6 @@
 from mocks import FakeAlloyDBClient
 from mocks import FakeConnectionInfo
 from mocks import FakeCredentials
-from mocks import write_static_info
 import pytest
 
 from google.cloud.alloydb.connector import AsyncConnector
@@ -334,29 +333,3 @@ async def test_Connector_remove_cached_no_ip_type(credentials: FakeCredentials)
             await connector.connect(instance_uri, "asyncpg", ip_type="private")
         # check that cache has been removed from dict
         assert instance_uri not in connector._cache
-
-
-async def test_Connector_static_connection_info(
-    credentials: FakeCredentials, fake_client: FakeAlloyDBClient
-) -> None:
-    """
-    Test that AsyncConnector.__init__() can specify a static connection info to
-    connect to an instance.
-    """
-    static_info = write_static_info(fake_client.instance)
-    async with AsyncConnector(
-        credentials=credentials, static_conn_info=static_info
-    ) as connector:
-        connector._client = fake_client
-        # patch db connection creation
-        with patch("google.cloud.alloydb.connector.asyncpg.connect") as mock_connect:
-            mock_connect.return_value = True
-            connection = await connector.connect(
-                fake_client.instance.uri(),
-                "asyncpg",
-                user="test-user",
-                password="test-password",
-                db="test-db",
-            )
-        # check connection is returned
-        assert connection is True