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)