Skip to content

Commit d4cc960

Browse files
kgala2kgala2
committed
feat: add support for psycopg
1 parent bcc1cbf commit d4cc960

6 files changed

Lines changed: 526 additions & 0 deletions

File tree

‎google/cloud/sql/connector/connector.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131

3232
from google.cloud.sql.connector import asyncpg
3333
from google.cloud.sql.connector import pg8000
34+
from google.cloud.sql.connector import psycopg
3435
from google.cloud.sql.connector import pymysql
3536
from google.cloud.sql.connector import pytds
3637
from google.cloud.sql.connector.client import CloudSQLClient
@@ -362,6 +363,7 @@ async def connect_async(
362363
"pg8000": pg8000.connect,
363364
"asyncpg": asyncpg.connect,
364365
"pytds": pytds.connect,
366+
"psycopg": psycopg.connect,
365367
}
366368

367369
# only accept supported database drivers

‎google/cloud/sql/connector/enums.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ class DriverMapping(Enum):
6262

6363
ASYNCPG = "POSTGRES"
6464
PG8000 = "POSTGRES" # noqa: PIE796
65+
PSYCOPG = "POSTGRES" # noqa: PIE796
6566
PYMYSQL = "MYSQL"
6667
PYTDS = "SQLSERVER"
6768

