fix: throw exception when connecting to closed Connector (#436)

Author: rhatgadkar-goog

google/cloud/alloydb/connector/async_connector.py

@@ -27,6 +27,7 @@
 from google.cloud.alloydb.connector.client import AlloyDBClient
 from google.cloud.alloydb.connector.enums import IPTypes
 from google.cloud.alloydb.connector.enums import RefreshStrategy
+from google.cloud.alloydb.connector.exceptions import ClosedConnectorError
 from google.cloud.alloydb.connector.instance import RefreshAheadCache
 from google.cloud.alloydb.connector.lazy import LazyRefreshCache
 from google.cloud.alloydb.connector.types import CacheTypes
@@ -102,6 +103,7 @@ def __init__(
         except RuntimeError:
             self._keys = None
         self._client: Optional[AlloyDBClient] = None
+        self._closed = False
 
     async def connect(
         self,
@@ -127,6 +129,10 @@ async def connect(
         Returns:
             connection: A DBAPI connection to the specified AlloyDB instance.
         """
+        if self._closed:
+            raise ClosedConnectorError(
+                "Connection attempt failed because the connector has already been closed."
+            )
         if self._keys is None:
             self._keys = asyncio.create_task(generate_keys())
         if self._client is None:
@@ -236,3 +242,4 @@ async def close(self) -> None:
         """Helper function to cancel RefreshAheadCaches' tasks
         and close client."""
         await asyncio.gather(*[cache.close() for cache in self._cache.values()])
+        self._closed = True

google/cloud/alloydb/connector/connector.py

@@ -32,6 +32,7 @@
 from google.cloud.alloydb.connector.client import AlloyDBClient
 from google.cloud.alloydb.connector.enums import IPTypes
 from google.cloud.alloydb.connector.enums import RefreshStrategy
+from google.cloud.alloydb.connector.exceptions import ClosedConnectorError
 from google.cloud.alloydb.connector.instance import RefreshAheadCache
 from google.cloud.alloydb.connector.lazy import LazyRefreshCache
 import google.cloud.alloydb.connector.pg8000 as pg8000
@@ -124,6 +125,7 @@ def __init__(
         )
         self._client: Optional[AlloyDBClient] = None
         self._static_conn_info = static_conn_info
+        self._closed = False
 
     def connect(self, instance_uri: str, driver: str, **kwargs: Any) -> Any:
         """
@@ -144,6 +146,10 @@ def connect(self, instance_uri: str, driver: str, **kwargs: Any) -> Any:
         Returns:
             connection: A DBAPI connection to the specified AlloyDB instance.
         """
+        if self._closed:
+            raise ClosedConnectorError(
+                "Connection attempt failed because the connector has already been closed."
+            )
         # call async connect and wait on result
         connect_task = asyncio.run_coroutine_threadsafe(
             self.connect_async(instance_uri, driver, **kwargs), self._loop
@@ -385,6 +391,7 @@ def close(self) -> None:
                 self._loop.call_soon_threadsafe(self._loop.stop)
             # wait for thread to finish closing (i.e. loop to stop)
             self._thread.join()
+        self._closed = True
 
     async def close_async(self) -> None:
         """Helper function to cancel RefreshAheadCaches' tasks

google/cloud/alloydb/connector/exceptions.py

@@ -19,3 +19,7 @@ class RefreshError(Exception):
 
 class IPTypeNotFoundError(Exception):
     pass
+
+
+class ClosedConnectorError(Exception):
+    pass

tests/unit/test_async_connector.py

@@ -24,6 +24,7 @@
 
 from google.cloud.alloydb.connector import AsyncConnector
 from google.cloud.alloydb.connector import IPTypes
+from google.cloud.alloydb.connector.exceptions import ClosedConnectorError
 from google.cloud.alloydb.connector.exceptions import IPTypeNotFoundError
 from google.cloud.alloydb.connector.instance import RefreshAheadCache
 
@@ -42,6 +43,7 @@ async def test_AsyncConnector_init(credentials: FakeCredentials) -> None:
     assert connector._client is None
     assert connector._credentials == credentials
     assert connector._enable_iam_auth is False
+    assert connector._closed is False
     await connector.close()
 
 
@@ -109,7 +111,7 @@ async def test_AsyncConnector_init_bad_ip_type(credentials: FakeCredentials) ->
     )
 
 
-def test_AsyncConnector_init_alloydb_api_endpoint_with_http_prefix(
+async def test_AsyncConnector_init_alloydb_api_endpoint_with_http_prefix(
     credentials: FakeCredentials,
 ) -> None:
     """
@@ -120,10 +122,10 @@ def test_AsyncConnector_init_alloydb_api_endpoint_with_http_prefix(
         alloydb_api_endpoint="http://alloydb.googleapis.com", credentials=credentials
     )
     assert connector._alloydb_api_endpoint == "alloydb.googleapis.com"
-    connector.close()
+    await connector.close()
 
 
-def test_AsyncConnector_init_alloydb_api_endpoint_with_https_prefix(
+async def test_AsyncConnector_init_alloydb_api_endpoint_with_https_prefix(
     credentials: FakeCredentials,
 ) -> None:
     """
@@ -134,7 +136,7 @@ def test_AsyncConnector_init_alloydb_api_endpoint_with_https_prefix(
         alloydb_api_endpoint="https://alloydb.googleapis.com", credentials=credentials
     )
     assert connector._alloydb_api_endpoint == "alloydb.googleapis.com"
-    connector.close()
+    await connector.close()
 
 
 @pytest.mark.asyncio
@@ -357,3 +359,26 @@ 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_close_sets_connector_as_closed(credentials: FakeCredentials) -> None:
+    """
+    Test that when connector is closed, it marked as closed.
+    """
+    async with AsyncConnector(credentials=credentials) as connector:
+        assert connector._closed is False
+    assert connector._closed is True
+
+
+async def test_connect_when_closed(credentials: FakeCredentials) -> None:
+    """
+    Test that connector.connect errors when the connection is closed.
+    """
+    connector = AsyncConnector(credentials=credentials)
+    await connector.close()
+    with pytest.raises(ClosedConnectorError) as exc_info:
+        await connector.connect("", "")
+    assert (
+        exc_info.value.args[0]
+        == "Connection attempt failed because the connector has already been closed."
+    )

tests/unit/test_connector.py

@@ -25,6 +25,7 @@
 
 from google.cloud.alloydb.connector import Connector
 from google.cloud.alloydb.connector import IPTypes
+from google.cloud.alloydb.connector.exceptions import ClosedConnectorError
 from google.cloud.alloydb.connector.exceptions import IPTypeNotFoundError
 from google.cloud.alloydb.connector.instance import RefreshAheadCache
 from google.cloud.alloydb.connector.utils import generate_keys
@@ -40,6 +41,7 @@ def test_Connector_init(credentials: FakeCredentials) -> None:
     assert connector._alloydb_api_endpoint == "alloydb.googleapis.com"
     assert connector._client is None
     assert connector._credentials == credentials
+    assert connector._closed is False
     connector.close()
 
 
@@ -150,15 +152,17 @@ def test_Connector_context_manager(credentials: FakeCredentials) -> None:
 def test_Connector_close(credentials: FakeCredentials) -> None:
     """
     Test that Connector's close method stops event loop and
-    background thread.
+    background thread, and sets the connector as closed.
     """
     with Connector(credentials) as connector:
         loop: asyncio.AbstractEventLoop = connector._loop
         thread: Thread = connector._thread
         assert loop.is_running() is True
         assert thread.is_alive() is True
+        assert connector._closed is False
     assert loop.is_running() is False
     assert thread.is_alive() is False
+    assert connector._closed is True
 
 
 @pytest.mark.usefixtures("proxy_server")
@@ -302,3 +306,17 @@ def test_Connector_static_connection_info(
             )
         # check connection is returned
         assert connection is True
+
+
+def test_connect_when_closed(credentials: FakeCredentials) -> None:
+    """
+    Test that connector.connect errors when the connection is closed.
+    """
+    connector = Connector(credentials=credentials)
+    connector.close()
+    with pytest.raises(ClosedConnectorError) as exc_info:
+        connector.connect("", "")
+    assert (
+        exc_info.value.args[0]
+        == "Connection attempt failed because the connector has already been closed."
+    )