diff --git a/README.md b/README.md index d6921425..e158daf0 100644 --- a/README.md +++ b/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) @@ -56,12 +57,16 @@ based on your database dialect. pip install "cloud-sql-python-connector[pymysql]" ``` ### Postgres -There are two different database drivers that are supported for the Postgres dialect: +There are three different database drivers that are supported for the Postgres dialect: #### pg8000 ``` pip install "cloud-sql-python-connector[pg8000]" ``` +#### psycopg +``` +pip install "cloud-sql-python-connector[psycopg]" +``` #### asyncpg ``` pip install "cloud-sql-python-connector[asyncpg]" @@ -137,6 +142,18 @@ pool = sqlalchemy.create_engine( db="my-db-name" ), ) + +# Or with Postgres (psycopg): +pool = sqlalchemy.create_engine( + "postgresql+psycopg://", + creator=lambda: connector.connect( + "project:region:instance", + "psycopg", + user="my-user", + password="my-password", + db="my-db-name" + ), +) ``` The returned connection pool engine can then be used to query and modify the database. diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py index 3a1df0ea..7a9964ca 100644 --- a/google/cloud/sql/connector/connector.py +++ b/google/cloud/sql/connector/connector.py @@ -31,6 +31,7 @@ from google.cloud.sql.connector import asyncpg from google.cloud.sql.connector import pg8000 +from google.cloud.sql.connector import psycopg from google.cloud.sql.connector import pymysql from google.cloud.sql.connector import pytds from google.cloud.sql.connector.client import CloudSQLClient @@ -362,6 +363,7 @@ async def connect_async( "pg8000": pg8000.connect, "asyncpg": asyncpg.connect, "pytds": pytds.connect, + "psycopg": psycopg.connect, } # only accept supported database drivers diff --git a/google/cloud/sql/connector/enums.py b/google/cloud/sql/connector/enums.py index 88b5bf47..4bfb0a44 100644 --- a/google/cloud/sql/connector/enums.py +++ b/google/cloud/sql/connector/enums.py @@ -62,6 +62,7 @@ class DriverMapping(Enum): ASYNCPG = "POSTGRES" PG8000 = "POSTGRES" # noqa: PIE796 + PSYCOPG = "POSTGRES" # noqa: PIE796 PYMYSQL = "MYSQL" PYTDS = "SQLSERVER" diff --git a/google/cloud/sql/connector/psycopg.py b/google/cloud/sql/connector/psycopg.py new file mode 100644 index 00000000..e5942951 --- /dev/null +++ b/google/cloud/sql/connector/psycopg.py @@ -0,0 +1,230 @@ +# Copyright 2026 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 +# +# http://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 logging +import os +import selectors +import socket +import ssl +import tempfile +import threading +from typing import Any, TYPE_CHECKING + +if TYPE_CHECKING: + import psycopg + +logger = logging.getLogger(name=__name__) + + +def _proxy(local: socket.socket, remote: "ssl.SSLSocket") -> None: + """Single-threaded selectors-based proxy to avoid SSLSocket thread-safety issues.""" + sel = selectors.DefaultSelector() + sel.register(local, selectors.EVENT_READ, data="local") + sel.register(remote, selectors.EVENT_READ, data="remote") + + def forward_pending() -> bool: + """Read any pending decrypted data from SSL buffer and forward it. + Returns True if EOF was reached or error occurred (should exit). + """ + if not hasattr(remote, "pending"): + return False + pending_bytes = remote.pending() + if not isinstance(pending_bytes, int): + return False + + while pending_bytes > 0: + try: + data = remote.recv(8192) + except OSError as e: + logger.debug("psycopg proxy: remote recv pending error: %s", e) + return True + if not data: + logger.debug("psycopg proxy: remote pending EOF") + return True + try: + local.sendall(data) + except OSError as e: + logger.debug("psycopg proxy: local send pending error: %s", e) + return True + try: + pending_bytes = remote.pending() + except OSError: + break + if not isinstance(pending_bytes, int): + break + return False + + try: + while True: + # First check if there is any pending data in SSL buffer + if forward_pending(): + break + + events = sel.select() + + for key, mask in events: + if key.data == "local": + try: + data = local.recv(8192) + except OSError as e: + logger.debug("psycopg proxy: local recv error: %s", e) + return + if not data: + logger.debug("psycopg proxy: local EOF") + return + try: + remote.sendall(data) + except OSError as e: + logger.debug("psycopg proxy: remote send error: %s", e) + return + elif key.data == "remote": + try: + data = remote.recv(8192) + except OSError as e: + logger.debug("psycopg proxy: remote recv error: %s", e) + return + if not data: + logger.debug("psycopg proxy: remote EOF") + return + try: + local.sendall(data) + except OSError as e: + logger.debug("psycopg proxy: local send error: %s", e) + return + except OSError as e: + logger.debug("psycopg proxy: OSError in loop: %s", e) + finally: + sel.close() + for s in (local, remote): + try: + s.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + s.close() + except OSError: + pass + + +def connect( + ip_address: str, remote_sock: "ssl.SSLSocket", **kwargs: Any +) -> "psycopg.Connection": + """Create a psycopg DBAPI connection object. + + Because psycopg does not accept a pre-connected socket, this function + creates a temporary Unix domain socket, tells psycopg to connect there, + and runs a background proxy that forwards bytes between that socket and + the already-established Cloud SQL TLS connection. + + Args: + ip_address (str): IP address of the Cloud SQL instance. + remote_sock (ssl.SSLSocket): SSL/TLS secure socket stream connected to the + Cloud SQL proxy server. + + Returns: + psycopg.Connection: A psycopg Connection object for the Cloud SQL instance. + """ + try: + import psycopg + except ImportError: + raise ImportError( + 'Unable to import module "psycopg." Please install and try again.' + ) + + tmpdir = tempfile.mkdtemp() + socket_path = os.path.join(tmpdir, ".s.PGSQL.5432") + logger.debug("psycopg: created Unix socket at %s", socket_path) + + local_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + local_sock.bind(socket_path) + local_sock.listen(1) + + def _accept_and_proxy() -> None: + """Accept one connection then proxy bytes until the connection closes.""" + unix_conn = None + try: + unix_conn, _ = local_sock.accept() + local_sock.close() + logger.debug("psycopg proxy: accepted connection, starting proxy") + _proxy(unix_conn, remote_sock) + except Exception as e: # noqa: BLE001 + logger.debug("psycopg proxy: error in accept/proxy thread: %s", e) + # Ensure cleanup on any exception + if unix_conn: + try: + unix_conn.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + unix_conn.close() + except OSError: + pass + try: + remote_sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + remote_sock.close() + except OSError: + pass + + threading.Thread(target=_accept_and_proxy, daemon=True).start() + + user = kwargs.pop("user") + db = kwargs.pop("db") + passwd = kwargs.pop("password", None) + # SSL is already handled by the underlying SSLSocket; disable it on the + # Unix socket so psycopg does not attempt a second TLS handshake. + kwargs.pop("sslmode", None) + timeout = kwargs.pop("timeout", None) + if timeout is not None: + kwargs["connect_timeout"] = int(timeout) + + logger.debug("psycopg: connecting as user=%s dbname=%s", user, db) + try: + conn = psycopg.connect( + user=user, + dbname=db, + password=passwd, + host=tmpdir, + port=5432, + sslmode="disable", + **kwargs, + ) + logger.debug("psycopg: connection established") + return conn + except Exception as e: + logger.debug("psycopg: connection failed: %s", e) + # psycopg never connected (or failed mid-handshake); close the server + # socket so the proxy thread unblocks and exits cleanly. + try: + local_sock.close() + except OSError: + pass + try: + remote_sock.close() + except OSError: + pass + raise + finally: + # The socket file and its parent directory are only needed during the + # initial connect() call; remove them now regardless of outcome. + try: + os.remove(socket_path) + except OSError: + pass + try: + os.rmdir(tmpdir) + except OSError: + pass diff --git a/pyproject.toml b/pyproject.toml index dcff67ae..93ed5eaf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,6 +63,7 @@ pymysql = ["PyMySQL>=1.1.0"] pg8000 = ["pg8000>=1.31.1"] pytds = ["python-tds>=1.15.0"] asyncpg = ["asyncpg>=0.30.0"] +psycopg = ["psycopg>=3.1.0"] [tool.setuptools.dynamic] version = { attr = "google.cloud.sql.connector.version.__version__" } diff --git a/requirements-test.txt b/requirements-test.txt index fcbb1bc4..ca0490e7 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -11,4 +11,6 @@ asyncpg==0.31.0 python-tds==1.17.1 aioresponses==0.7.9 pytest-aiohttp==1.1.1 +psycopg==3.3.4 +psycopg-binary==3.3.4 diff --git a/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py new file mode 100644 index 00000000..9424d86d --- /dev/null +++ b/tests/system/test_psycopg_connection.py @@ -0,0 +1,212 @@ +""" +Copyright 2026 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 + +import asyncio +from datetime import datetime +import os + +import pytest +import sqlalchemy + +from google.cloud.sql.connector import Connector +from google.cloud.sql.connector import DefaultResolver +from google.cloud.sql.connector import DnsResolver + + +def create_sqlalchemy_engine( + instance_connection_name: str, + user: str, + password: str, + db: str, + ip_type: str = "public", + refresh_strategy: str = "background", + resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, +) -> tuple[sqlalchemy.engine.Engine, Connector]: + """Creates a connection pool for a Cloud SQL instance and returns the pool + and the connector. + """ + connector = Connector(refresh_strategy=refresh_strategy, resolver=resolver) + + # create SQLAlchemy connection pool + engine = sqlalchemy.create_engine( + "postgresql+psycopg://", + creator=lambda: connector.connect( + instance_connection_name, + "psycopg", + user=user, + password=password, + db=db, + ip_type=ip_type, + ), + ) + return engine, connector + + +def test_psycopg_connection() -> None: + """Basic test to get time from database using psycopg.""" + inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"] + user = os.environ["POSTGRES_USER"] + password = os.environ["POSTGRES_PASS"] + db = os.environ["POSTGRES_DB"] + ip_type = os.environ.get("IP_TYPE", "public") + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, user, password, db, ip_type + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + +def test_lazy_psycopg_connection() -> None: + """Basic test to get time from database using psycopg and lazy refresh.""" + inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"] + user = os.environ["POSTGRES_USER"] + password = os.environ["POSTGRES_PASS"] + db = os.environ["POSTGRES_DB"] + ip_type = os.environ.get("IP_TYPE", "public") + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, user, password, db, ip_type, "lazy" + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + +def test_CAS_psycopg_connection() -> None: + """Basic test to get time from database using CAS.""" + inst_conn_name = os.environ.get("POSTGRES_CAS_CONNECTION_NAME") + user = os.environ["POSTGRES_USER"] + password = os.environ.get("POSTGRES_CAS_PASS") + db = os.environ["POSTGRES_DB"] + ip_type = os.environ.get("IP_TYPE", "public") + + if not inst_conn_name or not password: + pytest.skip("POSTGRES_CAS_CONNECTION_NAME or POSTGRES_CAS_PASS not set") + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, user, password, db, ip_type + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + +def test_customer_managed_CAS_psycopg_connection() -> None: + """Basic test to get time from database using Customer Managed CAS.""" + inst_conn_name = os.environ.get("POSTGRES_CUSTOMER_CAS_CONNECTION_NAME") + user = os.environ["POSTGRES_USER"] + password = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS") + db = os.environ["POSTGRES_DB"] + ip_type = os.environ.get("IP_TYPE", "public") + + if not inst_conn_name or not password: + pytest.skip("POSTGRES_CUSTOMER_CAS_CONNECTION_NAME or POSTGRES_CUSTOMER_CAS_PASS not set") + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, user, password, db, ip_type + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + +def test_custom_SAN_with_dns_psycopg_connection() -> None: + """Basic test to get time from database using Custom SAN with DNS.""" + inst_conn_name = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME") + user = os.environ["POSTGRES_USER"] + password = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS") + db = os.environ["POSTGRES_DB"] + ip_type = os.environ.get("IP_TYPE", "public") + + if not inst_conn_name or not password: + pytest.skip("POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME or POSTGRES_CUSTOMER_CAS_PASS not set") + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, user, password, db, ip_type, resolver=DnsResolver + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + +def test_MCP_psycopg_connection() -> None: + """Basic test to get time from database using MCP enabled instance.""" + inst_conn_name = os.environ.get("POSTGRES_MCP_CONNECTION_NAME") + user = os.environ["POSTGRES_USER"] + password = os.environ.get("POSTGRES_MCP_PASS") + db = os.environ["POSTGRES_DB"] + ip_type = os.environ.get("IP_TYPE", "public") + + if not inst_conn_name or not password: + pytest.skip("POSTGRES_MCP_CONNECTION_NAME or POSTGRES_MCP_PASS not set") + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, user, password, db, ip_type + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + +def test_system_psycopg_to_thread() -> None: + """Verify that running sync connect in asyncio.to_thread works.""" + inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"] + user = os.environ["POSTGRES_USER"] + password = os.environ["POSTGRES_PASS"] + db = os.environ["POSTGRES_DB"] + ip_type = os.environ.get("IP_TYPE", "public") + + async def run_connect(): + with Connector() as connector: + # Run the blocking connector.connect in a thread + conn = await asyncio.to_thread( + connector.connect, + inst_conn_name, + "psycopg", + user=user, + password=password, + db=db, + ip_type=ip_type, + ) + cursor = conn.cursor() + cursor.execute("SELECT NOW();") + result = cursor.fetchone() + assert result is not None + cursor.close() + conn.close() + + asyncio.run(run_connect()) diff --git a/tests/system/test_psycopg_iam_auth.py b/tests/system/test_psycopg_iam_auth.py new file mode 100644 index 00000000..ebb034db --- /dev/null +++ b/tests/system/test_psycopg_iam_auth.py @@ -0,0 +1,90 @@ +""" +Copyright 2026 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 datetime import datetime +import os + +import pytest +import sqlalchemy + +from google.cloud.sql.connector import Connector + +# Skip all tests in this file if POSTGRES_IAM_USER is not set +pytestmark = pytest.mark.skipif( + not os.environ.get("POSTGRES_IAM_USER"), + reason="POSTGRES_IAM_USER env var not set for IAM Authn tests", +) + + +def create_sqlalchemy_engine( + instance_connection_name: str, + user: str, + db: str, + ip_type: str = "public", + refresh_strategy: str = "background", +) -> tuple[sqlalchemy.engine.Engine, Connector]: + """Creates a connection pool for a Cloud SQL instance and returns the pool + and the connector. + """ + connector = Connector(refresh_strategy=refresh_strategy) + + # create SQLAlchemy connection pool + engine = sqlalchemy.create_engine( + "postgresql+psycopg://", + creator=lambda: connector.connect( + instance_connection_name, + "psycopg", + user=user, + db=db, + ip_type=ip_type, + enable_iam_auth=True, + ), + ) + return engine, connector + + +def test_psycopg_iam_authn_connection() -> None: + """Basic test to get time from database using psycopg and IAM Authn.""" + inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"] + user = os.environ["POSTGRES_IAM_USER"] + db = os.environ["POSTGRES_DB"] + ip_type = os.getenv("IP_TYPE", "public") + + engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, ip_type) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + +def test_lazy_psycopg_iam_authn_connection() -> None: + """Basic test to get time from database using psycopg, IAM Authn and lazy refresh.""" + inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"] + user = os.environ["POSTGRES_IAM_USER"] + db = os.environ["POSTGRES_DB"] + ip_type = os.getenv("IP_TYPE", "public") + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, user, db, ip_type, refresh_strategy="lazy" + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index edc2bf3e..9b621a69 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -1032,5 +1032,37 @@ async def test_Connector_connect_async_connection_error_triggers_force_refresh( mock_force_refresh.assert_called_once() +@pytest.mark.asyncio +async def test_Connector_connect_async_psycopg( + fake_credentials: Credentials, fake_client: CloudSQLClient +) -> None: + """Test that Connector.connect_async works with psycopg driver.""" + async with Connector( + credentials=fake_credentials, loop=asyncio.get_running_loop() + ) as connector: + connector._client = fake_client + + mock_sock = MagicMock() + with ( + patch( + "google.cloud.sql.connector.connector.socket.create_connection", + return_value=mock_sock, + ), + patch("ssl.SSLContext.wrap_socket", return_value=mock_sock), + patch("google.cloud.sql.connector.psycopg.connect") as mock_connect, + ): + mock_connect.return_value = True + + connection = await connector.connect_async( + "test-project:test-region:test-instance", + "psycopg", + user="my-user", + password="my-pass", + db="my-db", + ) + assert connection is True + mock_connect.assert_called_once() + + diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py new file mode 100644 index 00000000..6ed51b7e --- /dev/null +++ b/tests/unit/test_psycopg.py @@ -0,0 +1,733 @@ +# Copyright 2026 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 +# +# http://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 os +import socket +import ssl +import threading +from typing import Any +from unittest.mock import MagicMock +from unittest.mock import patch + +import pytest + +from google.cloud.sql.connector.psycopg import _proxy +from google.cloud.sql.connector.psycopg import connect + +pytestmark = pytest.mark.skipif( + not hasattr(socket, "AF_UNIX"), + reason="Unix domain sockets (AF_UNIX) not available on this platform", +) + + +class MockableSocket(socket.socket): + pass + + +def mockable_socketpair() -> tuple[MockableSocket, MockableSocket]: + """Create a socketpair wrapped in MockableSocket to allow method mocking.""" + s1, s2 = socket.socketpair() + fd1 = s1.detach() + fd2 = s2.detach() + ms1 = MockableSocket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd1) + ms2 = MockableSocket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd2) + return ms1, ms2 + + +def test_proxy_bidirectional() -> None: + """Test that _proxy forwards bytes in both directions and exits on EOF.""" + # local_client <-> local_server (simulates psycopg <-> proxy) + local_client, local_server = socket.socketpair() + # remote_client <-> remote_server (simulates proxy <-> Cloud SQL) + remote_client, remote_server = socket.socketpair() + + # Start proxy in a background thread because it blocks + proxy_thread = threading.Thread( + target=_proxy, args=(local_server, remote_client), daemon=True + ) + proxy_thread.start() + + # Test local -> remote + local_client.sendall(b"hello from local") + assert remote_server.recv(1024) == b"hello from local" + + # Test remote -> local + remote_server.sendall(b"hello from remote") + assert local_client.recv(1024) == b"hello from remote" + + # Close local client (EOF) + local_client.close() + + # Wait for proxy thread to finish + proxy_thread.join(timeout=2.0) + assert not proxy_thread.is_alive() + + # Verify remote socket was also closed by proxy + try: + data = remote_server.recv(1024) + assert data == b"" + except OSError: + pass # Closed socket error is also acceptable + + # Clean up remaining sockets + local_server.close() + remote_client.close() + remote_server.close() + + +def test_proxy_pending_data() -> None: + """Test that _proxy forwards pending SSL data correctly.""" + local_client, local_server = socket.socketpair() + remote_client, remote_server = socket.socketpair() + + class MockSSLSocket: + def __init__(self, sock: socket.socket) -> None: + self._sock = sock + self._pending_calls = [12, 0] # "pending data" is 12 bytes + + def pending(self) -> int: + if self._pending_calls: + return self._pending_calls.pop(0) + return 0 + + def recv(self, bufsize: int, flags: int = 0) -> bytes: + return self._sock.recv(bufsize, flags) + + def __getattr__(self, name: str) -> Any: + return getattr(self._sock, name) + + wrapped_remote = MockSSLSocket(remote_client) + + # Pre-populate the socket with data that will be read by forward_pending + remote_server.sendall(b"pending data") + + # Start proxy in background + proxy_thread = threading.Thread( + target=_proxy, args=(local_server, wrapped_remote), daemon=True + ) + proxy_thread.start() + + # Verify that local_client receives the pending data immediately + assert local_client.recv(1024) == b"pending data" + + # Clean up + local_client.close() + proxy_thread.join(timeout=2.0) + + local_server.close() + remote_client.close() + remote_server.close() + + +@patch("psycopg.connect") +def test_connect_wrapper(mock_psycopg_connect: MagicMock) -> None: + """Test connect wrapper creates temp socket and calls psycopg.connect with correct arguments.""" + mock_remote_sock = MagicMock(spec=ssl.SSLSocket) + + # We need to mock psycopg.connect to simulate a connection. + # To prevent the accept thread in connect() from hanging, we make the mock + # connect to the Unix socket before returning. + def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock: + host = kwargs.get("host") + socket_path = os.path.join(host, ".s.PGSQL.5432") + client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + client.connect(socket_path) + client.close() + return MagicMock() + + mock_psycopg_connect.side_effect = mock_connect_impl + + # Call the connect wrapper + conn = connect( + "127.0.0.1", + mock_remote_sock, + user="test_user", + db="test_db", + password="test_password", + sslmode="require", + timeout=30.5, + ) + + assert conn is not None + assert mock_psycopg_connect.called + + # Verify arguments passed to psycopg.connect + _, kwargs = mock_psycopg_connect.call_args + assert kwargs["user"] == "test_user" + assert kwargs["dbname"] == "test_db" + assert kwargs["password"] == "test_password" + assert kwargs["sslmode"] == "disable" + assert kwargs["connect_timeout"] == 30 + assert "timeout" not in kwargs + assert "host" in kwargs + assert kwargs["port"] == 5432 + + # Verify temp dir was cleaned up + assert not os.path.exists(kwargs["host"]) + + +@patch("psycopg.connect") +def test_connect_wrapper_failure(mock_psycopg_connect: MagicMock) -> None: + """Test that connect wrapper cleans up correctly when psycopg.connect fails.""" + mock_remote_sock = MagicMock(spec=ssl.SSLSocket) + mock_psycopg_connect.side_effect = Exception("connection failed simulated") + + # Call the connect wrapper and expect it to raise + with pytest.raises(Exception, match="connection failed simulated"): + connect( + "127.0.0.1", + mock_remote_sock, + user="test_user", + db="test_db", + password="test_password", + ) + + # Verify remote socket was closed + assert mock_remote_sock.close.called + + # Verify cleanup with mocked paths + with patch("google.cloud.sql.connector.psycopg.tempfile.mkdtemp") as mock_mkdtemp: + mock_mkdtemp.return_value = "/tmp/mock_temp_dir_failure" + + with ( + patch("google.cloud.sql.connector.psycopg.os.rmdir") as mock_rmdir, + patch("google.cloud.sql.connector.psycopg.os.remove") as mock_remove, + patch( + "google.cloud.sql.connector.psycopg.socket.socket" + ) as mock_socket_cls, + ): + # Mock the local socket to avoid real OS bind/listen + mock_local_sock = MagicMock() + mock_socket_cls.return_value = mock_local_sock + + with pytest.raises(Exception, match="connection failed simulated"): + connect( + "127.0.0.1", + mock_remote_sock, + user="test_user", + db="test_db", + password="test_password", + ) + + # Verify rmdir and remove were called for cleanup + mock_rmdir.assert_called_once_with("/tmp/mock_temp_dir_failure") + mock_remove.assert_called_once() + + +def test_proxy_remote_eof() -> None: + """Test that _proxy exits when remote socket receives EOF.""" + local_client, local_server = socket.socketpair() + remote_client, remote_server = socket.socketpair() + + # Start proxy in background + proxy_thread = threading.Thread( + target=_proxy, args=(local_server, remote_client), daemon=True + ) + proxy_thread.start() + + # Close remote server to trigger EOF on remote_client + remote_server.close() + + # local_client should receive EOF (b"") + assert local_client.recv(1024) == b"" + + # Wait for proxy thread to finish + proxy_thread.join(timeout=2.0) + assert not proxy_thread.is_alive() + + # Sockets should be closed + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + + +def test_proxy_local_recv_error() -> None: + """Test that _proxy exits when local.recv raises OSError.""" + local_client, local_server = mockable_socketpair() + remote_client, remote_server = mockable_socketpair() + + # Mock local_server.recv to raise OSError + local_server.recv = MagicMock(side_effect=OSError("local recv failed")) + + # Trigger selector by sending data to local_server + local_client.send(b"x") + + # Run proxy. It should detect local is readable, call local.recv() which raises OSError, and exit. + _proxy(local_server, remote_client) + + # Sockets should be closed + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + remote_server.close() + + +def test_proxy_remote_send_error() -> None: + """Test that _proxy exits when remote.sendall raises OSError.""" + local_client, local_server = mockable_socketpair() + remote_client, remote_server = mockable_socketpair() + + # Mock remote_client.sendall to raise OSError + remote_client.sendall = MagicMock(side_effect=OSError("remote send failed")) + + # Trigger selector by sending data from local_client -> local_server + local_client.send(b"x") + + # Run proxy. It reads "x" from local, tries to send to remote, fails, and exits. + _proxy(local_server, remote_client) + + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + remote_server.close() + + +def test_proxy_remote_recv_error() -> None: + """Test that _proxy exits when remote.recv raises OSError.""" + local_client, local_server = mockable_socketpair() + remote_client, remote_server = mockable_socketpair() + + # Mock remote_client.recv to raise OSError + remote_client.recv = MagicMock(side_effect=OSError("remote recv failed")) + + # Trigger selector by sending data from remote_server -> remote_client + remote_server.send(b"x") + + # Run proxy. It detects remote is readable, calls remote.recv() which fails, and exits. + _proxy(local_server, remote_client) + + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + remote_server.close() + + +def test_proxy_local_send_error() -> None: + """Test that _proxy exits when local.sendall raises OSError.""" + local_client, local_server = mockable_socketpair() + remote_client, remote_server = mockable_socketpair() + + # Mock local_server.sendall to raise OSError + local_server.sendall = MagicMock(side_effect=OSError("local send failed")) + + # Trigger selector by sending data from remote_server -> remote_client + remote_server.send(b"x") + + # Run proxy. It reads "x" from remote, tries to send to local, fails, and exits. + _proxy(local_server, remote_client) + + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + remote_server.close() + + +def test_proxy_pending_recv_error() -> None: + """Test that _proxy exits when remote.recv raises OSError during pending check.""" + local_client, local_server = socket.socketpair() + remote_client, remote_server = mockable_socketpair() + + # Mock remote_client (remote sock in proxy) + remote_client.pending = MagicMock(return_value=10) + remote_client.recv = MagicMock(side_effect=OSError("pending recv failed")) + + # Run proxy + _proxy(local_server, remote_client) + + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + remote_server.close() + + +def test_proxy_pending_local_send_error() -> None: + """Test that _proxy exits when local.sendall raises OSError during pending check.""" + local_client, local_server = mockable_socketpair() + remote_client, remote_server = mockable_socketpair() + + # Mock remote_client (remote sock in proxy) + remote_client.pending = MagicMock(return_value=10) + remote_client.recv = MagicMock(return_value=b"pending data") + + # Mock local_server.sendall to raise OSError + local_server.sendall = MagicMock(side_effect=OSError("local send pending failed")) + + # Run proxy + _proxy(local_server, remote_client) + + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + remote_server.close() + + +@patch.dict("sys.modules", {"psycopg": None}) +def test_connect_import_error() -> None: + """Test that connect raises ImportError if psycopg is not installed.""" + mock_remote_sock = MagicMock(spec=ssl.SSLSocket) + + with pytest.raises(ImportError, match='Unable to import module "psycopg."'): + connect("127.0.0.1", mock_remote_sock) + + +def test_connect_cleanup_errors() -> None: + """Test that connect ignores OSErrors when removing temp files/dirs during cleanup.""" + mock_remote_sock = MagicMock(spec=ssl.SSLSocket) + + def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock: + host = kwargs.get("host") + socket_path = os.path.join(host, ".s.PGSQL.5432") + client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + client.connect(socket_path) + client.close() + return MagicMock() + + with ( + patch("psycopg.connect", side_effect=mock_connect_impl), + patch( + "google.cloud.sql.connector.psycopg.os.remove", + side_effect=OSError("remove failed"), + ) as mock_remove, + patch( + "google.cloud.sql.connector.psycopg.os.rmdir", + side_effect=OSError("rmdir failed"), + ) as mock_rmdir, + ): + conn = connect( + "127.0.0.1", + mock_remote_sock, + user="test_user", + db="test_db", + password="test_password", + ) + + assert conn is not None + assert mock_remove.called + assert mock_rmdir.called + + +def test_proxy_happy_path_sequential() -> None: + """Test that _proxy forwards sequential request/response traffic cleanly without errors.""" + local_client, local_proxy = socket.socketpair( + socket.AF_UNIX, socket.SOCK_STREAM + ) + remote_proxy, remote_server = socket.socketpair( + socket.AF_UNIX, socket.SOCK_STREAM + ) + + t = threading.Thread( + target=_proxy, args=(local_proxy, remote_proxy), daemon=True + ) + t.start() + + # Step 1: Client sends query + local_client.sendall(b"SELECT 1;") + received_query = remote_server.recv(1024) + assert received_query == b"SELECT 1;" + + # Step 2: Server sends response + remote_server.sendall(b"RESULT 1") + received_resp = local_client.recv(1024) + assert received_resp == b"RESULT 1" + + # Step 3: Client closes connection cleanly + local_client.shutdown(socket.SHUT_RDWR) + local_client.close() + remote_server.close() + t.join(timeout=2.0) + assert not t.is_alive() + + +def test_proxy_backpressure_and_clean_teardown() -> None: + """Demonstrate that Cloud SQL's single-threaded selector proxy exits cleanly + + without deadlocking on teardown. + """ + local_client, local_proxy = socket.socketpair( + socket.AF_UNIX, socket.SOCK_STREAM + ) + remote_proxy, remote_server = socket.socketpair( + socket.AF_UNIX, socket.SOCK_STREAM + ) + + t_proxy = threading.Thread( + target=_proxy, args=(local_proxy, remote_proxy), daemon=True + ) + t_proxy.start() + + # Step 1: Client sends query + local_client.sendall(b"SELECT 1;") + assert remote_server.recv(1024) == b"SELECT 1;" + + # Step 2: Server sends response + remote_server.sendall(b"RESULT 1") + assert local_client.recv(1024) == b"RESULT 1" + + # Step 3: Client closes connection while server is open + local_client.close() + + # Single-threaded proxy terminates immediately without deadlocking + t_proxy.join(timeout=2.0) + assert ( + not t_proxy.is_alive() + ), "Cloud SQL proxy must terminate cleanly without deadlocks" + + remote_server.close() + + +def test_proxy_pending_oserror_in_loop() -> None: + """Test that _proxy pending loop handles OSError on pending() and continues.""" + local_client, local_proxy = mockable_socketpair() + remote_proxy, remote_server = mockable_socketpair() + + # Add pending mock to real socket. First call returns 10, second raises OSError. + remote_proxy.pending = MagicMock(side_effect=[10, OSError("pending failed")]) + + # Send data to be read by recv in pending loop + remote_server.sendall(b"pending data") + + # Start proxy in thread because it will block on select() after pending fails + t = threading.Thread(target=_proxy, args=(local_proxy, remote_proxy), daemon=True) + t.start() + + # Verify local_client received the pending data + assert local_client.recv(1024) == b"pending data" + + # Now trigger proxy exit by closing local_client + local_client.close() + + t.join(timeout=2.0) + assert not t.is_alive() + + remote_server.close() + + +def test_proxy_pending_non_int() -> None: + """Test that _proxy pending check handles non-int return values from pending().""" + local_client, local_proxy = mockable_socketpair() + remote_proxy, remote_server = mockable_socketpair() + + # First call returns non-int. Should return False (goes to select) + remote_proxy.pending = MagicMock(return_value="not an int") + + # We need to run it in a thread because it will block on select + t = threading.Thread(target=_proxy, args=(local_proxy, remote_proxy), daemon=True) + t.start() + + # Trigger exit. + local_client.close() + t.join(timeout=2.0) + assert not t.is_alive() + remote_server.close() + + +def test_proxy_pending_non_int_in_loop() -> None: + """Test that _proxy pending loop handles non-int return values from pending() in loop.""" + local_client, local_proxy = mockable_socketpair() + remote_proxy, remote_server = mockable_socketpair() + + # Second call returns non-int. + remote_proxy.pending = MagicMock(side_effect=[10, "not an int"]) + remote_server.sendall(b"data") + + t = threading.Thread(target=_proxy, args=(local_proxy, remote_proxy), daemon=True) + t.start() + + assert local_client.recv(1024) == b"data" + + local_client.close() + t.join(timeout=2.0) + assert not t.is_alive() + remote_server.close() + + +def test_proxy_pending_recv_eof() -> None: + """Test that _proxy pending loop handles remote recv EOF.""" + local_client, local_server = mockable_socketpair() + remote_client, remote_server = mockable_socketpair() + + remote_client.pending = MagicMock(return_value=10) + # Mock recv to return EOF + remote_client.recv = MagicMock(return_value=b"") + + # Run proxy. It should call forward_pending, which calls recv, gets EOF, + # and returns True. This breaks the loop and proxy exits. + _proxy(local_server, remote_client) + + # Sockets should be closed + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + remote_server.close() + + +def test_proxy_select_oserror() -> None: + """Test that _proxy loop handles OSError in select().""" + local_client, local_server = mockable_socketpair() + remote_client, remote_server = mockable_socketpair() + + with patch( + "google.cloud.sql.connector.psycopg.selectors.DefaultSelector" + ) as mock_sel_cls: + mock_sel = MagicMock() + mock_sel.select.side_effect = OSError("select failed") + mock_sel_cls.return_value = mock_sel + + # Run proxy. It should catch OSError on select() and exit loop to finally block. + _proxy(local_server, remote_client) + + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + local_client.close() + remote_server.close() + + +def test_proxy_finally_cleanup_errors() -> None: + """Test that _proxy finally block ignores OSErrors during socket shutdown/close.""" + local_client, local_server = mockable_socketpair() + remote_client, remote_server = mockable_socketpair() + + real_local_shutdown = local_server.shutdown + def mock_local_shutdown(how): + try: + real_local_shutdown(how) + except OSError: + pass + raise OSError("shutdown failed") + local_server.shutdown = MagicMock(side_effect=mock_local_shutdown) + + real_local_close = local_server.close + def mock_local_close(): + real_local_close() + raise OSError("close failed") + local_server.close = MagicMock(side_effect=mock_local_close) + + real_remote_shutdown = remote_client.shutdown + def mock_remote_shutdown(how): + try: + real_remote_shutdown(how) + except OSError: + pass + raise OSError("shutdown failed") + remote_client.shutdown = MagicMock(side_effect=mock_remote_shutdown) + + real_remote_close = remote_client.close + def mock_remote_close(): + real_remote_close() + raise OSError("close failed") + remote_client.close = MagicMock(side_effect=mock_remote_close) + + # Trigger exit by closing client + local_client.close() + + # Run proxy. It should handle the exceptions in finally block. + _proxy(local_server, remote_client) + + remote_server.close() + + +@patch("google.cloud.sql.connector.psycopg._proxy") +def test_accept_and_proxy_cleanup_errors(mock_proxy_fn: MagicMock) -> None: + """Test that accept_and_proxy thread handles errors during cleanup on exception.""" + # Make _proxy raise an exception to trigger the except block in _accept_and_proxy + mock_proxy_fn.side_effect = Exception("proxy crash") + + mock_unix_conn = MagicMock(spec=socket.socket) + mock_unix_conn.shutdown.side_effect = OSError("unix conn shutdown failed") + mock_unix_conn.close.side_effect = OSError("unix conn close failed") + + mock_remote_sock = MagicMock(spec=ssl.SSLSocket) + mock_remote_sock.shutdown.side_effect = OSError("remote shutdown failed") + + # Mock socket.socket to return a mock local_sock that returns our mock_unix_conn + real_socket = socket.socket + mock_local_sock = MagicMock() + mock_local_sock.accept.return_value = (mock_unix_conn, ("path",)) + + def socket_side_effect(family, type, proto=0, fileno=None): + if family == socket.AF_UNIX: + return mock_local_sock + return real_socket(family, type, proto, fileno) + + event = threading.Event() + def remote_close_fn(): + event.set() + raise OSError("remote close failed") + mock_remote_sock.close.side_effect = remote_close_fn + + with patch("socket.socket", side_effect=socket_side_effect), patch( + "psycopg.connect" + ) as mock_psycopg_connect: + mock_psycopg_connect.return_value = MagicMock() + + connect( + "127.0.0.1", + mock_remote_sock, + user="test_user", + db="test_db", + password="test_password", + ) + + # Wait for the accept_and_proxy thread to finish cleanup + assert event.wait(timeout=2.0) + + # Verify mock_unix_conn shutdown and close were called + mock_unix_conn.shutdown.assert_called_once_with(socket.SHUT_RDWR) + mock_unix_conn.close.assert_called_once() + mock_remote_sock.shutdown.assert_called_once_with(socket.SHUT_RDWR) + mock_remote_sock.close.assert_called_once() + + +@patch("psycopg.connect") +def test_connect_wrapper_failure_cleanup_errors(mock_psycopg_connect: MagicMock) -> None: + """Test that connect wrapper ignores OSErrors during cleanup on connection failure.""" + mock_remote_sock = MagicMock(spec=ssl.SSLSocket) + mock_remote_sock.close.side_effect = OSError("remote close failed") + + mock_psycopg_connect.side_effect = Exception("connection failed simulated") + + mock_local_sock = MagicMock() + mock_local_sock.close.side_effect = OSError("local close failed") + + real_socket = socket.socket + def socket_side_effect(family, type, proto=0, fileno=None): + if family == socket.AF_UNIX: + return mock_local_sock + return real_socket(family, type, proto, fileno) + + with ( + patch("socket.socket", side_effect=socket_side_effect), + pytest.raises(Exception, match="connection failed simulated"), + ): + connect( + "127.0.0.1", + mock_remote_sock, + user="test_user", + db="test_db", + password="test_password", + ) + + mock_local_sock.close.assert_called_once() + assert mock_remote_sock.close.call_count >= 1 + +