From 637f3317cdcb8bb645ae5369b92308ee8abf6558 Mon Sep 17 00:00:00 2001 From: Eno Compton Date: Thu, 5 Mar 2026 19:23:33 -0700 Subject: [PATCH] feat: add support for psycopg Psycopg does not support providing a pre-connected socket in its connect API as there is no matching API in the underlying libpq C interface. To work around this limitation, this commit creates an in-memory proxy that copies data between a Unix domain socket and the secure SSL Socket created by the Connector. This approach works, but suffers from poor performance and makes the driver comparable to pg8000 (a pure Python implementation). The best solution would be for psycopg (and libpq) to expose a way to provide a socket creator function as we have in the other drivers. Until then, this is better than nothing. Fixes #377. --- README.md | 39 ++- google/cloud/alloydbconnector/client.py | 8 +- google/cloud/alloydbconnector/connector.py | 6 +- google/cloud/alloydbconnector/psycopg.py | 167 ++++++++++++ pyproject.toml | 1 + requirements-test.txt | 2 + tests/system/test_psycopg_connection.py | 114 +++++++++ tests/system/test_psycopg_iam_authn.py | 106 ++++++++ tests/system/test_psycopg_psc.py | 90 +++++++ tests/system/test_psycopg_public_ip.py | 94 +++++++ tests/unit/test_client.py | 4 +- tests/unit/test_connector.py | 21 ++ tests/unit/test_psycopg.py | 281 +++++++++++++++++++++ 13 files changed, 925 insertions(+), 8 deletions(-) create mode 100644 google/cloud/alloydbconnector/psycopg.py create mode 100644 tests/system/test_psycopg_connection.py create mode 100644 tests/system/test_psycopg_iam_authn.py create mode 100644 tests/system/test_psycopg_psc.py create mode 100644 tests/system/test_psycopg_public_ip.py create mode 100644 tests/unit/test_psycopg.py diff --git a/README.md b/README.md index caf9acc4..0b9c1600 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,11 @@ Python applications. It provides: [iam-db-authn]: https://cloud.google.com/alloydb/docs/manage-iam-authn -**Supported drivers:** [`pg8000`](https://codeberg.org/tlocke/pg8000) (sync) · [`asyncpg`](https://magicstack.github.io/asyncpg) (async) +**Supported drivers:** [`pg8000`][pg8000] (sync) · [`psycopg`][psycopg] (sync) · [`asyncpg`][asyncpg] (async) + +[pg8000]: https://codeberg.org/tlocke/pg8000 +[psycopg]: https://www.psycopg.org/ +[asyncpg]: https://magicstack.github.io/asyncpg ## Quickstart @@ -67,6 +71,37 @@ with Connector() as connector: print(result) ``` +### Sync (psycopg + SQLAlchemy) + +**Install:** +```sh +pip install "google-cloud-alloydb-connector[psycopg]" sqlalchemy +``` + +**Connect:** +```python +import sqlalchemy +from google.cloud.alloydbconnector import Connector + +INSTANCE_URI = "projects/MY_PROJECT/locations/MY_REGION/clusters/MY_CLUSTER/instances/MY_INSTANCE" + +with Connector() as connector: + pool = sqlalchemy.create_engine( + "postgresql+psycopg://", + creator=lambda: connector.connect( + INSTANCE_URI, + "psycopg", + user="my-user", + password="my-password", + db="my-db", + ), + ) + + with pool.connect() as conn: + result = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + print(result) +``` + ### Async (asyncpg + SQLAlchemy) **Install:** @@ -188,7 +223,7 @@ database user][add-iam-user]. ```python connector.connect( INSTANCE_URI, - "pg8000", # or "asyncpg" + "pg8000", # or "psycopg" (sync) / "asyncpg" (async) user="service-account@my-project.iam", # omit .gserviceaccount.com suffix db="my-db", enable_iam_auth=True, diff --git a/google/cloud/alloydbconnector/client.py b/google/cloud/alloydbconnector/client.py index 5de91ac2..075a2af1 100644 --- a/google/cloud/alloydbconnector/client.py +++ b/google/cloud/alloydbconnector/client.py @@ -97,7 +97,7 @@ def __init__( self._is_sync = False if client: self._client = client - elif driver == "pg8000": + elif driver == "pg8000" or driver == "psycopg": self._client = v1beta.AlloyDBAdminClient( credentials=credentials, transport="grpc", @@ -126,7 +126,11 @@ def __init__( self._credentials = credentials # asyncpg does not currently support using metadata exchange # only use metadata exchange for pg8000 driver - self._use_metadata = True if driver == "pg8000" else False + use_metadata = True + if driver in ("asyncpg", None): + use_metadata = False + + self._use_metadata = use_metadata self._user_agent = user_agent async def _get_metadata( diff --git a/google/cloud/alloydbconnector/connector.py b/google/cloud/alloydbconnector/connector.py index db392f65..061f000e 100644 --- a/google/cloud/alloydbconnector/connector.py +++ b/google/cloud/alloydbconnector/connector.py @@ -24,7 +24,7 @@ import struct from threading import Thread from types import TracebackType -from typing import Any, Optional, TYPE_CHECKING +from typing import Any, Callable, Optional, TYPE_CHECKING from google.auth import default from google.auth.credentials import TokenState @@ -39,6 +39,7 @@ from google.cloud.alloydbconnector.instance import RefreshAheadCache from google.cloud.alloydbconnector.lazy import LazyRefreshCache import google.cloud.alloydbconnector.pg8000 as pg8000 +import google.cloud.alloydbconnector.psycopg as psycopg from google.cloud.alloydbconnector.static import StaticConnectionInfoCache from google.cloud.alloydbconnector.types import CacheTypes from google.cloud.alloydbconnector.utils import generate_keys @@ -227,8 +228,9 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) -> self._cache[instance_uri] = cache logger.debug(f"['{instance_uri}']: Connection info added to cache") - connect_func = { + connect_func: dict[str, Callable[..., Any]] = { "pg8000": pg8000.connect, + "psycopg": psycopg.connect, } # only accept supported database drivers try: diff --git a/google/cloud/alloydbconnector/psycopg.py b/google/cloud/alloydbconnector/psycopg.py new file mode 100644 index 00000000..89bac5fb --- /dev/null +++ b/google/cloud/alloydbconnector/psycopg.py @@ -0,0 +1,167 @@ +# 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 AlloyDB 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: + # shutdown is required on POSIX systems to forcefully + # interrupt the sibling thread's blocking recv_into() call. + # Simply closing the socket does not interrupt the system + # call and leads to leaked threads. + 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(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 AlloyDB TLS connection. + + Args: + remote_sock (ssl.SSLSocket): SSL/TLS secure socket stream connected to the + AlloyDB proxy server. + + Returns: + psycopg.Connection: A psycopg Connection object for the AlloyDB 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) + + 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 4cfa0e0c..b1bb87d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,7 @@ Changelog = "https://github.com/GoogleCloudPlatform/alloydb-python-connector/blo [project.optional-dependencies] pg8000 = ["pg8000>=1.31.1"] asyncpg = ["asyncpg>=0.31.0"] +psycopg = ["psycopg>=3.1.0"] [tool.setuptools.dynamic] version = { attr = "google.cloud.alloydbconnector.version.__version__" } diff --git a/requirements-test.txt b/requirements-test.txt index b92d5bd9..84168893 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -2,6 +2,8 @@ asyncpg==0.31.0 mock==5.2.0 pg8000==1.31.5 psycopg2-binary==2.9.11 +psycopg==3.3.3 +psycopg-binary==3.3.3 pytest==9.0.2 pytest-asyncio==1.3.0 pytest-cov==7.0.0 diff --git a/tests/system/test_psycopg_connection.py b/tests/system/test_psycopg_connection.py new file mode 100644 index 00000000..8981fa21 --- /dev/null +++ b/tests/system/test_psycopg_connection.py @@ -0,0 +1,114 @@ +# 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. + +from datetime import datetime +import os + +# [START alloydb_sqlalchemy_connect_connector_psycopg] +import sqlalchemy + +from google.cloud.alloydbconnector import Connector + + +def create_sqlalchemy_engine( + inst_uri: str, + user: str, + password: str, + db: str, + refresh_strategy: str = "background", +) -> tuple[sqlalchemy.engine.Engine, Connector]: + """Creates a connection pool for an AlloyDB instance and returns the pool + and the connector. Callers are responsible for closing the pool and the + connector. + + A sample invocation looks like: + + engine, connector = create_sqlalchemy_engine( + inst_uri, + user, + password, + db, + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + # do something with query result + connector.close() + + Args: + instance_uri (str): + The instance URI specifies the instance relative to the project, + region, and cluster. For example: + "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" + user (str): + The database user name, e.g., postgres + password (str): + The database user's password, e.g., secret-password + db (str): + The name of the database, e.g., mydb + refresh_strategy (Optional[str]): + Refresh strategy for the AlloyDB Connector. Can be one of "lazy" + or "background". For serverless environments use "lazy" to avoid + errors resulting from CPU being throttled. + """ + connector = Connector(refresh_strategy=refresh_strategy) + + # create SQLAlchemy connection pool + engine = sqlalchemy.create_engine( + "postgresql+psycopg://", + creator=lambda: connector.connect( + inst_uri, + "psycopg", + user=user, + password=password, + db=db, + ), + ) + return engine, connector + + +# [END alloydb_sqlalchemy_connect_connector_psycopg] + + +def test_psycopg_connection() -> None: + """Basic test to get time from database.""" + inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] + user = os.environ["ALLOYDB_USER"] + password = os.environ["ALLOYDB_PASS"] + db = os.environ["ALLOYDB_DB"] + + engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db) + 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.""" + inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] + user = os.environ["ALLOYDB_USER"] + password = os.environ["ALLOYDB_PASS"] + db = os.environ["ALLOYDB_DB"] + + engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db, "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/system/test_psycopg_iam_authn.py b/tests/system/test_psycopg_iam_authn.py new file mode 100644 index 00000000..508431b3 --- /dev/null +++ b/tests/system/test_psycopg_iam_authn.py @@ -0,0 +1,106 @@ +# 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. + +from datetime import datetime +import os + +# [START alloydb_sqlalchemy_connect_connector_psycopg_iam_authn] +import sqlalchemy + +from google.cloud.alloydbconnector import Connector + + +def create_sqlalchemy_engine( + inst_uri: str, user: str, db: str, refresh_strategy: str = "background" +) -> tuple[sqlalchemy.engine.Engine, Connector]: + """Creates a connection pool for an AlloyDB instance and returns the pool + and the connector. Callers are responsible for closing the pool and the + connector. + + A sample invocation looks like: + + engine, connector = create_sqlalchemy_engine( + inst_uri, + user, + db, + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + # do something with query result + connector.close() + + Args: + instance_uri (str): + The instance URI specifies the instance relative to the project, + region, and cluster. For example: + "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" + user (str): + The formatted IAM database username. + e.g., my-email@test.com, service-account@project-id.iam + db (str): + The name of the database, e.g., mydb + refresh_strategy (Optional[str]): + Refresh strategy for the AlloyDB Connector. Can be one of "lazy" + or "background". For serverless environments use "lazy" to avoid + errors resulting from CPU being throttled. + """ + connector = Connector(refresh_strategy=refresh_strategy) + + # create SQLAlchemy connection pool + engine = sqlalchemy.create_engine( + "postgresql+psycopg://", + creator=lambda: connector.connect( + inst_uri, + "psycopg", + user=user, + db=db, + enable_iam_auth=True, + ), + ) + return engine, connector + + +# [END alloydb_sqlalchemy_connect_connector_psycopg_iam_authn] + + +def test_psycopg_iam_authn_time() -> None: + """Basic test to get time from database.""" + inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] + user = os.environ["ALLOYDB_IAM_USER"] + db = os.environ["ALLOYDB_DB"] + + engine, connector = create_sqlalchemy_engine(inst_uri, user, db) + 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_psycopg_iam_authn_lazy() -> None: + """Basic test to get time from database.""" + inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] + user = os.environ["ALLOYDB_IAM_USER"] + db = os.environ["ALLOYDB_DB"] + + engine, connector = create_sqlalchemy_engine(inst_uri, user, db, "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/system/test_psycopg_psc.py b/tests/system/test_psycopg_psc.py new file mode 100644 index 00000000..e72e9be9 --- /dev/null +++ b/tests/system/test_psycopg_psc.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 +# +# 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. + +from datetime import datetime +import os + +import sqlalchemy + +from google.cloud.alloydbconnector import Connector + + +def create_sqlalchemy_engine( + inst_uri: str, + user: str, + password: str, + db: str, +) -> tuple[sqlalchemy.engine.Engine, Connector]: + """Creates a connection pool for an AlloyDB instance and returns the pool + and the connector. Callers are responsible for closing the pool and the + connector. + + A sample invocation looks like: + + engine, connector = create_sqlalchemy_engine( + inst_uri, + user, + password, + db, + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + # do something with query result + connector.close() + + Args: + instance_uri (str): + The instance URI specifies the instance relative to the project, + region, and cluster. For example: + "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" + user (str): + The database user name, e.g., postgres + password (str): + The database user's password, e.g., secret-password + db (str): + The name of the database, e.g., mydb + """ + connector = Connector() + + # create SQLAlchemy connection pool + engine = sqlalchemy.create_engine( + "postgresql+psycopg://", + creator=lambda: connector.connect( + inst_uri, + "psycopg", + user=user, + password=password, + db=db, + ip_type="PSC", + ), + ) + return engine, connector + + +def test_psycopg_time() -> None: + """Basic test to get time from database.""" + inst_uri = os.environ["ALLOYDB_PSC_INSTANCE_URI"] + user = os.environ["ALLOYDB_USER"] + password = os.environ["ALLOYDB_PASS"] + db = os.environ["ALLOYDB_DB"] + + engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db) + 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/system/test_psycopg_public_ip.py b/tests/system/test_psycopg_public_ip.py new file mode 100644 index 00000000..d2753acf --- /dev/null +++ b/tests/system/test_psycopg_public_ip.py @@ -0,0 +1,94 @@ +# 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. + +from datetime import datetime +import os + +# [START alloydb_sqlalchemy_connect_connector_psycopg_public_ip] +import sqlalchemy + +from google.cloud.alloydbconnector import Connector + + +def create_sqlalchemy_engine( + inst_uri: str, + user: str, + password: str, + db: str, +) -> tuple[sqlalchemy.engine.Engine, Connector]: + """Creates a connection pool for an AlloyDB instance and returns the pool + and the connector. Callers are responsible for closing the pool and the + connector. + + A sample invocation looks like: + + engine, connector = create_sqlalchemy_engine( + inst_uri, + user, + password, + db, + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + # do something with query result + connector.close() + + Args: + instance_uri (str): + The instance URI specifies the instance relative to the project, + region, and cluster. For example: + "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance" + user (str): + The database user name, e.g., postgres + password (str): + The database user's password, e.g., secret-password + db (str): + The name of the database, e.g., mydb + """ + connector = Connector() + + # create SQLAlchemy connection pool + engine = sqlalchemy.create_engine( + "postgresql+psycopg://", + creator=lambda: connector.connect( + inst_uri, + "psycopg", + user=user, + password=password, + db=db, + ip_type="PUBLIC", + ), + ) + return engine, connector + + +# [END alloydb_sqlalchemy_connect_connector_psycopg_public_ip] + + +def test_psycopg_time() -> None: + """Basic test to get time from database.""" + inst_uri = os.environ["ALLOYDB_INSTANCE_URI"] + user = os.environ["ALLOYDB_USER"] + password = os.environ["ALLOYDB_PASS"] + db = os.environ["ALLOYDB_DB"] + + engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db) + 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_client.py b/tests/unit/test_client.py index 5e682f21..2007b43e 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -243,7 +243,7 @@ async def test_AlloyDBClient_init_async_client(credentials: FakeCredentials) -> @pytest.mark.parametrize( "driver", - [None, "pg8000", "asyncpg"], + [None, "pg8000", "asyncpg", "psycopg"], ) @pytest.mark.asyncio async def test_AlloyDBClient_user_agent( @@ -266,7 +266,7 @@ async def test_AlloyDBClient_user_agent( @pytest.mark.parametrize( "driver, expected", - [(None, False), ("pg8000", True), ("asyncpg", False)], + [(None, False), ("pg8000", True), ("asyncpg", False), ("psycopg", True)], ) @pytest.mark.asyncio async def test_AlloyDBClient_use_metadata( diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index b950de61..5729be1e 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -273,6 +273,27 @@ def test_connect_bad_ip_type( ) +@pytest.mark.usefixtures("proxy_server") +def test_connect_psycopg( + credentials: FakeCredentials, fake_client: FakeAlloyDBClient +) -> None: + """ + Test that connector.connect returns a connection object when using psycopg. + """ + with Connector(credentials) as connector: + connector._client = fake_client + with patch("google.cloud.alloydbconnector.psycopg.connect") as mock_connect: + mock_connect.return_value = True + connection = connector.connect( + "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", + "psycopg", + user="test-user", + password="test-password", + db="test-db", + ) + assert connection is True + + def test_connect_unsupported_driver(credentials: FakeCredentials) -> None: """ Test that connector.connect errors with unsupported database driver. diff --git a/tests/unit/test_psycopg.py b/tests/unit/test_psycopg.py new file mode 100644 index 00000000..db127f79 --- /dev/null +++ b/tests/unit/test_psycopg.py @@ -0,0 +1,281 @@ +# 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 sys +import threading +import types +from typing import Any + +import pytest + +from google.cloud.alloydbconnector.psycopg import _proxy +from google.cloud.alloydbconnector.psycopg import connect + +pytestmark = pytest.mark.skipif( + not hasattr(socket, "AF_UNIX"), + reason="Unix domain sockets (AF_UNIX) not available on this platform", +) + + +def _socketpair() -> tuple[socket.socket, socket.socket]: + """Return a connected pair of Unix domain sockets.""" + return socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM) + + +def test_connect_calls_psycopg(monkeypatch: pytest.MonkeyPatch) -> None: + """connect() imports psycopg and forwards the right keyword arguments.""" + captured: dict = {} + + class FakeConn: + pass + + def fake_psycopg_connect(**kwargs: Any) -> FakeConn: + captured.update(kwargs) + return FakeConn() + + fake_module = types.ModuleType("psycopg") + fake_module.connect = fake_psycopg_connect # type: ignore[attr-defined] + + # Create a real socket pair so the server_sock.accept() can succeed + client_sock, server_sock = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM) + + # We also need a "ssl_sock" stand-in. Use a plain socket; the proxy + # just calls recv/sendall so the type doesn't matter in the test. + ssl_a, ssl_b = _socketpair() + + monkeypatch.setitem(__import__("sys").modules, "psycopg", fake_module) + + conn = connect( + ssl_a, # type: ignore[arg-type] + user="alice", + db="mydb", + password="secret", + ) + + assert isinstance(conn, FakeConn) + assert captured["user"] == "alice" + assert captured["dbname"] == "mydb" + assert captured["password"] == "secret" + assert captured["sslmode"] == "disable" + + client_sock.close() + server_sock.close() + ssl_b.close() + + +def test_connect_raises_on_missing_psycopg(monkeypatch: pytest.MonkeyPatch) -> None: + """connect() raises ImportError when psycopg is not installed.""" + import sys + + monkeypatch.setitem(sys.modules, "psycopg", None) # type: ignore[assignment] + + local_a, local_b = _socketpair() + + with pytest.raises(ImportError, match="psycopg"): + connect(local_b, user="u", db="d", password="p") # type: ignore[arg-type] + + local_a.close() + + +def _make_fake_psycopg_module(captured: dict) -> types.ModuleType: + """Return a fake psycopg module whose connect() captures kwargs into ``captured``.""" + + class FakeConn: + pass + + def fake_connect(**kw: Any) -> FakeConn: + captured.update(kw) + return FakeConn() + + mod = types.ModuleType("psycopg") + mod.connect = fake_connect # type: ignore[attr-defined] + return mod + + +def test_proxy_forwards_local_to_remote() -> None: + """_proxy() forwards bytes written to the local socket to the remote.""" + local_a, local_b = _socketpair() + remote_a, remote_b = _socketpair() + + remote_a.settimeout(2.0) + t = threading.Thread(target=_proxy, args=(local_b, remote_b), daemon=True) + t.start() + + local_a.sendall(b"hello") + received = remote_a.recv(5) + + # Trigger EOF so the proxy thread exits. + local_a.shutdown(socket.SHUT_RDWR) + local_a.close() + remote_a.close() + t.join(timeout=2) + + assert received == b"hello" + + +def test_proxy_forwards_remote_to_local() -> None: + """_proxy() forwards bytes written to the remote socket to the local.""" + local_a, local_b = _socketpair() + remote_a, remote_b = _socketpair() + + local_a.settimeout(2.0) + t = threading.Thread(target=_proxy, args=(local_b, remote_b), daemon=True) + t.start() + + remote_a.sendall(b"world") + received = local_a.recv(5) + + # Trigger EOF so the proxy thread exits. + local_a.shutdown(socket.SHUT_RDWR) + local_a.close() + remote_a.close() + t.join(timeout=2) + + assert received == b"world" + + +def test_proxy_eof_on_local_closes_remote() -> None: + """EOF on the local side causes _proxy() to close the remote socket.""" + local_a, local_b = _socketpair() + remote_a, remote_b = _socketpair() + + remote_a.settimeout(2.0) + t = threading.Thread(target=_proxy, args=(local_b, remote_b), daemon=True) + t.start() + + # Signal EOF from the local driver side. + local_a.shutdown(socket.SHUT_RDWR) + local_a.close() + + # remote_a should receive EOF once remote_b is shut down by the proxy. + eof = remote_a.recv(1) + remote_a.close() + t.join(timeout=2) + + assert eof == b"" + assert not t.is_alive() + + +def test_connect_strips_caller_sslmode(monkeypatch: pytest.MonkeyPatch) -> None: + """connect() replaces any caller-supplied sslmode with 'disable'.""" + captured: dict = {} + mod = _make_fake_psycopg_module(captured) + monkeypatch.setitem(sys.modules, "psycopg", mod) + + ssl_a, ssl_b = _socketpair() + connect(ssl_a, user="u", db="d", sslmode="require") # type: ignore[arg-type] + ssl_b.close() + + assert captured["sslmode"] == "disable" + + +def test_connect_passes_extra_kwargs(monkeypatch: pytest.MonkeyPatch) -> None: + """connect() forwards unrecognised kwargs to psycopg.connect().""" + captured: dict = {} + mod = _make_fake_psycopg_module(captured) + monkeypatch.setitem(sys.modules, "psycopg", mod) + + ssl_a, ssl_b = _socketpair() + connect( # type: ignore[arg-type] + ssl_a, + user="u", + db="d", + connect_timeout=10, + application_name="myapp", + ) + ssl_b.close() + + assert captured["connect_timeout"] == 10 + assert captured["application_name"] == "myapp" + + +def test_connect_cleanup_on_success(monkeypatch: pytest.MonkeyPatch) -> None: + """connect() removes the Unix socket file and tmpdir after a successful connect.""" + captured: dict = {} + mod = _make_fake_psycopg_module(captured) + monkeypatch.setitem(sys.modules, "psycopg", mod) + + ssl_a, ssl_b = _socketpair() + connect(ssl_a, user="u", db="d") # type: ignore[arg-type] + ssl_b.close() + + tmpdir = captured["host"] + socket_path = os.path.join(tmpdir, ".s.PGSQL.5432") + assert not os.path.exists(socket_path) + assert not os.path.exists(tmpdir) + + +def test_no_thread_leak(monkeypatch: pytest.MonkeyPatch) -> None: + """Opening and closing N connections must not permanently increase thread count. + + The fake psycopg actually connects to the Unix socket so _accept_and_proxy + can proceed and both proxy threads wind down naturally when the client + closes its end. + """ + import time + + N = 10 + baseline = threading.active_count() + + for _ in range(N): + + def fake_connect(**kw: Any) -> object: + # Connect to the Unix socket so _accept_and_proxy can accept(), then + # close immediately to send EOF through the proxy, letting both proxy + # threads unwind. + path = os.path.join(kw["host"], f".s.PGSQL.{kw['port']}") + c = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + c.connect(path) + c.close() + return object() + + mod = types.ModuleType("psycopg") + mod.connect = fake_connect # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "psycopg", mod) + + ssl_a, ssl_b = _socketpair() + connect(ssl_a, user="u", db="d") # type: ignore[arg-type] + ssl_b.close() + + # Give daemon threads a moment to observe EOF and exit. + time.sleep(0.2) + assert threading.active_count() <= baseline + 2 + + +def test_connect_cleanup_on_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """connect() removes the Unix socket file and tmpdir even when psycopg.connect raises.""" + captured_host: list[str] = [] + + class FakeConn: + pass + + def failing_connect(**kwargs: Any) -> FakeConn: + captured_host.append(kwargs["host"]) + raise RuntimeError("db unavailable") + + mod = types.ModuleType("psycopg") + mod.connect = failing_connect # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "psycopg", mod) + + ssl_a, ssl_b = _socketpair() + with pytest.raises(RuntimeError, match="db unavailable"): + connect(ssl_a, user="u", db="d") # type: ignore[arg-type] + ssl_b.close() + + tmpdir = captured_host[0] + socket_path = os.path.join(tmpdir, ".s.PGSQL.5432") + assert not os.path.exists(socket_path) + assert not os.path.exists(tmpdir)