Skip to content

Commit b34e9c6

Browse files
authored
feat: Add PSC DNS and Global Write Endpoint support to Python Connector (#1424)
The Private Service Connect (PSC) DNS Name Resolution & Fallback feature enables the connector to securely connect to Cloud SQL database instances using Private Service Connect (PSC) DNS and custom domain names (CNAMEs). The connector can now dial PSC instances using its DNS hostname (e.g., 0123456789ab.fedcba9876543.us-central1.sql-psc.goog) or a custom domain CNAME pointing to the instance PSC DNS hostname. The connectors will dynamically resolve this hostname to its real Cloud SQL instance connection name via the SQL Admin API. See also: GoogleCloudPlatform/cloud-sql-go-connector#1106
1 parent 53b3bca commit b34e9c6

34 files changed

Lines changed: 607 additions & 220 deletions

‎.gitignore‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,7 @@ dist/
1010
sponge_log.xml
1111
.envrc
1212
*.iml
13+
.mypy_cache/
14+
.nox/
15+
.pytest_cache/
16+
.ruff_cache/

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,11 @@
2323
from google.cloud.sql.connector.version import __version__
2424

2525
__all__ = [
26-
"__version__",
27-
"create_async_connector",
2826
"Connector",
2927
"DefaultResolver",
3028
"DnsResolver",
3129
"IPTypes",
3230
"RefreshStrategy",
31+
"__version__",
32+
"create_async_connector",
3333
]

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

Lines changed: 61 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import asyncio
1818
import datetime
1919
import logging
20-
from typing import Any, Optional, TYPE_CHECKING
20+
from typing import Any, TYPE_CHECKING
2121

2222
import aiohttp
2323
from cryptography.hazmat.backends import default_backend
@@ -42,7 +42,7 @@
4242
logger = logging.getLogger(name=__name__)
4343

4444

45-
def _format_user_agent(driver: Optional[str], custom: Optional[str]) -> str:
45+
def _format_user_agent(driver: str | None, custom: str | None) -> str:
4646
agent = f"{USER_AGENT}+{driver}" if driver else USER_AGENT
4747
if custom and isinstance(custom, str):
4848
agent = f"{agent} {custom}"
@@ -52,12 +52,12 @@ def _format_user_agent(driver: Optional[str], custom: Optional[str]) -> str:
5252
class CloudSQLClient:
5353
def __init__(
5454
self,
55-
sqladmin_api_endpoint: Optional[str],
56-
quota_project: Optional[str],
55+
sqladmin_api_endpoint: str | None,
56+
quota_project: str | None,
5757
credentials: Credentials,
58-
client: Optional[aiohttp.ClientSession] = None,
59-
driver: Optional[str] = None,
60-
user_agent: Optional[str] = None,
58+
client: aiohttp.ClientSession | None = None,
59+
driver: str | None = None,
60+
user_agent: str | None = None,
6161
) -> None:
6262
"""Establishes the client to be used for Cloud SQL Admin API requests.
6363
@@ -135,7 +135,7 @@ async def _get_metadata(
135135
if message:
136136
resp.reason = message
137137
# skip, raise_for_status will catch all errors in finally block
138-
except Exception:
138+
except Exception: # noqa: BLE001, S110
139139
pass
140140
finally:
141141
resp.raise_for_status()
@@ -146,7 +146,7 @@ async def _get_metadata(
146146
)
147147

148148
ip_addresses = (
149-
{ip["type"]: ip["ipAddress"] for ip in ret_dict["ipAddresses"]}
149+
{ip["type"]: [ip["ipAddress"]] for ip in ret_dict["ipAddresses"]}
150150
if "ipAddresses" in ret_dict
151151
else {}
152152
)
@@ -156,27 +156,70 @@ async def _get_metadata(
156156
if ret_dict.get("pscEnabled"):
157157
# Find PSC instance DNS name in the dns_names field
158158
psc_dns_names = [
159-
d["name"]
159+
d["name"].rstrip(".")
160160
for d in ret_dict.get("dnsNames", [])
161161
if d["connectionType"] == "PRIVATE_SERVICE_CONNECT"
162162
and d["dnsScope"] == "INSTANCE"
163163
]
164-
dns_name = psc_dns_names[0] if psc_dns_names else None
164+
# Sort: .sql-psc.goog first
165+
psc_dns_names.sort(key=lambda x: x.endswith(".sql-psc.goog"), reverse=True)
165166

166167
# Fall back do dns_name field if dns_names is not set
167-
if dns_name is None:
168+
if not psc_dns_names:
168169
dns_name = ret_dict.get("dnsName", None)
170+
if dns_name:
171+
psc_dns_names = [dns_name.rstrip(".")]
169172

170-
# Remove trailing period from DNS name. Required for SSL in Python
171-
if dns_name:
172-
ip_addresses["PSC"] = dns_name.rstrip(".")
173+
if psc_dns_names:
174+
ip_addresses["PSC"] = psc_dns_names
173175

174176
return {
175177
"ip_addresses": ip_addresses,
176178
"server_ca_cert": ret_dict["serverCaCert"]["cert"],
177179
"database_version": ret_dict["databaseVersion"],
178180
}
179181

182+
async def resolve_connect_settings(
183+
self,
184+
dns_name: str,
185+
location: str,
186+
) -> dict[str, Any]:
187+
"""Asynchronously calls the resolveConnectSettings endpoint to resolve a
188+
PSC DNS name to a connection name.
189+
190+
Args:
191+
dns_name (str): The DNS name of the Cloud SQL instance.
192+
location (str): The region/location of the instance.
193+
194+
Returns:
195+
A dictionary containing the resolve response (e.g. connectionName).
196+
"""
197+
# before making Cloud SQL Admin API calls, refresh creds if required
198+
if self._credentials.token_state != TokenState.FRESH:
199+
self._credentials.refresh(requests.Request())
200+
201+
headers = {
202+
"Authorization": f"Bearer {self._credentials.token}",
203+
}
204+
205+
url = f"{self._sqladmin_api_endpoint}/sql/{API_VERSION}/locations/{location}/dns/{dns_name}:resolveConnectSettings"
206+
207+
resp = await self._client.get(url, headers=headers)
208+
if resp.status >= 500:
209+
resp = await retry_50x(self._client.get, url, headers=headers)
210+
try:
211+
ret_dict = await resp.json()
212+
if resp.status >= 400:
213+
message = ret_dict.get("error", {}).get("message")
214+
if message:
215+
resp.reason = message
216+
except Exception: # noqa: BLE001, S110
217+
pass
218+
finally:
219+
resp.raise_for_status()
220+
221+
return ret_dict
222+
180223
async def _get_ephemeral(
181224
self,
182225
project: str,
@@ -223,7 +266,7 @@ async def _get_ephemeral(
223266
if message:
224267
resp.reason = message
225268
# skip, raise_for_status will catch all errors in finally block
226-
except Exception:
269+
except Exception: # noqa: BLE001, S110
227270
pass
228271
finally:
229272
resp.raise_for_status()
@@ -244,8 +287,7 @@ async def _get_ephemeral(
244287
# Ref: https://github.com/googleapis/google-auth-library-python/blob/49a5ff7411a2ae4d32a7d11700f9f961c55406a9/google/auth/_helpers.py#L93-L99
245288
token_expiration = token_expiration.replace(tzinfo=datetime.timezone.utc)
246289

247-
if expiration > token_expiration:
248-
expiration = token_expiration
290+
expiration = min(expiration, token_expiration)
249291
return ephemeral_cert, expiration
250292

251293
async def get_connection_info(
@@ -274,7 +316,7 @@ async def get_connection_info(
274316
"""
275317
priv_key, pub_key = await keys
276318
# before making Cloud SQL Admin API calls, refresh creds if required
277-
if not self._credentials.token_state == TokenState.FRESH:
319+
if self._credentials.token_state != TokenState.FRESH:
278320
self._credentials.refresh(requests.Request())
279321

280322
metadata_task = asyncio.create_task(

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

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
from dataclasses import dataclass
1919
import logging
2020
import ssl
21-
from typing import Any, Optional, TYPE_CHECKING
21+
from typing import Any, TYPE_CHECKING
2222

2323
from google.cloud.sql.connector.connection_name import ConnectionName
2424
from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
@@ -67,7 +67,7 @@ class ConnectionInfo:
6767
ip_addrs: dict[str, Any]
6868
database_version: str
6969
expiration: datetime.datetime
70-
context: Optional[ssl.SSLContext] = None
70+
context: ssl.SSLContext | None = None
7171

7272
async def create_ssl_context(self, enable_iam_auth: bool = False) -> ssl.SSLContext:
7373
"""Constructs a SSL/TLS context for the given connection info.
@@ -121,6 +121,17 @@ def get_preferred_ip(self, ip_type: IPTypes) -> str:
121121
"""Returns the first IP address for the instance, according to the preference
122122
supplied by ip_type. If no IP addressess with the given preference are found,
123123
an error is raised."""
124+
if ip_type.value in self.ip_addrs:
125+
return self.ip_addrs[ip_type.value][0]
126+
raise CloudSQLIPTypeError(
127+
"Cloud SQL instance does not have any IP addresses matching "
128+
f"preference: {ip_type.value}"
129+
)
130+
131+
def get_preferred_ips(self, ip_type: IPTypes) -> list[str]:
132+
"""Returns all IP addresses for the instance, according to the preference
133+
supplied by ip_type. If no IP addressess with the given preference are found,
134+
an error is raised."""
124135
if ip_type.value in self.ip_addrs:
125136
return self.ip_addrs[ip_type.value]
126137
raise CloudSQLIPTypeError(

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
# Instance connection name is the format <PROJECT>:<REGION>:<INSTANCE_NAME>
1919
# Additionally, we have to support legacy "domain-scoped" projects
2020
# (e.g. "google.com:PROJECT")
21-
CONN_NAME_REGEX = re.compile(("([^:]+(:[^:]+)?):([^:]+):([^:]+)"))
21+
CONN_NAME_REGEX = re.compile("([^:]+(:[^:]+)?):([^:]+):([^:]+)")
2222
# The domain name pattern in accordance with RFC 1035, RFC 1123 and RFC 2181.
2323
DOMAIN_NAME_REGEX = re.compile(
2424
r"^(?:[_a-z0-9](?:[_a-z0-9-]{0,61}[a-z0-9])?\.)+(?:[a-z](?:[a-z0-9-]{0,61}[a-z0-9])?)?$"
@@ -48,9 +48,7 @@ def get_connection_string(self) -> str:
4848

4949

5050
def _is_valid_domain(domain_name: str) -> bool:
51-
if DOMAIN_NAME_REGEX.fullmatch(domain_name) is None:
52-
return False
53-
return True
51+
return DOMAIN_NAME_REGEX.fullmatch(domain_name) is not None
5452

5553

5654
def _parse_connection_name(connection_name: str) -> ConnectionName:

0 commit comments

Comments
 (0)