‎google/cloud/sql/connector/psycopg.py‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import logging
16+
import os
17+
import socket
18+
import ssl
19+
import tempfile
20+
import threading
21+
from typing import Any, TYPE_CHECKING
22+
23+
if TYPE_CHECKING:
24+
import psycopg
25+
26+
logger = logging.getLogger(name=__name__)
27+
28+
_CHUNK_SIZE = 8 * 1024 # bytes per recv() call inside the proxy forwarding loop
29+
30+
31+
def _proxy(local: socket.socket, remote: "ssl.SSLSocket") -> None:
32+
"""Bidirectionally proxy bytes between a local Unix socket and a remote
33+
SSL socket.
34+
35+
Spawns one daemon thread for the remote→local direction and runs the
36+
local→remote direction in the calling thread. Blocks until the calling
37+
thread's direction reaches EOF or a socket error, at which point both
38+
sockets are closed so the other thread also unblocks and exits.
39+
40+
Args:
41+
local: The Unix domain socket connected to the database driver.
42+
remote: The SSL socket connected to the Cloud SQL proxy server.
43+
"""
44+
def forward(src: Any, dst: Any) -> None:
45+
buf = bytearray(_CHUNK_SIZE)
46+
view = memoryview(buf)
47+
try:
48+
while True:
49+
n = src.recv_into(view)
50+
if n == 0:
51+
logger.debug("psycopg proxy: EOF on %s, closing both sockets", src)
52+
break
53+
dst.sendall(view[:n])
54+
except (OSError, ssl.SSLError) as e:
55+
logger.debug("psycopg proxy: socket error on %s: %s", src, e)
56+
finally:
57+
# Close both ends so the sibling thread also unblocks.
58+
for s in (local, remote):
59+
try:
60+
s.shutdown(socket.SHUT_RDWR)
61+
except OSError:
62+
pass
63+
try:
64+
s.close()
65+
except OSError:
66+
pass
67+
68+
threading.Thread(target=forward, args=(remote, local), daemon=True).start()
69+
forward(local, remote) # run in calling thread rather than spawning a third
70+
71+
72+
def connect(
73+
ip_address: str, remote_sock: "ssl.SSLSocket", **kwargs: Any
74+
) -> "psycopg.Connection":
75+
"""Create a psycopg DBAPI connection object.
76+
77+
Because psycopg does not accept a pre-connected socket, this function
78+
creates a temporary Unix domain socket, tells psycopg to connect there,
79+
and runs a background proxy that forwards bytes between that socket and
80+
the already-established Cloud SQL TLS connection.
81+
82+
Args:
83+
ip_address (str): IP address of the Cloud SQL instance.
84+
remote_sock (ssl.SSLSocket): SSL/TLS secure socket stream connected to the
85+
Cloud SQL proxy server.
86+
87+
Returns:
88+
psycopg.Connection: A psycopg Connection object for the Cloud SQL instance.
89+
"""
90+
try:
91+
import psycopg
92+
except ImportError:
93+
raise ImportError(
94+
'Unable to import module "psycopg." Please install and try again.'
95+
)
96+
97+
tmpdir = tempfile.mkdtemp()
98+
socket_path = os.path.join(tmpdir, ".s.PGSQL.5432")
99+
logger.debug("psycopg: created Unix socket at %s", socket_path)
100+
101+
local_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
102+
local_sock.bind(socket_path)
103+
local_sock.listen(1)
104+
105+
def _accept_and_proxy() -> None:
106+
"""Accept one connection then proxy bytes until the connection closes."""
107+
try:
108+
unix_conn, _ = local_sock.accept()
109+
local_sock.close()
110+
logger.debug("psycopg proxy: accepted connection, starting proxy")
111+
except OSError as e:
112+
logger.debug("psycopg proxy: accept failed: %s", e)
113+
try:
114+
remote_sock.close()
115+
except OSError:
116+
pass
117+
return
118+
_proxy(unix_conn, remote_sock)
119+
120+
threading.Thread(target=_accept_and_proxy, daemon=True).start()
121+
122+
user = kwargs.pop("user")
123+
db = kwargs.pop("db")
124+
passwd = kwargs.pop("password", None)
125+
# SSL is already handled by the underlying SSLSocket; disable it on the
126+
# Unix socket so psycopg does not attempt a second TLS handshake.
127+
kwargs.pop("sslmode", None)
128+
timeout = kwargs.pop("timeout", None)
129+
if timeout is not None:
130+
kwargs["connect_timeout"] = int(timeout)
131+
132+
logger.debug("psycopg: connecting as user=%s dbname=%s", user, db)
133+
try:
134+
conn = psycopg.connect(
135+
user=user,
136+
dbname=db,
137+
password=passwd,
138+
host=tmpdir,
139+
port=5432,
140+
sslmode="disable",
141+
**kwargs,
142+
)
143+
logger.debug("psycopg: connection established")
144+
return conn
145+
except Exception as e:
146+
logger.debug("psycopg: connection failed: %s", e)
147+
# psycopg never connected (or failed mid-handshake); close the server
148+
# socket so the proxy thread unblocks and exits cleanly.
149+
try:
150+
local_sock.close()
151+
except OSError:
152+
pass
153+
try:
154+
remote_sock.close()
155+
except OSError:
156+
pass
157+
raise
158+
finally:
159+
# The socket file and its parent directory are only needed during the
160+
# initial connect() call; remove them now regardless of outcome.
161+
try:
162+
os.remove(socket_path)
163+
except OSError:
164+
pass
165+
try:
166+
os.rmdir(tmpdir)
167+
except OSError:
168+
pass

