refactor: rename Instance to RefreshAheadCache (#1068)

Author: jackwotherspoon

google/cloud/sql/connector/connector.py

@@ -32,8 +32,8 @@
 from google.cloud.sql.connector.client import CloudSQLClient
 from google.cloud.sql.connector.exceptions import ConnectorLoopError
 from google.cloud.sql.connector.exceptions import DnsNameResolutionError
-from google.cloud.sql.connector.instance import Instance
 from google.cloud.sql.connector.instance import IPTypes
+from google.cloud.sql.connector.instance import RefreshAheadCache
 import google.cloud.sql.connector.pg8000 as pg8000
 import google.cloud.sql.connector.pymysql as pymysql
 import google.cloud.sql.connector.pytds as pytds
@@ -113,7 +113,7 @@ def __init__(
                 asyncio.run_coroutine_threadsafe(generate_keys(), self._loop),
                 loop=self._loop,
             )
-        self._instances: Dict[str, Instance] = {}
+        self._cache: Dict[str, RefreshAheadCache] = {}
         self._client: Optional[CloudSQLClient] = None
 
         # initialize credentials
@@ -255,23 +255,23 @@ async def connect_async(
                 driver=driver,
             )
         enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)
-        if instance_connection_string in self._instances:
-            instance = self._instances[instance_connection_string]
-            if enable_iam_auth != instance._enable_iam_auth:
+        if instance_connection_string in self._cache:
+            cache = self._cache[instance_connection_string]
+            if enable_iam_auth != cache._enable_iam_auth:
                 raise ValueError(
                     f"connect() called with 'enable_iam_auth={enable_iam_auth}', "
-                    f"but previously used 'enable_iam_auth={instance._enable_iam_auth}'. "
+                    f"but previously used 'enable_iam_auth={cache._enable_iam_auth}'. "
                     "If you require both for your use case, please use a new "
                     "connector.Connector object."
                 )
         else:
