chore: attempt moving socket into ConnectionInfo

Author: jackwotherspoon

google/cloud/sql/connector/__init__.py

@@ -1,4 +1,4 @@
-""""
+"""
 Copyright 2019 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");

google/cloud/sql/connector/connection_info.py

@@ -17,20 +17,21 @@
 import abc
 from dataclasses import dataclass
 import logging
+import socket
 import ssl
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, Optional, TYPE_CHECKING, Union
 
 from aiofiles.tempfile import TemporaryDirectory
 
 from google.cloud.sql.connector.connection_name import ConnectionName
+from google.cloud.sql.connector.enums import IPTypes
 from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
 from google.cloud.sql.connector.exceptions import TLSVersionError
 from google.cloud.sql.connector.utils import write_to_file
 
 if TYPE_CHECKING:
     import datetime
 
-    from google.cloud.sql.connector.enums import IPTypes
 
 logger = logging.getLogger(name=__name__)
 
@@ -69,13 +70,21 @@ class ConnectionInfo:
     database_version: str
     expiration: datetime.datetime
     context: Optional[ssl.SSLContext] = None
+    sock: Optional[ssl.SSLSocket] = None
 
-    async def create_ssl_context(self, enable_iam_auth: bool = False) -> ssl.SSLContext:
+    async def create_ssl_context(
+        self, enable_iam_auth: bool = False, return_socket: bool = False
+    ) -> Union[ssl.SSLContext, ssl.SSLSocket]:
         """Constructs a SSL/TLS context for the given connection info.
 
         Cache the SSL context to ensure we don't read from disk repeatedly when
         configuring a secure connection.
         """
+        # Return socket if socket is cached and return_socket is set to True
+        if self.sock is not None and return_socket:
+            logger.debug("Socket in cache, returning it!")
+            return self.sock
+
         # if SSL context is cached, use it
         if self.context is not None:
             return self.context
@@ -116,6 +125,15 @@ async def create_ssl_context(self, enable_iam_auth: bool = False) -> ssl.SSLCont
             context.load_verify_locations(cafile=ca_filename)
         # set class attribute to cache context for subsequent calls
         self.context = context
+        # If return_socket is True, cache socket and return it
+        if return_socket:
+            logger.debug("Returning socket instead of context!")
+            sock = self.context.wrap_socket(
+                socket.create_connection((self.get_preferred_ip(IPTypes.PUBLIC), 3307)),
+                server_hostname="blah",
+            )
+            self.sock = sock
+            return sock
         return context
 
     def get_preferred_ip(self, ip_type: IPTypes) -> str:

google/cloud/sql/connector/connector.py

@@ -390,7 +390,7 @@ async def connect_async(
 
         except Exception:
             # with any exception, we attempt a force refresh, then throw the error
-            await cache.force_refresh()
+            await monitored_cache.force_refresh()
             raise
 
     async def _remove_cached(

google/cloud/sql/connector/monitored_cache.py

@@ -36,7 +36,7 @@ def __init__(
         self.resolver = resolver
         self.cache = cache
         self.domain_name_ticker: Optional[asyncio.Task] = None
-        self.open_conns_count: int = 0
+        self.open_conns: int = 0
 
         if self.cache.conn_name.domain_name:
             self.domain_name_ticker = asyncio.create_task(
@@ -66,6 +66,12 @@ 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.

google/cloud/sql/connector/pymysql.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
 ) -> "pymysql.connections.Connection":
     """Helper function to create a pymysql 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.
 
     :rtype: pymysql.Connection
@@ -50,11 +49,6 @@ def connect(
     # allow automatic IAM database authentication to not require password
     kwargs["password"] = kwargs["password"] if "password" in kwargs else None
 
-    # Create socket and wrap with context.
-    sock = ctx.wrap_socket(
-        socket.create_connection((ip_address, SERVER_PROXY_PORT)),
-        server_hostname=ip_address,
-    )
     # pop timeout as timeout arg is called 'connect_timeout' for pymysql
     timeout = kwargs.pop("timeout")
     kwargs["connect_timeout"] = kwargs.get("connect_timeout", timeout)

google/cloud/sql/connector/pytds.py

@@ -15,7 +15,6 @@
 """
 
 import platform
-import socket
 import ssl
 from typing import Any, TYPE_CHECKING
 
@@ -27,15 +26,15 @@
     import pytds
 
 
-def connect(ip_address: str, ctx: ssl.SSLContext, **kwargs: Any) -> "pytds.Connection":
+def connect(ip_address: str, sock: ssl.SSLSocket, **kwargs: Any) -> "pytds.Connection":
     """Helper function to create a pytds 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.
 
 
@@ -51,11 +50,6 @@ def connect(ip_address: str, ctx: ssl.SSLContext, **kwargs: Any) -> "pytds.Conne
 
     db = kwargs.pop("db", None)
 
-    # Create socket and wrap with context.
-    sock = ctx.wrap_socket(
-        socket.create_connection((ip_address, SERVER_PROXY_PORT)),
-        server_hostname=ip_address,
-    )
     if kwargs.pop("active_directory_auth", False):
         if platform.system() == "Windows":
             # Ignore username and password if using active directory auth

tests/conftest.py

@@ -1,4 +1,4 @@
-""""
+"""
 Copyright 2021 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");

tests/system/test_connector_object.py

@@ -1,4 +1,4 @@
-""""
+"""
 Copyright 2021 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");

tests/system/test_ip_types.py

@@ -1,4 +1,4 @@
-""""
+"""
 Copyright 2021 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");

tests/system/test_pymysql_connection.py

@@ -1,4 +1,4 @@
-""""
+"""
 Copyright 2021 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");