‎tests/system/test_psycopg_connection.py‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
import asyncio
17+
import time
18+
import pytest
19+
import psutil
20+
from google.cloud.sql.connector import Connector
21+
22+
# These will be set from environment variables or default to our test instance
23+
INSTANCE_CONNECTION_NAME = os.getenv(
24+
"DB_CONNECTION_NAME", "galakp-playground:us-east7:pg-us-east7-psycopg"
25+
)
26+
DB_USER = os.getenv("DB_USER", "postgres")
27+
DB_PASSWORD = os.getenv("DB_PASSWORD", "SuperPass123!")
28+
DB_NAME = os.getenv("DB_NAME", "postgres")
29+
30+
31+
def test_system_psycopg_resource_leak() -> None:
32+
"""Benchmark test to verify no resource leaks (threads, FDs, memory) between iteration 20 and 100."""
33+
print("\nStarting resource leak benchmark...")
34+
35+
process = psutil.Process(os.getpid())
36+
37+
def get_metrics():
38+
return {
39+
"threads": process.num_threads(),
40+
"fds": process.num_fds(),
41+
"rss_mb": process.memory_info().rss / (1024 * 1024),
42+
}
43+
44+
warmup_iterations = 20
45+
total_iterations = 100
46+
47+
baseline_metrics = None
48+
active_final_metrics = None
49+
50+
with Connector() as connector:
51+
for i in range(1, total_iterations + 1):
52+
conn = connector.connect(
53+
INSTANCE_CONNECTION_NAME,
54+
"psycopg",
55+
user=DB_USER,
56+
password=DB_PASSWORD,
57+
db=DB_NAME,
58+
)
59+
cursor = conn.cursor()
60+
cursor.execute("SELECT 1;")
61+
cursor.fetchone()
62+
cursor.close()
63+
conn.close()
64+
65+
if i == warmup_iterations:
66+
time.sleep(0.5)
67+
baseline_metrics = get_metrics()
68+
print(f"Baseline Metrics (Iteration {i}): {baseline_metrics}")
69+
70+
if i == total_iterations:
71+
time.sleep(0.5)
72+
active_final_metrics = get_metrics()
73+
print(f"Active Final Metrics (Iteration {i}): {active_final_metrics}")
74+
75+
if i % 10 == 0 and warmup_iterations < i < total_iterations:
76+
current_metrics = get_metrics()
77+
print(f"Iteration {i:3d}/{total_iterations}: {current_metrics}")
78+
79+
# Post-close metrics
80+
time.sleep(1)
81+
post_close_metrics = get_metrics()
82+
print(f"Post-Close Metrics: {post_close_metrics}")
83+
84+
assert baseline_metrics is not None
85+
assert active_final_metrics is not None
86+
87+
# Assertions: compare Active Final (100) vs Baseline (20)
88+
# Threads should not grow
89+
assert active_final_metrics["threads"] <= baseline_metrics["threads"] + 1, f"Thread leak: {baseline_metrics} -> {active_final_metrics}"
90+
# FDs should not grow
91+
assert active_final_metrics["fds"] <= baseline_metrics["fds"] + 1, f"FD leak: {baseline_metrics} -> {active_final_metrics}"
92+
# Memory growth should be minimal (allow < 5MB growth for minor fragmentation)
93+
assert active_final_metrics["rss_mb"] <= baseline_metrics["rss_mb"] + 5, f"Memory leak: {baseline_metrics} -> {active_final_metrics}"
94+
95+
print("Resource leak benchmark passed successfully.")
96+
97+
98+
def test_system_psycopg_basic() -> None:
99+
"""Basic system test to verify connection and query."""
100+
print(f"\nConnecting to {INSTANCE_CONNECTION_NAME}...")
101+
with Connector() as connector:
102+
conn = connector.connect(
103+
INSTANCE_CONNECTION_NAME,
104+
"psycopg",
105+
user=DB_USER,
106+
password=DB_PASSWORD,
107+
db=DB_NAME,
108+
)
109+
110+
cursor = conn.cursor()
111+
cursor.execute("SELECT version();")
112+
result = cursor.fetchone()
113+
print(f"Database version: {result[0]}")
114+
assert result is not None
115+
cursor.close()
116+
conn.close()
117+
print("Connection closed successfully.")
118+
119+
120+
121+
122+
123+
def test_system_psycopg_to_thread() -> None:
124+
"""Verify that running sync connect in asyncio.to_thread works."""
125+
print(f"\nConnecting via asyncio.to_thread to {INSTANCE_CONNECTION_NAME}...")
126+
127+
async def run_connect():
128+
with Connector() as connector:
129+
# Run the blocking connector.connect in a thread
130+
conn = await asyncio.to_thread(
131+
connector.connect,
132+
INSTANCE_CONNECTION_NAME,
133+
"psycopg",
134+
user=DB_USER,
135+
password=DB_PASSWORD,
136+
db=DB_NAME,
137+
)
138+
cursor = conn.cursor()
139+
cursor.execute("SELECT version();")
140+
result = cursor.fetchone()
141+
print(f"Database version (to_thread): {result[0]}")
142+
assert result is not None
143+
cursor.close()
144+
conn.close()
145+
146+
asyncio.run(run_connect())
147+
print("to_thread connection closed successfully.")

0 commit comments

Comments
 (0)