chore: move socket initialization to Connector level

Author: jackwotherspoon

google/cloud/sql/connector/connection_info.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import abc
 from dataclasses import dataclass
 import logging
 import ssl
@@ -34,6 +35,27 @@
 logger = logging.getLogger(name=__name__)
 
 
+class ConnectionInfoCache(abc.ABC):
+    """Abstract class for Connector connection info caches."""
+
+    @abc.abstractmethod
+    async def connect_info(self) -> ConnectionInfo:
+        pass
+
+    @abc.abstractmethod
+    async def force_refresh(self) -> None:
+        pass
+
+    @abc.abstractmethod
+    async def close(self) -> None:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def closed(self) -> bool:
+        pass
+
+
 @dataclass
 class ConnectionInfo:
     """Contains all necessary information to connect securely to the

google/cloud/sql/connector/connector.py

@@ -20,6 +20,7 @@
 from functools import partial
 import logging
 import os
+import socket
 from threading import Thread
 from types import TracebackType
 from typing import Any, Optional, Union
@@ -47,6 +48,7 @@
 logger = logging.getLogger(name=__name__)
 
 ASYNC_DRIVERS = ["asyncpg"]
+SERVER_PROXY_PORT = 3307
 _DEFAULT_SCHEME = "https://"
 _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
 _SQLADMIN_HOST_TEMPLATE = "sqladmin.{universe_domain}"
@@ -291,10 +293,11 @@ async def connect_async(
                 driver=driver,
             )
         enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)
-        if (instance_connection_string, enable_iam_auth) in self._cache:
-            monitored_cache = self._cache[(instance_connection_string, enable_iam_auth)]
+
+        conn_name = await self._resolver.resolve(instance_connection_string)
+        if (str(conn_name), enable_iam_auth) in self._cache:
+            monitored_cache = self._cache[(str(conn_name), enable_iam_auth)]
         else:
-            conn_name = await self._resolver.resolve(instance_connection_string)
             if self._refresh_strategy == RefreshStrategy.LAZY:
                 logger.debug(
                     f"['{conn_name}']: Refresh strategy is set to lazy refresh"
@@ -322,7 +325,7 @@ async def connect_async(
                 self._resolver,
             )
             logger.debug(f"['{conn_name}']: Connection info added to cache")
-            self._cache[(instance_connection_string, enable_iam_auth)] = monitored_cache
+            self._cache[(str(conn_name), enable_iam_auth)] = monitored_cache
 
         connect_func = {
             "pymysql": pymysql.connect,
@@ -358,7 +361,7 @@ async def connect_async(
         except Exception:
             # with an error from Cloud SQL Admin API call or IP type, invalidate
             # the cache and re-raise the error
-            await self._remove_cached(instance_connection_string, enable_iam_auth)
+            await self._remove_cached(str(conn_name), enable_iam_auth)
             raise
         logger.debug(f"['{conn_info.conn_name}']: Connecting to {ip_address}:3307")
         # format `user` param for automatic IAM database authn
@@ -379,11 +382,21 @@ async def connect_async(
                     await conn_info.create_ssl_context(enable_iam_auth),
                     **kwargs,
                 )
-            # synchronous drivers are blocking and run using executor
+            # Create socket with SSLContext for sync drivers
+            ctx = await conn_info.create_ssl_context(enable_iam_auth)
+            sock = ctx.wrap_socket(
+                socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+                server_hostname=ip_address,
+            )
+            # If this connection was opened using a domain name, then store it
+            # for later in case we need to forcibly close it on failover.
+            if conn_info.conn_name.domain_name:
+                monitored_cache.sockets.append(sock)
+            # Synchronous drivers are blocking and run using executor
             connect_partial = partial(
                 connector,
                 ip_address,
-                await conn_info.create_ssl_context(enable_iam_auth),
+                sock,
                 **kwargs,
             )
             return await self._loop.run_in_executor(None, connect_partial)
@@ -468,6 +481,7 @@ async def create_async_connector(
     universe_domain: Optional[str] = None,
     refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
     resolver: type[DefaultResolver] | type[DnsResolver] = DefaultResolver,
+    failover_period: int = 30,
 ) -> Connector:
     """Helper function to create Connector object for asyncio connections.
 
