fix(main): Make local proxy to accept multiple connections (WIP)

Author: Uziel Silva

google/cloud/sql/connector/connector.py

@@ -38,10 +38,10 @@
 from google.cloud.sql.connector.exceptions import ConnectorLoopError
 from google.cloud.sql.connector.instance import RefreshAheadCache
 from google.cloud.sql.connector.lazy import LazyRefreshCache
+import google.cloud.sql.connector.local_unix_socket as local_unix_socket
 from google.cloud.sql.connector.monitored_cache import MonitoredCache
 import google.cloud.sql.connector.pg8000 as pg8000
-import google.cloud.sql.connector.proxy as proxy
-import google.cloud.sql.connector.psycopg as psycopg
+from google.cloud.sql.connector.proxy import Proxy
 import google.cloud.sql.connector.pymysql as pymysql
 import google.cloud.sql.connector.pytds as pytds
 from google.cloud.sql.connector.resolver import DefaultResolver
@@ -52,7 +52,6 @@
 logger = logging.getLogger(name=__name__)
 
 ASYNC_DRIVERS = ["asyncpg"]
-LOCAL_PROXY_DRIVERS = ["psycopg"]
 SERVER_PROXY_PORT = 3307
 _DEFAULT_SCHEME = "https://"
 _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
@@ -160,7 +159,7 @@ def __init__(
         self._cache: dict[tuple[str, bool], MonitoredCache] = {}
         self._client: Optional[CloudSQLClient] = None
         self._closed: bool = False
-        self._proxy: Optional[asyncio.Task] = None
+        self._proxies: Optional[Proxy] = None
 
         # initialize credentials
         scopes = ["https://www.googleapis.com/auth/sqlservice.admin"]
@@ -221,6 +220,29 @@ def __init__(
     def universe_domain(self) -> str:
         return self._universe_domain or _DEFAULT_UNIVERSE_DOMAIN
 
+    def start_unix_socket_proxy_async(
+        self,
+        instance_connection_name: str,
+        local_socket_path: str,
+        **kwargs: Any
+    ) -> None:
+        """Creates a new Proxy instance and stores it to properly disposal
+
+        Args:
+            instance_connection_string (str): The instance connection name of the
+                Cloud SQL instance to connect to. Takes the form of
+                "project-id:region:instance-name"
+
+                Example: "my-project:us-central1:my-instance"
+
+            local_socket_path (str): A string representing the location of the local socket.
+
+            **kwargs: Any driver-specific arguments to pass to the underlying
+                driver .connect call.
+        """
+        # TODO: validates the local socket path is not the same as other invocation
+        self._proxies.append(new Proxy(self, instance_connection_name, local_socket_path, self.loop, **kwargs))
+
     def connect(
         self, instance_connection_string: str, driver: str, **kwargs: Any
     ) -> Any:
@@ -238,7 +260,7 @@ def connect(
                 Example: "my-project:us-central1:my-instance"
 
             driver (str): A string representing the database driver to connect
-                with. Supported drivers are pymysql, pg8000, psycopg, and pytds.
+                with. Supported drivers are pymysql, pg8000, local_unix_socket, and pytds.
 
             **kwargs: Any driver-specific arguments to pass to the underlying
                 driver .connect call.
@@ -280,7 +302,7 @@ async def connect_async(
                 Example: "my-project:us-central1:my-instance"
 
             driver (str): A string representing the database driver to connect
-                with. Supported drivers are pymysql, asyncpg, pg8000, psycopg, and
+                with. Supported drivers are pymysql, asyncpg, pg8000, local_unix_socket, and
                 pytds.
 
             **kwargs: Any driver-specific arguments to pass to the underlying
@@ -293,7 +315,7 @@ async def connect_async(
             ValueError: Connection attempt with built-in database authentication
                 and then subsequent attempt with IAM database authentication.
             KeyError: Unsupported database driver Must be one of pymysql, asyncpg,
-                pg8000, psycopg, and pytds.
+                pg8000, local_unix_socket, and pytds.
             RuntimeError: Connector has been closed. Cannot connect using a closed
                 Connector.
         """
@@ -362,7 +384,7 @@ async def connect_async(
         connect_func = {
             "pymysql": pymysql.connect,
             "pg8000": pg8000.connect,
-            "psycopg": psycopg.connect,
+            "local_unix_socket": local_unix_socket.connect,
             "asyncpg": asyncpg.connect,
             "pytds": pytds.connect,
         }
@@ -449,17 +471,6 @@ async def connect_async(
                 server_hostname=ip_address,
             )
 
-            host = ip_address
-            # start local proxy if driver needs it
-            if driver in LOCAL_PROXY_DRIVERS:
-                local_socket_path = kwargs.pop("local_socket_path", "/tmp/connector-socket")
-                host = local_socket_path
-                self._proxy = proxy.start_local_proxy(
-                    sock,
-                    socket_path=f"{local_socket_path}/.s.PGSQL.{SERVER_PROXY_PORT}",
-                    loop=self._loop
-                )
-
             # 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:
@@ -543,6 +554,7 @@ async def close_async(self) -> None:
             await self._client.close()
         await asyncio.gather(*[cache.close() for cache in self._cache.values()])
 
+
 async def create_async_connector(
     ip_type: str | IPTypes = IPTypes.PUBLIC,
     enable_iam_auth: bool = False,

google/cloud/sql/connector/enums.py

@@ -62,7 +62,7 @@ class DriverMapping(Enum):
 
     ASYNCPG = "POSTGRES"
     PG8000 = "POSTGRES"
-    PSYCOPG = "POSTGRES"
+    LOCAL_UNIX_SOCKET = "ANY"
     PYMYSQL = "MYSQL"
     PYTDS = "SQLSERVER"
 
@@ -79,7 +79,7 @@ def validate_engine(driver: str, engine_version: str) -> None:
                 the given engine.
         """
         mapping = DriverMapping[driver.upper()]
-        if not engine_version.startswith(mapping.value):
+        if not mapping.value == "ANY" and not engine_version.startswith(mapping.value):
             raise IncompatibleDriverError(
                 f"Database driver '{driver}' is incompatible with database "
                 f"version '{engine_version}'. Given driver can "

google/cloud/sql/connector/psycopg.py …cloud/sql/connector/local_unix_socket.pygoogle/cloud/sql/connector/psycopg.py renamed to google/cloud/sql/connector/local_unix_socket.py google/cloud/sql/connector/psycopg.py renamed to google/cloud/sql/connector/local_unix_socket.py

@@ -19,44 +19,19 @@
 
 SERVER_PROXY_PORT = 3307
 
-if TYPE_CHECKING:
-    import psycopg
-
-
 def connect(
     host: str, sock: ssl.SSLSocket, **kwargs: Any
-) -> "psycopg.Connection":
-    """Helper function to create a psycopg DB-API connection object.
+) -> "ssl.SSLSocket":
+    """Helper function to retrieve the socket for local UNIX sockets.
 
     Args:
         host (str): A string containing the socket path used by the local proxy.
         sock (ssl.SSLSocket): An SSLSocket object created from the Cloud SQL
             server CA cert and ephemeral cert.
-        kwargs: Additional arguments to pass to the psycopg connect method.
+        kwargs: Additional arguments to pass to the local UNIX socket connect method.
 
     Returns:
-        psycopg.Connection: A psycopg connection to the Cloud SQL
-            instance.
-
-    Raises:
-        ImportError: The psycopg module cannot be imported.
+        ssl.SSLSocket: The same socket
     """
-    try:
-        from psycopg import Connection
-    except ImportError:
-        raise ImportError(
-            'Unable to import module "psycopg." Please install and try again.'
-        )
-
-    user = kwargs.pop("user")
-    db = kwargs.pop("db")
-    passwd = kwargs.pop("password", None)
-
-    kwargs.pop("timeout", None)
-
-    conn = Connection.connect(
-        f"host={host} port={SERVER_PROXY_PORT} dbname={db} user={user} password={passwd} sslmode=require",
-        **kwargs
-    )
-
-    return conn
+    
+    return sock

google/cloud/sql/connector/proxy.py

@@ -18,73 +18,125 @@
 import os
 from pathlib import Path
 import socket
+import selectors
 import ssl
 
+from google.cloud.sql.connector import Connector
 from google.cloud.sql.connector.exceptions import LocalProxyStartupError
 
 SERVER_PROXY_PORT = 3307
 LOCAL_PROXY_MAX_MESSAGE_SIZE = 10485760
 
-def start_local_proxy(
-    ssl_sock: ssl.SSLSocket,
-    socket_path: str,
-    loop: asyncio.AbstractEventLoop
-) -> asyncio.Task:
-    """Helper function to start a UNIX based local proxy for
-    transport messages through the SSL Socket.
-
-    Args:
-        ssl_sock (ssl.SSLSocket): An SSLSocket object created from the Cloud SQL
-            server CA cert and ephemeral cert.
-        socket_path: A system path that is going to be used to store the socket.
-        loop (asyncio.AbstractEventLoop): Event loop to run asyncio tasks.
-
-    Returns:
-        asyncio.Task: The asyncio task containing the proxy server process.
-
-    Raises:
-        LocalProxyStartupError: Local UNIX socket based proxy was not able to
-        get started.
-    """
-    unix_socket = None
-
-    try:
-        path_parts = socket_path.rsplit('/', 1)
-        parent_directory = '/'.join(path_parts[:-1])
-
-        desired_path = Path(parent_directory)
-        desired_path.mkdir(parents=True, exist_ok=True)
-
-        if os.path.exists(socket_path):
-            os.remove(socket_path)
-        unix_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-
-        unix_socket.bind(socket_path)
-        unix_socket.listen(1)
-        unix_socket.setblocking(False)
-        os.chmod(socket_path, 0o600)
-    except Exception:
-        raise LocalProxyStartupError(
-            'Local UNIX socket based proxy was not able to get started.'
-        )
-
-    return loop.create_task(local_communication(unix_socket, ssl_sock, socket_path, loop))
-
-
-async def local_communication(
-  unix_socket, ssl_sock, socket_path, loop
-):
-    client, _ = await loop.sock_accept(unix_socket)
-    
-    try:
+
+class Proxy:
+    """Creates an "accept loop" async task which will open the unix server socket and listen for new connections."""
+
+    def __init__(
+        self,
+        connector: Connector,
+        instance_connection_string: str,
+        socket_path: str,
+        loop: asyncio.AbstractEventLoop,
+        **kwargs: Any
+    ) -> None:
+        """Keeps track of all the async tasks and starts the accept loop for new connections.
+        
+        Args:
+            connector (Connector): The instance where this Proxy class was created.
+
+            instance_connection_string (str): The instance connection name of the
+                Cloud SQL instance to connect to. Takes the form of
+                "project-id:region:instance-name"
+
+                Example: "my-project:us-central1:my-instance"
+
+            socket_path (str): A system path that is going to be used to store the socket.
+
+            loop (asyncio.AbstractEventLoop): Event loop to run asyncio tasks.
+
+            **kwargs: Any driver-specific arguments to pass to the underlying
+                driver .connect call.
+        """
+        self._connection_tasks = []
+        self._addr = instance_connection_string
+        self._kwargs = kwargs
+        self._connector = connector
+        self._task = loop.create_task(accept_loop(socket_path, loop, **kwargs))
+
+    async def accept_loop(
+        self
+        socket_path: str,
+        loop: asyncio.AbstractEventLoop
+    ) -> asyncio.Task:
+        """Starts a UNIX based local proxy for transporting messages through
+        the SSL Socket, and waits until there is a new connection to accept, to register it
+        and keep track of it.
+        
+        Args:
+            socket_path: A system path that is going to be used to store the socket.
+
+            loop (asyncio.AbstractEventLoop): Event loop to run asyncio tasks.
+
+        Raises:
+            LocalProxyStartupError: Local UNIX socket based proxy was not able to
+            get started.
+        """
+        unix_socket = None
+        sel = selectors.DefaultSelector()
+
+        try:
+            path_parts = socket_path.rsplit('/', 1)
+            parent_directory = '/'.join(path_parts[:-1])
+
+            desired_path = Path(parent_directory)
+            desired_path.mkdir(parents=True, exist_ok=True)
+
+            if os.path.exists(socket_path):
+                os.remove(socket_path)
+
+            unix_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+
+            unix_socket.bind(socket_path)
+            unix_socket.listen(1)
+            unix_socket.setblocking(False)
+            os.chmod(socket_path, 0o600)
+            
+            sel.register(unix_socket, selectors.EVENT_READ, data=None)
+
+        except Exception:
+            raise LocalProxyStartupError(
+                'Local UNIX socket based proxy was not able to get started.'
+            )
+
         while True:
-            data = await loop.sock_recv(client, LOCAL_PROXY_MAX_MESSAGE_SIZE)
-            if not data:
-              client.close()
-              break
-            ssl_sock.sendall(data)
-            response = ssl_sock.recv(LOCAL_PROXY_MAX_MESSAGE_SIZE)
-            await loop.sock_sendall(client, response)
-    finally:
-        client.close()
-        os.remove(socket_path) # Clean up the socket file
+            client, _ = await loop.sock_accept(unix_socket)
+            self._connection_tasks.append(loop.create_task(self.client_socket(client, unix_socket, socket_path, loop))) 
+
+    async def close_async(self):
+        proxy_task = asyncio.gather(self._task)
+        try:
+            await asyncio.wait_for(proxy_task, timeout=0.1)
+        except (asyncio.CancelledError, asyncio.TimeoutError, TimeoutError):
+            pass # This task runs forever so it is expected to throw this exception
+
+
+    async def client_socket(
+        self, client, unix_socket, socket_path, loop
+    ):
+        try:
+            ssl_sock = self.connector.connect(
+                self._addr,
+                'local_unix_socket',
+                **self._kwargs
+            )
+            while True:
+                data = await loop.sock_recv(client, LOCAL_PROXY_MAX_MESSAGE_SIZE)
+                if not data:
+                    client.close()
+                    break
+                ssl_sock.sendall(data)
+                response = ssl_sock.recv(LOCAL_PROXY_MAX_MESSAGE_SIZE)
+                await loop.sock_sendall(client, response)
+        finally:
+            client.close()
+            os.remove(socket_path) # Clean up the socket file

pyproject.toml

@@ -59,7 +59,6 @@ Changelog = "https://github.com/GoogleCloudPlatform/cloud-sql-python-connector/b
 [project.optional-dependencies]
 pymysql = ["PyMySQL>=1.1.0"]
 pg8000 = ["pg8000>=1.31.1"]
-psycopg = ["psycopg>=3.2.9"]
 pytds = ["python-tds>=1.15.0"]
 asyncpg = ["asyncpg>=0.30.0"]