feat(main): Add support for psycopg (Issue #214)

README.md

@@ -42,6 +42,7 @@ The Cloud SQL Python Connector is a package to be used alongside a database driv
 Currently supported drivers are:
  - [`pymysql`](https://github.com/PyMySQL/PyMySQL) (MySQL)
  - [`pg8000`](https://github.com/tlocke/pg8000) (PostgreSQL)
+ - [`psycopg`](https://github.com/psycopg/psycopg) (PostgreSQL)
  - [`asyncpg`](https://github.com/MagicStack/asyncpg) (PostgreSQL)
  - [`pytds`](https://github.com/denisenkom/pytds) (SQL Server)
 
@@ -587,7 +588,7 @@ async def main():
     # acquire connection and query Cloud SQL database
     async with pool.acquire() as conn:
         res = await conn.fetch("SELECT NOW()")
-    
+
     # close Connector
     await connector.close_async()
 ```

google/cloud/sql/connector/client.py

@@ -58,6 +58,7 @@ def __init__(
         client: Optional[aiohttp.ClientSession] = None,
         driver: Optional[str] = None,
         user_agent: Optional[str] = None,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
     ) -> None:
         """Establishes the client to be used for Cloud SQL Admin API requests.
 
@@ -84,8 +85,7 @@ def __init__(
         }
         if quota_project:
             headers["x-goog-user-project"] = quota_project
-
-        self._client = client if client else aiohttp.ClientSession(headers=headers)
+        self._client = client if client else aiohttp.ClientSession(headers=headers, loop=loop)
         self._credentials = credentials
         if sqladmin_api_endpoint is None:
             self._sqladmin_api_endpoint = DEFAULT_SERVICE_ENDPOINT

google/cloud/sql/connector/enums.py

@@ -62,6 +62,7 @@ class DriverMapping(Enum):
 
     ASYNCPG = "POSTGRES"
     PG8000 = "POSTGRES"
+    PSYCOPG = "POSTGRES"
     PYMYSQL = "MYSQL"
     PYTDS = "SQLSERVER"
 

google/cloud/sql/connector/exceptions.py

@@ -84,3 +84,9 @@ class CacheClosedError(Exception):
     Exception to be raised when a ConnectionInfoCache can not be accessed after
     it is closed.
     """
+
+
+class LocalProxyStartupError(Exception):
+    """
+    Exception to be raised when a the local UNIX-socket based proxy can not be started.
+    """

google/cloud/sql/connector/local_unix_socket.py

@@ -0,0 +1,35 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import ssl
+from typing import Any, TYPE_CHECKING
+
+def connect(
+    host: str, sock: ssl.SSLSocket, **kwargs: Any
+) -> "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 local UNIX socket connect method.
+
+    Returns:
+        ssl.SSLSocket: The same socket
+    """
+
+    return sock

google/cloud/sql/connector/proxy.py

@@ -0,0 +1,266 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+import asyncio
+from functools import partial
+import logging
+import os
+from pathlib import Path
+from typing import Callable, List
+
+logger = logging.getLogger(name=__name__)
+
+
+class BaseProxyProtocol(asyncio.Protocol):
+    """
+    A protocol to proxy data between two transports.
+    """
+
+    def __init__(self, proxy: Proxy):
+        super().__init__()
+        self.proxy = proxy
+        self._buffer = bytearray()
+        self._target: asyncio.Transport | None = None
+        self.transport: asyncio.Transport | None = None
+        self._cached: List[bytes] = []
+        logger.debug(f"__init__  {self}")
+
+    def connection_made(self, transport):
+        logger.debug(f"connection_made {self}")
+        self.transport = transport
+
+    def data_received(self, data):
+        if self._target is None:
+            self._cached.append(data)
+        else:
+            self._target.write(data)
+
+    def set_target(self, target: asyncio.Transport):
+        logger.debug(f"set_target {self}")
+        self._target = target
+        if self._cached:
+            self._target.writelines(self._cached)
+            self._cached = []
+
+    def eof_received(self):
+        logger.debug(f"eof_received {self}")
+        if self._target is not None:
+            self._target.write_eof()
+
+    def connection_lost(self, exc: Exception | None):
+        logger.debug(f"connection_lost {exc} {self}")
+        if self._target is not None:
+            self._target.close()
+
+
+class ProxyClientConnection:
+    """
+    Holds all of the tasks and details for a client proxy
+    """
+
+    def __init__(
+        self,
+        client_transport: asyncio.Transport,
+        client_protocol: ClientToServerProtocol,
+    ):
+        self.client_transport = client_transport
+        self.client_protocol = client_protocol
+        self.server_transport: asyncio.Transport | None = None
+        self.server_protocol: ServerToClientProtocol | None = None
+        self.task: asyncio.Task | None = None
+
+    def close(self):
+        logger.debug(f"closing {self}")
+        if self.client_transport is not None:
+            self._close_transport(self.client_transport)
+        if self.server_transport is not None:
+            self._close_transport(self.server_transport)
+
+    def _close_transport(self, transport:asyncio.Transport):
+        if transport.is_closing():
+            return
+        if transport.can_write_eof():
+            transport.write_eof()
+        else:
+            transport.close()
+
+class ClientToServerProtocol(BaseProxyProtocol):
+    """
+    Protocol to copy bytes from the unix socket client to the database server
+    """
+
+    def __init__(self, proxy: Proxy):
+        super().__init__(proxy)
+        self._buffer = bytearray()
+        self._target: asyncio.Transport | None = None
+        logger.debug(f"__init__ {self}")
+
+    def connection_made(self, transport):
+        # When a connection is made, open the server connection
+        super().connection_made(transport)
+        self.proxy._handle_client_connection(transport, self)
+
+
+class ServerToClientProtocol(BaseProxyProtocol):
+    """
+    Protocol to copy bytes from the database server to the client socket
+    """
+
+    def __init__(self, proxy: Proxy, cconn: ProxyClientConnection):
+        super().__init__(proxy)
+        self._buffer = bytearray()
+        self._target = cconn.client_transport
+        self._client_protocol = cconn.client_protocol
+        logger.debug(f"__init__ {self}")
+
+    def connection_made(self, transport):
+        super().connection_made(transport)
+        self._client_protocol.set_target(transport)
+
+    def connection_lost(self, exc: Exception | None):
+        super().connection_lost(exc)
+        self.proxy._handle_server_connection_lost()
+
+class ServerConnectionFactory(ABC):
+    """
+    ServerConnectionFactory is an abstract class that provides connections to the service.
+    """
+    @abstractmethod
+    async def connect(self, protocol_fn: Callable[[], asyncio.Protocol]):
+        """
+        Establishes a connection to the server and configures it to use the protocol
+        returned from protocol_fn, with asyncio.EventLoop.create_connection().
+        :param protocol_fn: the protocol function
+        :return: None
+        """
+        pass
+
+class Proxy:
+    """
+    A class to represent a local Unix socket proxy for a Cloud SQL instance.
+    This class manages a Unix socket that listens for incoming connections and
+    proxies them to a Cloud SQL instance.
+    """
+
+    def __init__(
+        self,
+        unix_socket_path: str,
+        server_connection_factory: ServerConnectionFactory,
+        loop: asyncio.AbstractEventLoop,
+    ):
+        """
+        Creates a new Proxy
+        :param unix_socket_path: the path to listen for the proxy connection
+        :param loop: The event loop
+        :param instance_connect: A function that will establish the async connection to the server
+
+        The instance_connect function is an asynchronous function that should set up a new connection.
+        It takes one argument - another function that
+        """
+        self.unix_socket_path = unix_socket_path
+        self.alive = True
+        self._loop = loop
+        self._server: asyncio.AbstractServer | None = None
+        self._client_connections: set[ProxyClientConnection] = set()
+        self._server_connection_factory = server_connection_factory
+
+    async def start(self) -> None:
+        """Starts the Unix socket server."""
+        if os.path.exists(self.unix_socket_path):
+            os.remove(self.unix_socket_path)
+
+        parent_dir = Path(self.unix_socket_path).parent
+        parent_dir.mkdir(parents=True, exist_ok=True)
+
+        def new_protocol() -> ClientToServerProtocol:
+            return ClientToServerProtocol(self)
+
+        logger.debug(f"Socket path: {self.unix_socket_path}")
+        self._server = await self._loop.create_unix_server(
+            new_protocol, path=self.unix_socket_path
+        )
+        self._loop.create_task(self._server.serve_forever())
+
+    def _handle_client_connection(
+        self,
+        client_transport: asyncio.Transport,
+        client_protocol: ClientToServerProtocol,
+    ) -> None:
+        """
+        Register a new client connection and initiate the task to create a database connection.
+        This is called by ClientToServerProtocol.connection_made
+
+        :param client_transport: the client transport for the client unix socket
+        :param client_protocol:  the instance for the
+        :return: None
+        """
+        conn = ProxyClientConnection(client_transport, client_protocol)
+        self._client_connections.add(conn)
+        conn.task = self._loop.create_task(self._create_db_instance_connection(conn))
+        conn.task.add_done_callback(lambda _: self._client_connections.discard(conn))
+
+    def _handle_server_connection_lost(
+        self,
+    ) -> None:
+        """
+        Closes the proxy server if the connection to the server is lost
+
+        :return: None
+        """
+        logger.debug(f"Closing proxy server due to lost connection")
+        self._loop.create_task(self.close())
+
+    async def _create_db_instance_connection(self, conn: ProxyClientConnection) -> None:
+        """
+        Manages a single proxy connection from a client to the Cloud SQL instance.
+        """
+        try:
+            logger.debug("_proxy_connection() started")
+            new_protocol = partial(ServerToClientProtocol, self, conn)
+
+            # Establish connection to the database
+            await self._server_connection_factory.connect(new_protocol)
+            logger.debug("_proxy_connection() succeeded")
+
+        except Exception as e:
+            logger.error(f"Error handling proxy connection: {e}")
+            await self.close()
+            raise e
+
+    async def close(self) -> None:
+        """
+        Shuts down the proxy server and cleans up resources.
+        """
+        logger.info(f"Closing Unix socket proxy at {self.unix_socket_path}")
+
+        if self._server:
+            self._server.close()
+            await self._server.wait_closed()
+
+        if self._client_connections:
+            for conn in list(self._client_connections):
+                conn.close()
+            await asyncio.wait([c.task for c in self._client_connections if c.task is not None], timeout=0.1)
+
+        if os.path.exists(self.unix_socket_path):
+            os.remove(self.unix_socket_path)
+
+        logger.info(f"Unix socket proxy for {self.unix_socket_path} closed.")
+        self.alive = False

pyproject.toml

@@ -82,6 +82,10 @@ exclude = ['docs/*', 'samples/*']
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
+log_cli = true
+log_cli_level = "DEBUG"
+log_cli_format = "%(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s)"
+log_cli_date_format = "%Y-%m-%d %H:%M:%S.%f"
 
 [tool.ruff.lint]
 extend-select = ["I"]

requirements-test.txt

@@ -7,6 +7,7 @@ sqlalchemy-pytds==1.0.2
 sqlalchemy-stubs==0.4
 PyMySQL==1.1.1
 pg8000==1.31.4
+psycopg[binary]==3.2.9
 asyncpg==0.30.0
 python-tds==1.16.1
 aioresponses==0.7.8