@@ -519,6 +533,11 @@ async def create_async_connector(
             DnsResolver.
             Default: DefaultResolver
 
+        failover_period (int): The time interval in seconds between each
+            attempt to check if a failover has occured for a given instance.
+            Must be used with `resolver=DnsResolver` to have any effect.
+            Default: 30
+
     Returns:
         A Connector instance configured with running event loop.
     """
@@ -537,4 +556,5 @@ async def create_async_connector(
         universe_domain=universe_domain,
         refresh_strategy=refresh_strategy,
         resolver=resolver,
+        failover_period=failover_period,
     )

google/cloud/sql/connector/monitored_cache.py

@@ -14,6 +14,7 @@
 
 import asyncio
 import logging
+import ssl
 from typing import Any, Callable, Optional, Union
 
 from google.cloud.sql.connector.connection_info import ConnectionInfo
@@ -36,7 +37,7 @@ def __init__(
         self.resolver = resolver
         self.cache = cache
         self.domain_name_ticker: Optional[asyncio.Task] = None
-        self.open_conns: int = 0
+        self.sockets: list[ssl.SSLSocket] = []
 
         if self.cache.conn_name.domain_name:
             self.domain_name_ticker = asyncio.create_task(
@@ -51,6 +52,15 @@ def __init__(
     def closed(self) -> bool:
         return self.cache.closed
 
+    async def _purge_closed_sockets(self) -> None:
+        open_sockets = []
+        for socket in self.sockets:
+            # Check fileno as method to check if socket is closed. Will return
+            # -1 on failure, which will be used to signal socket closed.
+            if socket.fileno() != -1:
+                open_sockets.append(socket)
+        self.sockets = open_sockets
+
     async def _check_domain_name(self) -> None:
         try:
             # Resolve domain name and see if Cloud SQL instance connection name
@@ -66,12 +76,6 @@ async def _check_domain_name(self) -> None:
                     "connections!"
                 )
                 await self.close()
-                conn_info = await self.connect_info()
-                if conn_info.sock:
-                    logger.debug(f"Socket type: {type(conn_info.sock)}")
-                    conn_info.sock.close()
-            else:
-                logger.debug("Domain name mapping has not changed!")
 
         except Exception as e:
             # Domain name checks should not be fatal, log error and continue.
@@ -97,10 +101,20 @@ async def close(self) -> None:
                 logger.debug(
                     f"['{self.cache.conn_name}']: Cancelled domain name polling task."
                 )
-
+            finally:
+                self.domain_name_ticker = None
         # If cache is already closed, no further work.
-        if self.cache.closed:
+        if self.closed:
             return
+
+        # Close any still open sockets
+        for socket in self.sockets:
+            # Check fileno as method to check if socket is closed. Will return
+            # -1 on failure, which will be used to signal socket closed.
+            if socket.fileno() != -1:
+                socket.close()
+
+        # Close underyling ConnectionInfoCache
         await self.cache.close()
 
 

google/cloud/sql/connector/pg8000.py

@@ -14,7 +14,6 @@
 limitations under the License.
 """
 
-import socket
 import ssl
 from typing import Any, TYPE_CHECKING
 
@@ -25,16 +24,16 @@
 
 
 def connect(
-    ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
+    ip_address: str, sock: ssl.SSLSocket, **kwargs: Any
 ) -> "pg8000.dbapi.Connection":
     """Helper function to create a pg8000 DB-API connection object.
 
     :type ip_address: str
     :param ip_address: A string containing an IP address for the Cloud SQL
         instance.
 
-    :type ctx: ssl.SSLContext
-    :param ctx: An SSLContext object created from the Cloud SQL server CA
+    :type sock: ssl.SSLSocket
+    :param sock: An SSLSocket object created from the Cloud SQL server CA
         cert and ephemeral cert.
 
 
@@ -48,12 +47,6 @@ def connect(
             'Unable to import module "pg8000." Please install and try again.'
         )
 
-    # Create socket and wrap with context.
-    sock = ctx.wrap_socket(
-        socket.create_connection((ip_address, SERVER_PROXY_PORT)),
-        server_hostname=ip_address,
-    )
-
     user = kwargs.pop("user")
     db = kwargs.pop("db")
     passwd = kwargs.pop("password", None)