feat: invalidate cache on bad connection info and failed IP lookup (#389

)

This change is analogous to CloudSQL's GoogleCloudPlatform/cloud-sql-python-connector#1118.

The Connector caches connection info for future connections and schedules refresh operations. For unrecoverable errors, the cache should be invalidated to stop future bad refreshes. The cache should also be invalidated on failed IP lookups. This PR does these things.

google/cloud/alloydb/connector/async_connector.py

@@ -177,8 +177,14 @@ async def connect(
         # if ip_type is str, convert to IPTypes enum
         if isinstance(ip_type, str):
             ip_type = IPTypes(ip_type.upper())
-        conn_info = await cache.connect_info()
-        ip_address = conn_info.get_preferred_ip(ip_type)
+        try:
+            conn_info = await cache.connect_info()
+            ip_address = conn_info.get_preferred_ip(ip_type)
+        except Exception:
+            # with an error from AlloyDB API call or IP type, invalidate the
+            # cache and re-raise the error
+            await self._remove_cached(instance_uri)
+            raise
         logger.debug(f"['{instance_uri}']: Connecting to {ip_address}:5433")
 
         # callable to be used for auto IAM authn
@@ -202,6 +208,15 @@ def get_authentication_token() -> str:
             await cache.force_refresh()
             raise
 
+    async def _remove_cached(self, instance_uri: str) -> None:
+        """Stops all background refreshes and deletes the connection
+        info cache from the map of caches.
+        """
+        logger.debug(f"['{instance_uri}']: Removing connection info from cache")
+        # remove cache from stored caches and close it
+        cache = self._cache.pop(instance_uri)
+        await cache.close()
+
     async def __aenter__(self) -> Any:
         """Enter async context manager by returning Connector object"""
         return self

google/cloud/alloydb/connector/connector.py

@@ -206,8 +206,14 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
         # if ip_type is str, convert to IPTypes enum
         if isinstance(ip_type, str):
             ip_type = IPTypes(ip_type.upper())
-        conn_info = await cache.connect_info()
-        ip_address = conn_info.get_preferred_ip(ip_type)
+        try:
+            conn_info = await cache.connect_info()
+            ip_address = conn_info.get_preferred_ip(ip_type)
+        except Exception:
+            # with an error from AlloyDB API call or IP type, invalidate the
+            # cache and re-raise the error
+            await self._remove_cached(instance_uri)
+            raise
         logger.debug(f"['{instance_uri}']: Connecting to {ip_address}:5433")
 
         # synchronous drivers are blocking and run using executor
@@ -334,6 +340,15 @@ def metadata_exchange(
 
         return sock
 
+    async def _remove_cached(self, instance_uri: str) -> None:
+        """Stops all background refreshes and deletes the connection
+        info cache from the map of caches.
+        """
+        logger.debug(f"['{instance_uri}']: Removing connection info from cache")
+        # remove cache from stored caches and close it
+        cache = self._cache.pop(instance_uri)
+        await cache.close()
+
     def __enter__(self) -> "Connector":
         """Enter context manager by returning Connector object"""
         return self

tests/unit/test_async_connector.py

@@ -15,6 +15,7 @@
 import asyncio
 from typing import Union
 
+from aiohttp import ClientResponseError
 from mock import patch
 from mocks import FakeAlloyDBClient
 from mocks import FakeConnectionInfo
@@ -23,6 +24,8 @@
 
 from google.cloud.alloydb.connector import AsyncConnector
 from google.cloud.alloydb.connector import IPTypes
+from google.cloud.alloydb.connector.exceptions import IPTypeNotFoundError
+from google.cloud.alloydb.connector.instance import RefreshAheadCache
 
 ALLOYDB_API_ENDPOINT = "https://alloydb.googleapis.com"
 
@@ -294,3 +297,39 @@ async def test_async_connect_bad_ip_type(
             exc_info.value.args[0]
             == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'PUBLIC', 'PRIVATE', 'PSC'."
         )
+
+
+async def test_Connector_remove_cached_bad_instance(
+    credentials: FakeCredentials,
+) -> None:
+    """When a Connector attempts to retrieve connection info for a
+    non-existent instance, it should delete the instance from
+    the cache and ensure no background refresh happens (which would be
+    wasted cycles).
+    """
+    instance_uri = "projects/test-project/locations/test-region/clusters/test-cluster/instances/bad-test-instance"
+    async with AsyncConnector(credentials=credentials) as connector:
+        with pytest.raises(ClientResponseError):
+            await connector.connect(instance_uri, "asyncpg")
+        assert instance_uri not in connector._cache
+
+
+async def test_Connector_remove_cached_no_ip_type(credentials: FakeCredentials) -> None:
+    """When a Connector attempts to connect and preferred IP type is not present,
+    it should delete the instance from the cache and ensure no background refresh
+    happens (which would be wasted cycles).
+    """
+    instance_uri = "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance"
+    # set instance to only have Public IP
+    fake_client = FakeAlloyDBClient()
+    fake_client.instance.ip_addrs = {"PUBLIC": "127.0.0.1"}
+    async with AsyncConnector(credentials=credentials) as connector:
+        connector._client = fake_client
+        # populate cache
+        cache = RefreshAheadCache(instance_uri, fake_client, connector._keys)
+        connector._cache[instance_uri] = cache
+        # test instance does not have Private IP, thus should invalidate cache
+        with pytest.raises(IPTypeNotFoundError):
+            await connector.connect(instance_uri, "asyncpg", ip_type="private")
+        # check that cache has been removed from dict
+        assert instance_uri not in connector._cache

tests/unit/test_connector.py

@@ -16,13 +16,17 @@
 from threading import Thread
 from typing import Union
 
+from aiohttp import ClientResponseError
 from mock import patch
 from mocks import FakeAlloyDBClient
 from mocks import FakeCredentials
 import pytest
 
 from google.cloud.alloydb.connector import Connector
 from google.cloud.alloydb.connector import IPTypes
+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
 
 
 def test_Connector_init(credentials: FakeCredentials) -> None:
@@ -203,3 +207,50 @@ def test_Connector_close_called_multiple_times(credentials: FakeCredentials) ->
     assert connector._thread.is_alive() is False
     # call connector.close a second time
     connector.close()
+
+
+async def test_Connector_remove_cached_bad_instance(
+    credentials: FakeCredentials,
+) -> None:
+    """When a Connector attempts to retrieve connection info for a
+    non-existent instance, it should delete the instance from
+    the cache and ensure no background refresh happens (which would be
+    wasted cycles).
+    """
+    instance_uri = "projects/test-project/locations/test-region/clusters/test-cluster/instances/bad-test-instance"
+    with Connector(credentials) as connector:
+        connector._keys = asyncio.wrap_future(
+            asyncio.run_coroutine_threadsafe(
+                generate_keys(), asyncio.get_running_loop()
+            ),
+            loop=asyncio.get_running_loop(),
+        )
+        with pytest.raises(ClientResponseError):
+            await connector.connect_async(instance_uri, "pg8000")
+        assert instance_uri not in connector._cache
+
+
+async def test_Connector_remove_cached_no_ip_type(credentials: FakeCredentials) -> None:
+    """When a Connector attempts to connect and preferred IP type is not present,
+    it should delete the instance from the cache and ensure no background refresh
+    happens (which would be wasted cycles).
+    """
+    instance_uri = "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance"
+    # set instance to only have Public IP
+    fake_client = FakeAlloyDBClient()
+    fake_client.instance.ip_addrs = {"PUBLIC": "127.0.0.1"}
+    with Connector(credentials=credentials) as connector:
+        connector._client = fake_client
+        connector._keys = asyncio.wrap_future(
+            asyncio.run_coroutine_threadsafe(
+                generate_keys(), asyncio.get_running_loop()
+            ),
+            loop=asyncio.get_running_loop(),
+        )
+        cache = RefreshAheadCache(instance_uri, fake_client, connector._keys)
+        connector._cache[instance_uri] = cache
+        # test instance does not have Private IP, thus should invalidate cache
+        with pytest.raises(IPTypeNotFoundError):
+            await connector.connect_async(instance_uri, "pg8000", ip_type="private")
+        # check that cache has been removed from dict
+        assert instance_uri not in connector._cache