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..f78d07ae --- /dev/null +++ b/google/cloud/sql/connector/psycopg.py @@ -0,0 +1,233 @@ +# 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(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( + 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 221e8f34..1ae258e6 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_psycopg.py b/tests/unit/test_psycopg.py new file mode 100644 index 00000000..e3437790 --- /dev/null +++ b/tests/unit/test_psycopg.py @@ -0,0 +1,447 @@ +# 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 + +from google.cloud.sql.connector.psycopg import _proxy +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) + 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"]) + + +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() + + +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 +