1717import asyncio
1818import datetime
1919import logging
20- from typing import Any , Optional , TYPE_CHECKING
20+ from typing import Any , TYPE_CHECKING
2121
2222import aiohttp
2323from cryptography .hazmat .backends import default_backend
4242logger = 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:
5252class 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 (
0 commit comments