-            instance = Instance(
+            cache = RefreshAheadCache(
                 instance_connection_string,
                 self._client,
                 self._keys,
                 enable_iam_auth,
             )
-            self._instances[instance_connection_string] = instance
+            self._cache[instance_connection_string] = cache
 
         connect_func = {
             "pymysql": pymysql.connect,
@@ -300,7 +300,7 @@ async def connect_async(
 
         # attempt to make connection to Cloud SQL instance
         try:
-            instance_data, ip_address = await instance.connect_info(ip_type)
+            instance_data, ip_address = await cache.connect_info(ip_type)
             # resolve DNS name into IP address for PSC
             if ip_type.value == "PSC":
                 addr_info = await self._loop.getaddrinfo(
@@ -339,7 +339,7 @@ async def connect_async(
 
         except Exception:
             # with any exception, we attempt a force refresh, then throw the error
-            await instance.force_refresh()
+            await cache.force_refresh()
             raise
 
     def __enter__(self) -> Any:
@@ -385,11 +385,9 @@ def close(self) -> None:
             self._thread.join()
 
     async def close_async(self) -> None:
-        """Helper function to cancel Instances' tasks
+        """Helper function to cancel the cache's tasks
         and close aiohttp.ClientSession."""
-        await asyncio.gather(
-            *[instance.close() for instance in self._instances.values()]
-        )
+        await asyncio.gather(*[cache.close() for cache in self._cache.values()])
         if self._client:
             await self._client.close()
 

google/cloud/sql/connector/instance.py

@@ -144,32 +144,13 @@ def get_preferred_ip(self, ip_type: IPTypes) -> str:
         )
 
 
-class Instance:
-    """A class to manage the details of the connection to a Cloud SQL
-    instance, including refreshing the credentials.
-
-    :param instance_connection_string:
-        The Google Cloud SQL Instance's connection
-        string.
-    :type instance_connection_string: str
-
-    :param enable_iam_auth
-        Enables automatic IAM database authentication for Postgres or MySQL
-        instances.
-    :type enable_iam_auth: bool
-    """
-
-    _enable_iam_auth: bool
-    _keys: asyncio.Future
-    _instance_connection_string: str
-    _instance: str
-    _project: str
-    _region: str
+class RefreshAheadCache:
+    """Cache that refreshes connection info in the background prior to expiration.
 
-    _refresh_rate_limiter: AsyncRateLimiter
-    _refresh_in_progress: asyncio.locks.Event
-    _current: asyncio.Task  # task wraps coroutine that returns ConnectionInfo
-    _next: asyncio.Task  # task wraps coroutine that returns another task
+    Background tasks are used to schedule refresh attempts to get a new
+    ephemeral certificate and Cloud SQL metadata (IP addresses, etc.) ahead of
+    expiration.
+    """
 
     def __init__(
         self,
@@ -178,6 +159,18 @@ def __init__(
         keys: asyncio.Future,
         enable_iam_auth: bool = False,
     ) -> None:
+        """Initializes a RefreshAheadCache instance.
+
+        Args:
+            instance_connection_string (str): The Cloud SQL Instance's
+                connection string (also known as an instance connection name).
+            client (CloudSQLClient): The Cloud SQL Client instance.
+            keys (asyncio.Future): A future to the client's public-private key
+                pair.
+            enable_iam_auth (bool): Enables automatic IAM database authentication
+                (Postgres and MySQL) as the default authentication method for all
+                connections.
+        """
         # validate and parse instance connection name
         self._project, self._region, self._instance = _parse_instance_connection_name(
             instance_connection_string
@@ -192,8 +185,8 @@ def __init__(
             rate=1 / 30,
         )
         self._refresh_in_progress = asyncio.locks.Event()
-        self._current = self._schedule_refresh(0)
-        self._next = self._current
+        self._current: asyncio.Task = self._schedule_refresh(0)
+        self._next: asyncio.Task = self._current
 
     async def force_refresh(self) -> None:
         """
@@ -211,10 +204,11 @@ async def _perform_refresh(self) -> ConnectionInfo:
         """Retrieves instance metadata and ephemeral certificate from the
         Cloud SQL Instance.
 
-        :rtype: ConnectionInfo
-        :returns: A dataclass containing a string representing the ephemeral certificate, a dict
-            containing the instances IP adresses, a string representing a PEM-encoded private key
-            and a string representing a PEM-encoded certificate authority.
+        Returns:
+            A ConnectionInfo instance containing a string representing the
+            ephemeral certificate, a dict containing the instances IP adresses,
+            a string representing a PEM-encoded private key and a string
+            representing a PEM-encoded certificate authority.
         """
         self._refresh_in_progress.set()
         logger.debug(
@@ -290,15 +284,14 @@ def _schedule_refresh(self, delay: int) -> asyncio.Task:
         """
         Schedule task to sleep and then perform refresh to get ConnectionInfo.
 
-        :type delay: int
-        :param delay
-            Time in seconds to sleep before running _perform_refresh.
+        Args:
+            delay (int): Time in seconds to sleep before performing a refresh.
 
-        :rtype: asyncio.Task
-        :returns: A Task representing the scheduled _perform_refresh.
+        Returns:
+            An asyncio.Task representing the scheduled refresh.
         """
 
-        async def _refresh_task(self: Instance, delay: int) -> ConnectionInfo:
+        async def _refresh_task(self: RefreshAheadCache, delay: int) -> ConnectionInfo:
             """
             A coroutine that sleeps for the specified amount of time before
             running _perform_refresh.
@@ -349,16 +342,14 @@ async def connect_info(
         """Retrieve instance metadata and ip address required
         for making connection to Cloud SQL instance.
 
-        :type ip_type: IPTypes
-        :param ip_type: Enum specifying whether to look for public
-            or private IP address.
-
-        :rtype instance_data: ConnectionInfo
-        :returns: Instance metadata for Cloud SQL instance.
+        Args:
+            ip_type (IPTypes): Enum specifying type of IP address to lookup and
+                use for connection.
 
-        :rtype ip_address: str
-        :returns: A string representing the IP address of
-            the given Cloud SQL instance.
+        Returns:
+            A tuple with the first item being the ConnectionInfo instance for
+            establishing the connection, and the second item being the IP
+            address of the Cloud SQL instance matching the specified IP type.
         """
         logger.debug(
             f"['{self._instance_connection_string}']: Entered connect_info method"

tests/conftest.py

@@ -25,7 +25,7 @@
 from unit.mocks import FakeCSQLInstance  # type: ignore
 
 from google.cloud.sql.connector.client import CloudSQLClient
-from google.cloud.sql.connector.instance import Instance
+from google.cloud.sql.connector.instance import RefreshAheadCache
 from google.cloud.sql.connector.utils import generate_keys
 
 SCOPES = ["https://www.googleapis.com/auth/sqlservice.admin"]
@@ -137,12 +137,12 @@ async def fake_client(
 
 
 @pytest.fixture
-async def instance(fake_client: CloudSQLClient) -> AsyncGenerator[Instance, None]:
+async def cache(fake_client: CloudSQLClient) -> AsyncGenerator[RefreshAheadCache, None]:
     keys = asyncio.create_task(generate_keys())
-    instance = Instance(
+    cache = RefreshAheadCache(
         "test-project:test-region:test-instance",
         client=fake_client,
         keys=keys,
     )
-    yield instance
-    await instance.close()
+    yield cache
+    await cache.close()

tests/system/test_connector_object.py

@@ -83,11 +83,11 @@ def test_multiple_connectors() -> None:
             conn.execute(sqlalchemy.text("SELECT 1"))
 
         instance_connection_string = os.environ["MYSQL_CONNECTION_NAME"]
-        assert instance_connection_string in first_connector._instances
-        assert instance_connection_string in second_connector._instances
+        assert instance_connection_string in first_connector._cache
+        assert instance_connection_string in second_connector._cache
         assert (
-            first_connector._instances[instance_connection_string]
-            != second_connector._instances[instance_connection_string]
+            first_connector._cache[instance_connection_string]
+            != second_connector._cache[instance_connection_string]
         )
     except Exception as e:
         logging.exception("Failed to connect with multiple Connector objects!", e)

tests/unit/test_connector.py

@@ -25,18 +25,18 @@
 from google.cloud.sql.connector import IPTypes
 from google.cloud.sql.connector.client import CloudSQLClient
 from google.cloud.sql.connector.exceptions import ConnectorLoopError
-from google.cloud.sql.connector.instance import Instance
+from google.cloud.sql.connector.instance import RefreshAheadCache
 
 
 def test_connect_enable_iam_auth_error(
-    fake_credentials: Credentials, instance: Instance
+    fake_credentials: Credentials, cache: RefreshAheadCache
 ) -> None:
     """Test that calling connect() with different enable_iam_auth
     argument values throws error."""
     connect_string = "test-project:test-region:test-instance"
     connector = Connector(credentials=fake_credentials)
-    # set instance
-    connector._instances[connect_string] = instance
+    # set cache
+    connector._cache[connect_string] = cache
     # try to connect using enable_iam_auth=True, should raise error
     with pytest.raises(ValueError) as exc_info:
         connector.connect(connect_string, "pg8000", enable_iam_auth=True)
@@ -46,8 +46,8 @@ def test_connect_enable_iam_auth_error(
         "If you require both for your use case, please use a new "
         "connector.Connector object."
     )
-    # remove instance to avoid destructor warnings
-    connector._instances = {}
+    # remove cache entrry to avoid destructor warnings
+    connector._cache = {}
 
 
 def test_connect_with_unsupported_driver(fake_credentials: Credentials) -> None: