From b528fee07def56c8e80f094efa8de0ed16cf29f3 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 16:11:42 +0000 Subject: [PATCH 01/19] feat: add support for psycopg --- google/cloud/sql/connector/connector.py | 2 + google/cloud/sql/connector/enums.py | 1 + google/cloud/sql/connector/psycopg.py | 168 ++++++++++++++++++++++++ tests/system/test_psycopg_connection.py | 147 +++++++++++++++++++++ tests/system/test_psycopg_iam_auth.py | 96 ++++++++++++++ tests/unit/test_psycopg.py | 112 ++++++++++++++++ 6 files changed, 526 insertions(+) create mode 100644 google/cloud/sql/connector/psycopg.py create mode 100644 tests/system/test_psycopg_connection.py create mode 100644 tests/system/test_psycopg_iam_auth.py create mode 100644 tests/unit/test_psycopg.py 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..c930f7b2 --- /dev/null +++ b/google/cloud/sql/connector/psycopg.py @@ -0,0 +1,168 @@ +# 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 socket +import ssl +import tempfile +import threading +from typing import Any, TYPE_CHECKING + +if TYPE_CHECKING: + import psycopg + +logger = logging.getLogger(name=__name__) + +_CHUNK_SIZE = 8 * 1024 # bytes per recv() call inside the proxy forwarding loop + + +def _proxy(local: socket.socket, remote: "ssl.SSLSocket") -> None: + """Bidirectionally proxy bytes between a local Unix socket and a remote + SSL socket. + + Spawns one daemon thread for the remote→local direction and runs the + local→remote direction in the calling thread. Blocks until the calling + thread's direction reaches EOF or a socket error, at which point both + sockets are closed so the other thread also unblocks and exits. + + Args: + local: The Unix domain socket connected to the database driver. + remote: The SSL socket connected to the Cloud SQL proxy server. + """ + def forward(src: Any, dst: Any) -> None: + buf = bytearray(_CHUNK_SIZE) + view = memoryview(buf) + try: + while True: + n = src.recv_into(view) + if n == 0: + logger.debug("psycopg proxy: EOF on %s, closing both sockets", src) + break + dst.sendall(view[:n]) + except (OSError, ssl.SSLError) as e: + logger.debug("psycopg proxy: socket error on %s: %s", src, e) + finally: + # Close both ends so the sibling thread also unblocks. + for s in (local, remote): + try: + s.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + s.close() + except OSError: + pass + + threading.Thread(target=forward, args=(remote, local), daemon=True).start() + forward(local, remote) # run in calling thread rather than spawning a third + + +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.""" + try: + unix_conn, _ = local_sock.accept() + local_sock.close() + logger.debug("psycopg proxy: accepted connection, starting proxy") + except OSError as e: + logger.debug("psycopg proxy: accept failed: %s", e) + try: + remote_sock.close() + except OSError: + pass + return + _proxy(unix_conn, remote_sock) + + 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/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py new file mode 100644 index 00000000..60fdff08 --- /dev/null +++ b/tests/system/test_psycopg_connection.py @@ -0,0 +1,147 @@ +# 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 asyncio +import time +import pytest +import psutil +from google.cloud.sql.connector import Connector + +# These will be set from environment variables or default to our test instance +INSTANCE_CONNECTION_NAME = os.getenv( + "DB_CONNECTION_NAME", "galakp-playground:us-east7:pg-us-east7-psycopg" +) +DB_USER = os.getenv("DB_USER", "postgres") +DB_PASSWORD = os.getenv("DB_PASSWORD", "SuperPass123!") +DB_NAME = os.getenv("DB_NAME", "postgres") + + +def test_system_psycopg_resource_leak() -> None: + """Benchmark test to verify no resource leaks (threads, FDs, memory) between iteration 20 and 100.""" + print("\nStarting resource leak benchmark...") + + process = psutil.Process(os.getpid()) + + def get_metrics(): + return { + "threads": process.num_threads(), + "fds": process.num_fds(), + "rss_mb": process.memory_info().rss / (1024 * 1024), + } + + warmup_iterations = 20 + total_iterations = 100 + + baseline_metrics = None + active_final_metrics = None + + with Connector() as connector: + for i in range(1, total_iterations + 1): + conn = connector.connect( + INSTANCE_CONNECTION_NAME, + "psycopg", + user=DB_USER, + password=DB_PASSWORD, + db=DB_NAME, + ) + cursor = conn.cursor() + cursor.execute("SELECT 1;") + cursor.fetchone() + cursor.close() + conn.close() + + if i == warmup_iterations: + time.sleep(0.5) + baseline_metrics = get_metrics() + print(f"Baseline Metrics (Iteration {i}): {baseline_metrics}") + + if i == total_iterations: + time.sleep(0.5) + active_final_metrics = get_metrics() + print(f"Active Final Metrics (Iteration {i}): {active_final_metrics}") + + if i % 10 == 0 and warmup_iterations < i < total_iterations: + current_metrics = get_metrics() + print(f"Iteration {i:3d}/{total_iterations}: {current_metrics}") + + # Post-close metrics + time.sleep(1) + post_close_metrics = get_metrics() + print(f"Post-Close Metrics: {post_close_metrics}") + + assert baseline_metrics is not None + assert active_final_metrics is not None + + # Assertions: compare Active Final (100) vs Baseline (20) + # Threads should not grow + assert active_final_metrics["threads"] <= baseline_metrics["threads"] + 1, f"Thread leak: {baseline_metrics} -> {active_final_metrics}" + # FDs should not grow + assert active_final_metrics["fds"] <= baseline_metrics["fds"] + 1, f"FD leak: {baseline_metrics} -> {active_final_metrics}" + # Memory growth should be minimal (allow < 5MB growth for minor fragmentation) + assert active_final_metrics["rss_mb"] <= baseline_metrics["rss_mb"] + 5, f"Memory leak: {baseline_metrics} -> {active_final_metrics}" + + print("Resource leak benchmark passed successfully.") + + +def test_system_psycopg_basic() -> None: + """Basic system test to verify connection and query.""" + print(f"\nConnecting to {INSTANCE_CONNECTION_NAME}...") + with Connector() as connector: + conn = connector.connect( + INSTANCE_CONNECTION_NAME, + "psycopg", + user=DB_USER, + password=DB_PASSWORD, + db=DB_NAME, + ) + + cursor = conn.cursor() + cursor.execute("SELECT version();") + result = cursor.fetchone() + print(f"Database version: {result[0]}") + assert result is not None + cursor.close() + conn.close() + print("Connection closed successfully.") + + + + + +def test_system_psycopg_to_thread() -> None: + """Verify that running sync connect in asyncio.to_thread works.""" + print(f"\nConnecting via asyncio.to_thread to {INSTANCE_CONNECTION_NAME}...") + + async def run_connect(): + with Connector() as connector: + # Run the blocking connector.connect in a thread + conn = await asyncio.to_thread( + connector.connect, + INSTANCE_CONNECTION_NAME, + "psycopg", + user=DB_USER, + password=DB_PASSWORD, + db=DB_NAME, + ) + cursor = conn.cursor() + cursor.execute("SELECT version();") + result = cursor.fetchone() + print(f"Database version (to_thread): {result[0]}") + assert result is not None + cursor.close() + conn.close() + + asyncio.run(run_connect()) + print("to_thread connection closed successfully.") diff --git a/tests/system/test_psycopg_iam_auth.py b/tests/system/test_psycopg_iam_auth.py new file mode 100644 index 00000000..bacdede2 --- /dev/null +++ b/tests/system/test_psycopg_iam_auth.py @@ -0,0 +1,96 @@ +""" +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.getenv( + "POSTGRES_CONNECTION_NAME", + "galakp-playground:us-east7:pg-us-east7-psycopg", + ) + user = os.environ["POSTGRES_IAM_USER"] + db = os.getenv("POSTGRES_DB", "postgres") + 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.getenv( + "POSTGRES_CONNECTION_NAME", + "galakp-playground:us-east7:pg-us-east7-psycopg", + ) + user = os.environ["POSTGRES_IAM_USER"] + db = os.getenv("POSTGRES_DB", "postgres") + 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_psycopg.py b/tests/unit/test_psycopg.py new file mode 100644 index 00000000..f8e5e99e --- /dev/null +++ b/tests/unit/test_psycopg.py @@ -0,0 +1,112 @@ +# 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 +import time +from typing import Any +from unittest.mock import MagicMock, patch +import pytest + +from google.cloud.sql.connector.psycopg import _proxy, connect + + +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() + + +@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"]) From c4bd9f1e1be3501bff021e7221f5767d52457c87 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 17:03:04 +0000 Subject: [PATCH 02/19] chore: add psycopg dependencies for testing --- pyproject.toml | 1 + requirements-test.txt | 2 ++ 2 files changed, 3 insertions(+) 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 From c5ae4a75da638656e831fe667d47a9eef90ef43b Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 17:04:55 +0000 Subject: [PATCH 03/19] chore: fix lint --- tests/system/test_psycopg_connection.py | 5 +++-- tests/unit/test_psycopg.py | 8 ++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py index 60fdff08..63fac553 100644 --- a/tests/system/test_psycopg_connection.py +++ b/tests/system/test_psycopg_connection.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os import asyncio +import os import time -import pytest + import psutil + from google.cloud.sql.connector import Connector # These will be set from environment variables or default to our test instance diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py index f8e5e99e..5a032240 100644 --- a/tests/unit/test_psycopg.py +++ b/tests/unit/test_psycopg.py @@ -16,12 +16,12 @@ import socket import ssl import threading -import time from typing import Any -from unittest.mock import MagicMock, patch -import pytest +from unittest.mock import MagicMock +from unittest.mock import patch -from google.cloud.sql.connector.psycopg import _proxy, connect +from google.cloud.sql.connector.psycopg import _proxy +from google.cloud.sql.connector.psycopg import connect def test_proxy_bidirectional() -> None: From 177a7761228af78bc0eccdd5143e40d099971be7 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 17:09:14 +0000 Subject: [PATCH 04/19] fix: make psutil dependency optional in system tests --- tests/system/test_psycopg_connection.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py index 63fac553..451f25e7 100644 --- a/tests/system/test_psycopg_connection.py +++ b/tests/system/test_psycopg_connection.py @@ -16,7 +16,12 @@ import os import time -import psutil +try: + import psutil +except ImportError: + psutil = None + +import pytest from google.cloud.sql.connector import Connector @@ -29,6 +34,7 @@ DB_NAME = os.getenv("DB_NAME", "postgres") +@pytest.mark.skipif(psutil is None, reason="psutil package is not installed") def test_system_psycopg_resource_leak() -> None: """Benchmark test to verify no resource leaks (threads, FDs, memory) between iteration 20 and 100.""" print("\nStarting resource leak benchmark...") From 62c2bf097210d449fc124722a9eb974841d7f5f2 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 17:20:18 +0000 Subject: [PATCH 05/19] test: remove psycopg resource leak test and add remaining system tests --- tests/system/test_psycopg_connection.py | 322 ++++++++++++++---------- 1 file changed, 195 insertions(+), 127 deletions(-) diff --git a/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py index 451f25e7..607692f3 100644 --- a/tests/system/test_psycopg_connection.py +++ b/tests/system/test_psycopg_connection.py @@ -1,154 +1,222 @@ -# 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. +""" +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. +""" import asyncio +from datetime import datetime import os -import time - -try: - import psutil -except ImportError: - psutil = None +from typing import Union import pytest +import sqlalchemy from google.cloud.sql.connector import Connector - -# These will be set from environment variables or default to our test instance -INSTANCE_CONNECTION_NAME = os.getenv( - "DB_CONNECTION_NAME", "galakp-playground:us-east7:pg-us-east7-psycopg" -) -DB_USER = os.getenv("DB_USER", "postgres") -DB_PASSWORD = os.getenv("DB_PASSWORD", "SuperPass123!") -DB_NAME = os.getenv("DB_NAME", "postgres") - - -@pytest.mark.skipif(psutil is None, reason="psutil package is not installed") -def test_system_psycopg_resource_leak() -> None: - """Benchmark test to verify no resource leaks (threads, FDs, memory) between iteration 20 and 100.""" - print("\nStarting resource leak benchmark...") - - process = psutil.Process(os.getpid()) - - def get_metrics(): - return { - "threads": process.num_threads(), - "fds": process.num_fds(), - "rss_mb": process.memory_info().rss / (1024 * 1024), - } - - warmup_iterations = 20 - total_iterations = 100 - - baseline_metrics = None - active_final_metrics = None - - with Connector() as connector: - for i in range(1, total_iterations + 1): - conn = connector.connect( - INSTANCE_CONNECTION_NAME, - "psycopg", - user=DB_USER, - password=DB_PASSWORD, - db=DB_NAME, - ) - cursor = conn.cursor() - cursor.execute("SELECT 1;") - cursor.fetchone() - cursor.close() - conn.close() - - if i == warmup_iterations: - time.sleep(0.5) - baseline_metrics = get_metrics() - print(f"Baseline Metrics (Iteration {i}): {baseline_metrics}") - - if i == total_iterations: - time.sleep(0.5) - active_final_metrics = get_metrics() - print(f"Active Final Metrics (Iteration {i}): {active_final_metrics}") - - if i % 10 == 0 and warmup_iterations < i < total_iterations: - current_metrics = get_metrics() - print(f"Iteration {i:3d}/{total_iterations}: {current_metrics}") - - # Post-close metrics - time.sleep(1) - post_close_metrics = get_metrics() - print(f"Post-Close Metrics: {post_close_metrics}") - - assert baseline_metrics is not None - assert active_final_metrics is not None - - # Assertions: compare Active Final (100) vs Baseline (20) - # Threads should not grow - assert active_final_metrics["threads"] <= baseline_metrics["threads"] + 1, f"Thread leak: {baseline_metrics} -> {active_final_metrics}" - # FDs should not grow - assert active_final_metrics["fds"] <= baseline_metrics["fds"] + 1, f"FD leak: {baseline_metrics} -> {active_final_metrics}" - # Memory growth should be minimal (allow < 5MB growth for minor fragmentation) - assert active_final_metrics["rss_mb"] <= baseline_metrics["rss_mb"] + 5, f"Memory leak: {baseline_metrics} -> {active_final_metrics}" - - print("Resource leak benchmark passed successfully.") - - -def test_system_psycopg_basic() -> None: - """Basic system test to verify connection and query.""" - print(f"\nConnecting to {INSTANCE_CONNECTION_NAME}...") - with Connector() as connector: - conn = connector.connect( - INSTANCE_CONNECTION_NAME, +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: Union[type[DefaultResolver], type[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=DB_USER, - password=DB_PASSWORD, - db=DB_NAME, - ) - - cursor = conn.cursor() - cursor.execute("SELECT version();") - result = cursor.fetchone() - print(f"Database version: {result[0]}") - assert result is not None - cursor.close() - conn.close() - print("Connection closed successfully.") - - - + user=user, + password=password, + db=db, + ip_type=ip_type, + ), + ) + return engine, connector + + +# Fallback to playground values if env vars are missing +def get_env(key: str, default: str = "") -> str: + # Map standard env vars to our playground values as defaults + defaults = { + "POSTGRES_CONNECTION_NAME": "galakp-playground:us-east7:pg-us-east7-psycopg", + "POSTGRES_USER": "postgres", + "POSTGRES_PASS": "SuperPass123!", + "POSTGRES_DB": "postgres", + } + return os.getenv(key, defaults.get(key, default)) + + +def test_psycopg_connection() -> None: + """Basic test to get time from database using psycopg.""" + inst_conn_name = get_env("POSTGRES_CONNECTION_NAME") + user = get_env("POSTGRES_USER") + password = get_env("POSTGRES_PASS") + db = get_env("POSTGRES_DB") + ip_type = os.getenv("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 = get_env("POSTGRES_CONNECTION_NAME") + user = get_env("POSTGRES_USER") + password = get_env("POSTGRES_PASS") + db = get_env("POSTGRES_DB") + ip_type = os.getenv("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 = get_env("POSTGRES_USER") + password = os.environ.get("POSTGRES_CAS_PASS") + db = get_env("POSTGRES_DB") + ip_type = os.getenv("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 = get_env("POSTGRES_USER") + password = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS") + db = get_env("POSTGRES_DB") + ip_type = os.getenv("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 = get_env("POSTGRES_USER") + password = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS") + db = get_env("POSTGRES_DB") + ip_type = os.getenv("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 = get_env("POSTGRES_USER") + password = os.environ.get("POSTGRES_MCP_PASS") + db = get_env("POSTGRES_DB") + ip_type = os.getenv("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.""" - print(f"\nConnecting via asyncio.to_thread to {INSTANCE_CONNECTION_NAME}...") + inst_conn_name = get_env("POSTGRES_CONNECTION_NAME") + user = get_env("POSTGRES_USER") + password = get_env("POSTGRES_PASS") + db = get_env("POSTGRES_DB") async def run_connect(): with Connector() as connector: # Run the blocking connector.connect in a thread conn = await asyncio.to_thread( connector.connect, - INSTANCE_CONNECTION_NAME, + inst_conn_name, "psycopg", - user=DB_USER, - password=DB_PASSWORD, - db=DB_NAME, + user=user, + password=password, + db=db, ) cursor = conn.cursor() - cursor.execute("SELECT version();") + cursor.execute("SELECT NOW();") result = cursor.fetchone() - print(f"Database version (to_thread): {result[0]}") assert result is not None cursor.close() conn.close() asyncio.run(run_connect()) - print("to_thread connection closed successfully.") From 8eb3f09e8c7c25d2468655ab8d727c530e1728de Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 17:40:54 +0000 Subject: [PATCH 06/19] test: remove hardcoded playground credentials from system tests --- tests/system/test_psycopg_connection.py | 64 ++++++++++--------------- tests/system/test_psycopg_iam_auth.py | 14 ++---- 2 files changed, 30 insertions(+), 48 deletions(-) diff --git a/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py index 607692f3..09dc584b 100644 --- a/tests/system/test_psycopg_connection.py +++ b/tests/system/test_psycopg_connection.py @@ -56,25 +56,13 @@ def create_sqlalchemy_engine( return engine, connector -# Fallback to playground values if env vars are missing -def get_env(key: str, default: str = "") -> str: - # Map standard env vars to our playground values as defaults - defaults = { - "POSTGRES_CONNECTION_NAME": "galakp-playground:us-east7:pg-us-east7-psycopg", - "POSTGRES_USER": "postgres", - "POSTGRES_PASS": "SuperPass123!", - "POSTGRES_DB": "postgres", - } - return os.getenv(key, defaults.get(key, default)) - - def test_psycopg_connection() -> None: """Basic test to get time from database using psycopg.""" - inst_conn_name = get_env("POSTGRES_CONNECTION_NAME") - user = get_env("POSTGRES_USER") - password = get_env("POSTGRES_PASS") - db = get_env("POSTGRES_DB") - ip_type = os.getenv("IP_TYPE", "public") + 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 @@ -89,11 +77,11 @@ def test_psycopg_connection() -> None: def test_lazy_psycopg_connection() -> None: """Basic test to get time from database using psycopg and lazy refresh.""" - inst_conn_name = get_env("POSTGRES_CONNECTION_NAME") - user = get_env("POSTGRES_USER") - password = get_env("POSTGRES_PASS") - db = get_env("POSTGRES_DB") - ip_type = os.getenv("IP_TYPE", "public") + 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" @@ -109,10 +97,10 @@ def test_lazy_psycopg_connection() -> None: 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 = get_env("POSTGRES_USER") + user = os.environ["POSTGRES_USER"] password = os.environ.get("POSTGRES_CAS_PASS") - db = get_env("POSTGRES_DB") - ip_type = os.getenv("IP_TYPE", "public") + 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") @@ -131,10 +119,10 @@ def test_CAS_psycopg_connection() -> None: 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 = get_env("POSTGRES_USER") + user = os.environ["POSTGRES_USER"] password = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS") - db = get_env("POSTGRES_DB") - ip_type = os.getenv("IP_TYPE", "public") + 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") @@ -153,10 +141,10 @@ def test_customer_managed_CAS_psycopg_connection() -> None: 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 = get_env("POSTGRES_USER") + user = os.environ["POSTGRES_USER"] password = os.environ.get("POSTGRES_CUSTOMER_CAS_PASS") - db = get_env("POSTGRES_DB") - ip_type = os.getenv("IP_TYPE", "public") + 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") @@ -175,10 +163,10 @@ def test_custom_SAN_with_dns_psycopg_connection() -> None: 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 = get_env("POSTGRES_USER") + user = os.environ["POSTGRES_USER"] password = os.environ.get("POSTGRES_MCP_PASS") - db = get_env("POSTGRES_DB") - ip_type = os.getenv("IP_TYPE", "public") + 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") @@ -196,10 +184,10 @@ def test_MCP_psycopg_connection() -> None: def test_system_psycopg_to_thread() -> None: """Verify that running sync connect in asyncio.to_thread works.""" - inst_conn_name = get_env("POSTGRES_CONNECTION_NAME") - user = get_env("POSTGRES_USER") - password = get_env("POSTGRES_PASS") - db = get_env("POSTGRES_DB") + inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"] + user = os.environ["POSTGRES_USER"] + password = os.environ["POSTGRES_PASS"] + db = os.environ["POSTGRES_DB"] async def run_connect(): with Connector() as connector: diff --git a/tests/system/test_psycopg_iam_auth.py b/tests/system/test_psycopg_iam_auth.py index bacdede2..ebb034db 100644 --- a/tests/system/test_psycopg_iam_auth.py +++ b/tests/system/test_psycopg_iam_auth.py @@ -58,12 +58,9 @@ def create_sqlalchemy_engine( def test_psycopg_iam_authn_connection() -> None: """Basic test to get time from database using psycopg and IAM Authn.""" - inst_conn_name = os.getenv( - "POSTGRES_CONNECTION_NAME", - "galakp-playground:us-east7:pg-us-east7-psycopg", - ) + inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"] user = os.environ["POSTGRES_IAM_USER"] - db = os.getenv("POSTGRES_DB", "postgres") + db = os.environ["POSTGRES_DB"] ip_type = os.getenv("IP_TYPE", "public") engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, ip_type) @@ -77,12 +74,9 @@ def test_psycopg_iam_authn_connection() -> None: 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.getenv( - "POSTGRES_CONNECTION_NAME", - "galakp-playground:us-east7:pg-us-east7-psycopg", - ) + inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"] user = os.environ["POSTGRES_IAM_USER"] - db = os.getenv("POSTGRES_DB", "postgres") + db = os.environ["POSTGRES_DB"] ip_type = os.getenv("IP_TYPE", "public") engine, connector = create_sqlalchemy_engine( From ac288ebef8bd03c81c27962ea42739e7c84e0df4 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 17:44:59 +0000 Subject: [PATCH 07/19] fix: resolve ruff linter errors in psycopg system tests --- tests/system/test_psycopg_connection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py index 09dc584b..ebbe910e 100644 --- a/tests/system/test_psycopg_connection.py +++ b/tests/system/test_psycopg_connection.py @@ -13,11 +13,11 @@ 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 -from typing import Union import pytest import sqlalchemy @@ -34,7 +34,7 @@ def create_sqlalchemy_engine( db: str, ip_type: str = "public", refresh_strategy: str = "background", - resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver, + 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. From bc20b3456196991c42a3f7c94bc80c6dde4c63f9 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 17:53:31 +0000 Subject: [PATCH 08/19] fix: pass ip_type to connector.connect in to_thread system test --- tests/system/test_psycopg_connection.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py index ebbe910e..9424d86d 100644 --- a/tests/system/test_psycopg_connection.py +++ b/tests/system/test_psycopg_connection.py @@ -188,6 +188,7 @@ def test_system_psycopg_to_thread() -> None: 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: @@ -199,6 +200,7 @@ async def run_connect(): user=user, password=password, db=db, + ip_type=ip_type, ) cursor = conn.cursor() cursor.execute("SELECT NOW();") From fb6449fdcfdc8b43a26e9bec8a1d185ac5969352 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 18:01:53 +0000 Subject: [PATCH 09/19] fix: refactor psycopg proxy to be single-threaded selectors-based, resolving SSLSocket thread-safety deadlocks --- google/cloud/sql/connector/psycopg.py | 117 +++++++++++++++++--------- 1 file changed, 78 insertions(+), 39 deletions(-) diff --git a/google/cloud/sql/connector/psycopg.py b/google/cloud/sql/connector/psycopg.py index c930f7b2..df93f078 100644 --- a/google/cloud/sql/connector/psycopg.py +++ b/google/cloud/sql/connector/psycopg.py @@ -14,6 +14,7 @@ import logging import os +import selectors import socket import ssl import tempfile @@ -25,48 +26,86 @@ logger = logging.getLogger(name=__name__) -_CHUNK_SIZE = 8 * 1024 # bytes per recv() call inside the proxy forwarding loop - def _proxy(local: socket.socket, remote: "ssl.SSLSocket") -> None: - """Bidirectionally proxy bytes between a local Unix socket and a remote - SSL socket. - - Spawns one daemon thread for the remote→local direction and runs the - local→remote direction in the calling thread. Blocks until the calling - thread's direction reaches EOF or a socket error, at which point both - sockets are closed so the other thread also unblocks and exits. + """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). + """ + while remote.pending() > 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 + return False - Args: - local: The Unix domain socket connected to the database driver. - remote: The SSL socket connected to the Cloud SQL proxy server. - """ - def forward(src: Any, dst: Any) -> None: - buf = bytearray(_CHUNK_SIZE) - view = memoryview(buf) - try: - while True: - n = src.recv_into(view) - if n == 0: - logger.debug("psycopg proxy: EOF on %s, closing both sockets", src) - break - dst.sendall(view[:n]) - except (OSError, ssl.SSLError) as e: - logger.debug("psycopg proxy: socket error on %s: %s", src, e) - finally: - # Close both ends so the sibling thread also unblocks. - for s in (local, remote): - try: - s.shutdown(socket.SHUT_RDWR) - except OSError: - pass - try: - s.close() - except OSError: - pass - - threading.Thread(target=forward, args=(remote, local), daemon=True).start() - forward(local, remote) # run in calling thread rather than spawning a third + try: + while True: + # First check if there is any pending data in SSL buffer + if forward_pending(): + break + + events = sel.select(timeout=30) + if not events: + logger.debug("psycopg proxy: inactivity timeout (30s)") + break + + 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( From 0222334ee26dacc33128458eb85e91e8a9e485f8 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 18:13:33 +0000 Subject: [PATCH 10/19] fix: make proxy robust to raw sockets and mocks in unit tests --- google/cloud/sql/connector/psycopg.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/google/cloud/sql/connector/psycopg.py b/google/cloud/sql/connector/psycopg.py index df93f078..2f46e7f5 100644 --- a/google/cloud/sql/connector/psycopg.py +++ b/google/cloud/sql/connector/psycopg.py @@ -37,7 +37,16 @@ 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). """ - while remote.pending() > 0: + if not hasattr(remote, "pending"): + return False + try: + pending_bytes = remote.pending() + except AttributeError: + return False + if not isinstance(pending_bytes, int): + return False + + while pending_bytes > 0: try: data = remote.recv(8192) except OSError as e: @@ -51,6 +60,12 @@ def forward_pending() -> bool: except OSError as e: logger.debug("psycopg proxy: local send pending error: %s", e) return True + try: + pending_bytes = remote.pending() + except (AttributeError, OSError): + break + if not isinstance(pending_bytes, int): + break return False try: From 9a8cd412f9bab716ca4ce0f709d44375ba3af9e9 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 18:46:11 +0000 Subject: [PATCH 11/19] test: add unit test for psycopg proxy pending data forwarding --- tests/unit/test_psycopg.py | 44 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py index 5a032240..001edff3 100644 --- a/tests/unit/test_psycopg.py +++ b/tests/unit/test_psycopg.py @@ -65,6 +65,50 @@ def test_proxy_bidirectional() -> None: 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.""" From 4c2c10151e5835b61dd3d22dbdfac64a52770264 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 19:43:59 +0000 Subject: [PATCH 12/19] refactor: simplify pending bytes checks by removing redundant try-except AttributeError --- google/cloud/sql/connector/psycopg.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/google/cloud/sql/connector/psycopg.py b/google/cloud/sql/connector/psycopg.py index 2f46e7f5..eaa6381d 100644 --- a/google/cloud/sql/connector/psycopg.py +++ b/google/cloud/sql/connector/psycopg.py @@ -39,10 +39,7 @@ def forward_pending() -> bool: """ if not hasattr(remote, "pending"): return False - try: - pending_bytes = remote.pending() - except AttributeError: - return False + pending_bytes = remote.pending() if not isinstance(pending_bytes, int): return False @@ -62,7 +59,7 @@ def forward_pending() -> bool: return True try: pending_bytes = remote.pending() - except (AttributeError, OSError): + except OSError: break if not isinstance(pending_bytes, int): break From f1e1c8b16976f8460581d306e21132a399a59f69 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 19:50:39 +0000 Subject: [PATCH 13/19] fix: make accept/proxy thread robust against exceptions during setup, preventing resource leaks --- google/cloud/sql/connector/psycopg.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/google/cloud/sql/connector/psycopg.py b/google/cloud/sql/connector/psycopg.py index eaa6381d..f78d07ae 100644 --- a/google/cloud/sql/connector/psycopg.py +++ b/google/cloud/sql/connector/psycopg.py @@ -155,18 +155,32 @@ def connect( 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") - except OSError as e: - logger.debug("psycopg proxy: accept failed: %s", e) + _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 - return - _proxy(unix_conn, remote_sock) threading.Thread(target=_accept_and_proxy, daemon=True).start() From 23516dd82cd9d85a078bff8305f252212f84fc3e Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 19:54:59 +0000 Subject: [PATCH 14/19] test: add coverage tests for psycopg proxy timeout, handshake failure, and remote EOF --- tests/unit/test_psycopg.py | 105 +++++++++++++++++++++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py index 001edff3..23390127 100644 --- a/tests/unit/test_psycopg.py +++ b/tests/unit/test_psycopg.py @@ -154,3 +154,108 @@ def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock: # Verify temp dir was cleaned up assert not os.path.exists(kwargs["host"]) + + +def test_proxy_timeout() -> None: + """Test that _proxy exits and cleans up on selectors timeout.""" + local_client, local_server = socket.socketpair() + remote_client, remote_server = socket.socketpair() + + # Mock selectors.DefaultSelector.select to return empty list (timeout) + with patch( + "google.cloud.sql.connector.psycopg.selectors.DefaultSelector" + ) as mock_selector_cls: + mock_selector = MagicMock() + mock_selector.select.return_value = [] # Timeout + mock_selector_cls.return_value = mock_selector + + # Run proxy + _proxy(local_server, remote_client) + + # Sockets should be closed + assert local_server.fileno() == -1 + assert remote_client.fileno() == -1 + + # Clean up outer sockets + local_client.close() + remote_server.close() + + +@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 + import pytest + + 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() From 4fa7756f4f6fdb8317dd68c288e7cd866093f6e3 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 20:06:06 +0000 Subject: [PATCH 15/19] test: add comprehensive socket error and cleanup coverage tests for psycopg --- tests/unit/test_psycopg.py | 186 +++++++++++++++++++++++++++++++++++++ 1 file changed, 186 insertions(+) diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py index 23390127..e3437790 100644 --- a/tests/unit/test_psycopg.py +++ b/tests/unit/test_psycopg.py @@ -24,6 +24,20 @@ from google.cloud.sql.connector.psycopg import connect +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) @@ -259,3 +273,175 @@ def test_proxy_remote_eof() -> None: 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) + + import pytest + + 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 + From ab1d9543220afa8a0546c268836fa786e42e7f57 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Tue, 4 Aug 2026 20:34:47 +0000 Subject: [PATCH 16/19] docs: add psycopg to supported drivers and usage examples in README --- README.md | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) 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. From a30f5adef0ed64cfde1179b36619d0ba3ba116ac Mon Sep 17 00:00:00 2001 From: kgala2 Date: Thu, 13 Aug 2026 17:04:53 +0000 Subject: [PATCH 17/19] fix: remove inactivity timeout in psycopg proxy to preserve idle connections --- google/cloud/sql/connector/psycopg.py | 5 +- tests/unit/test_psycopg.py | 77 ++++++++++----------------- 2 files changed, 28 insertions(+), 54 deletions(-) diff --git a/google/cloud/sql/connector/psycopg.py b/google/cloud/sql/connector/psycopg.py index f78d07ae..e5942951 100644 --- a/google/cloud/sql/connector/psycopg.py +++ b/google/cloud/sql/connector/psycopg.py @@ -71,10 +71,7 @@ def forward_pending() -> bool: if forward_pending(): break - events = sel.select(timeout=30) - if not events: - logger.debug("psycopg proxy: inactivity timeout (30s)") - break + events = sel.select() for key, mask in events: if key.data == "local": diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py index e3437790..a065d5c1 100644 --- a/tests/unit/test_psycopg.py +++ b/tests/unit/test_psycopg.py @@ -20,9 +20,16 @@ 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 @@ -170,31 +177,6 @@ def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock: assert not os.path.exists(kwargs["host"]) -def test_proxy_timeout() -> None: - """Test that _proxy exits and cleans up on selectors timeout.""" - local_client, local_server = socket.socketpair() - remote_client, remote_server = socket.socketpair() - - # Mock selectors.DefaultSelector.select to return empty list (timeout) - with patch( - "google.cloud.sql.connector.psycopg.selectors.DefaultSelector" - ) as mock_selector_cls: - mock_selector = MagicMock() - mock_selector.select.return_value = [] # Timeout - mock_selector_cls.return_value = mock_selector - - # Run proxy - _proxy(local_server, remote_client) - - # Sockets should be closed - assert local_server.fileno() == -1 - assert remote_client.fileno() == -1 - - # Clean up outer sockets - local_client.close() - remote_server.close() - - @patch("psycopg.connect") def test_connect_wrapper_failure(mock_psycopg_connect: MagicMock) -> None: """Test that connect wrapper cleans up correctly when psycopg.connect fails.""" @@ -202,8 +184,6 @@ def test_connect_wrapper_failure(mock_psycopg_connect: MagicMock) -> None: mock_psycopg_connect.side_effect = Exception("connection failed simulated") # Call the connect wrapper and expect it to raise - import pytest - with pytest.raises(Exception, match="connection failed simulated"): connect( "127.0.0.1", @@ -217,18 +197,16 @@ def test_connect_wrapper_failure(mock_psycopg_connect: MagicMock) -> None: assert mock_remote_sock.close.called # Verify cleanup with mocked paths - with patch( - "google.cloud.sql.connector.psycopg.tempfile.mkdtemp" - ) as mock_mkdtemp: + 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: + 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 @@ -389,9 +367,7 @@ def test_proxy_pending_local_send_error() -> None: 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") - ) + local_server.sendall = MagicMock(side_effect=OSError("local send pending failed")) # Run proxy _proxy(local_server, remote_client) @@ -408,8 +384,6 @@ def test_connect_import_error() -> None: """Test that connect raises ImportError if psycopg is not installed.""" mock_remote_sock = MagicMock(spec=ssl.SSLSocket) - import pytest - with pytest.raises(ImportError, match='Unable to import module "psycopg."'): connect("127.0.0.1", mock_remote_sock) @@ -426,13 +400,17 @@ def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock: 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: + 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, @@ -444,4 +422,3 @@ def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock: assert conn is not None assert mock_remove.called assert mock_rmdir.called - From a77d7688f5b2e433a27a32beeebaf30c06fa9e39 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Fri, 14 Aug 2026 18:18:09 +0000 Subject: [PATCH 18/19] test(psycopg): add unit tests for happy path and error cleanup paths to reach 100% coverage --- tests/unit/test_psycopg.py | 309 +++++++++++++++++++++++++++++++++++++ 1 file changed, 309 insertions(+) diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py index a065d5c1..6ed51b7e 100644 --- a/tests/unit/test_psycopg.py +++ b/tests/unit/test_psycopg.py @@ -422,3 +422,312 @@ def mock_connect_impl(*args: Any, **kwargs: Any) -> MagicMock: 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 + + From 248d68dd0d3a44ece7c4620f508db0a3fe0ae0d3 Mon Sep 17 00:00:00 2001 From: kgala2 Date: Fri, 14 Aug 2026 18:53:30 +0000 Subject: [PATCH 19/19] test(connector): add unit test for Connector.connect_async with psycopg driver --- tests/unit/test_connector.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) 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() + +