cloud-sql-python-connector 1.21.0__py3-none-any.whl → 1.23.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -17,13 +17,16 @@ limitations under the License.
17
17
  from __future__ import annotations
18
18
 
19
19
  import asyncio
20
+ from collections.abc import Callable
20
21
  from functools import partial
21
22
  import logging
22
23
  import os
24
+ import random
23
25
  import socket
24
26
  from threading import Thread
27
+ import time
25
28
  from types import TracebackType
26
- from typing import Any, Callable
29
+ from typing import Any
27
30
 
28
31
  import google.auth
29
32
  from google.auth.credentials import Credentials
@@ -31,19 +34,24 @@ from google.auth.credentials import with_scopes_if_required
31
34
 
32
35
  from google.cloud.sql.connector import asyncpg
33
36
  from google.cloud.sql.connector import pg8000
37
+ from google.cloud.sql.connector import psycopg
34
38
  from google.cloud.sql.connector import pymysql
35
39
  from google.cloud.sql.connector import pytds
36
40
  from google.cloud.sql.connector.client import CloudSQLClient
41
+ from google.cloud.sql.connector.connection_name import ConnectionName
37
42
  from google.cloud.sql.connector.enums import DriverMapping
38
43
  from google.cloud.sql.connector.enums import IPTypes
39
44
  from google.cloud.sql.connector.enums import RefreshStrategy
40
45
  from google.cloud.sql.connector.exceptions import ClosedConnectorError
41
46
  from google.cloud.sql.connector.exceptions import ConnectorLoopError
47
+ from google.cloud.sql.connector.exceptions import IncompatibleDriverError
48
+ from google.cloud.sql.connector.exceptions import ResourceExhaustedError
42
49
  from google.cloud.sql.connector.instance import RefreshAheadCache
43
50
  from google.cloud.sql.connector.lazy import LazyRefreshCache
44
51
  from google.cloud.sql.connector.monitored_cache import MonitoredCache
45
52
  from google.cloud.sql.connector.resolver import DefaultResolver
46
53
  from google.cloud.sql.connector.resolver import DnsResolver
54
+ from google.cloud.sql.connector.sqldata_client import SqlDataClient
47
55
  from google.cloud.sql.connector.utils import format_database_user
48
56
  from google.cloud.sql.connector.utils import generate_keys
49
57
 
@@ -56,6 +64,49 @@ _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
56
64
  _SQLADMIN_HOST_TEMPLATE = "sqladmin.{universe_domain}"
57
65
 
58
66
 
67
+ class SqlDataConnState:
68
+ """Tracks connection state, fallback status, and resource exhaustion cooldown for SQL Data Service."""
69
+
70
+ def __init__(self) -> None:
71
+ self.allowed: bool = True
72
+ self.cooldown_until: float | None = None
73
+ self.backoff_counter: int = 0
74
+ self.last_err: Exception | None = None
75
+
76
+ def is_cooldown_active(self) -> bool:
77
+ """Returns True if the instance connection is in cooldown due to resource exhaustion."""
78
+ return bool(
79
+ self.allowed
80
+ and self.cooldown_until
81
+ and time.time() < self.cooldown_until
82
+ )
83
+
84
+ def record_exhausted(self, err: Exception, base_cooldown: float) -> float:
85
+ """Records a resource exhaustion error, increments backoff counter, and returns the cooldown duration."""
86
+ if self.backoff_counter < 5:
87
+ self.backoff_counter += 1
88
+ backoff = _cooldown_backoff(base_cooldown, self.backoff_counter)
89
+ self.cooldown_until = time.time() + backoff
90
+ self.last_err = err
91
+ return backoff
92
+
93
+ def record_success(self) -> None:
94
+ """Resets cooldown and backoff state on successful communication."""
95
+ self.backoff_counter = 0
96
+ self.cooldown_until = None
97
+ self.last_err = None
98
+
99
+ def record_fallback(self) -> None:
100
+ """Marks SQL Data Service as not allowed for this instance."""
101
+ self.allowed = False
102
+
103
+
104
+ def _cooldown_backoff(base_cooldown: float, attempt: int) -> float:
105
+ multi = 1.618
106
+ exp = float(attempt - 1) + random.random()
107
+ return base_cooldown * (multi**exp)
108
+
109
+
59
110
  class Connector:
60
111
  """Configure and create secure connections to Cloud SQL."""
61
112
 
@@ -73,14 +124,20 @@ class Connector:
73
124
  refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
74
125
  resolver: type[DefaultResolver | DnsResolver] = DefaultResolver,
75
126
  failover_period: int = 30,
127
+ sql_data_endpoint: str | None = None,
128
+ sql_data_stream_timeout: int = 7200,
129
+ resource_exhausted_cooldown_period: float = 5.0,
76
130
  ) -> None:
77
131
  """Initializes a Connector instance.
78
132
 
79
133
  Args:
80
134
  ip_type (str | IPTypes): The default IP address type used to connect to
81
135
  Cloud SQL instances. Can be one of the following:
82
- IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), or
83
- IPTypes.PSC ("PSC"). Default: IPTypes.PUBLIC
136
+ IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"),
137
+ IPTypes.PSC ("PSC"), or IPTypes.SQL_DATA ("SQL_DATA").
138
+ Note: IPTypes.SQL_DATA is not currently supported with the
139
+ asyncpg driver.
140
+ Default: IPTypes.PUBLIC
84
141
 
85
142
  enable_iam_auth (bool): Enables automatic IAM database authentication
86
143
  (Postgres and MySQL) as the default authentication method for all
@@ -125,6 +182,16 @@ class Connector:
125
182
  attempt to check if a failover has occured for a given instance.
126
183
  Must be used with `resolver=DnsResolver` to have any effect.
127
184
  Default: 30
185
+
186
+ sql_data_endpoint (str | None): Endpoint host for SQL Data Service
187
+ calls. If not specified, derived from `sqladmin_api_endpoint`.
188
+ Default: "sqladmin.googleapis.com".
189
+
190
+ sql_data_stream_timeout (int): Timeout in seconds for the SQL Data
191
+ Service gRPC stream. Default: 7200.
192
+
193
+ resource_exhausted_cooldown_period (float): Cooldown period in seconds
194
+ after a ResourceExhausted error. Default: 5.0.
128
195
  """
129
196
  # if refresh_strategy is str, convert to RefreshStrategy enum
130
197
  if isinstance(refresh_strategy, str):
@@ -213,6 +280,22 @@ class Connector:
213
280
  "configured the universe domain explicitly, `googleapis.com` "
214
281
  "is the default."
215
282
  )
283
+ if sql_data_endpoint:
284
+ self._sql_data_endpoint = sql_data_endpoint
285
+ else:
286
+ self._sql_data_endpoint = (
287
+ self._sqladmin_api_endpoint.split("://", 1)[-1].split("/", 1)[0]
288
+ )
289
+ self._sql_data_stream_timeout = sql_data_stream_timeout
290
+ self._resource_exhausted_cooldown_period = (
291
+ resource_exhausted_cooldown_period
292
+ )
293
+ self._sql_data_fallback_cache: set[str] = set()
294
+ self._sql_data_conn_state: dict[str, SqlDataConnState] = {}
295
+ self._sqldata_clients: set[Any] = set()
296
+ self._sqldata_client: SqlDataClient | None = None
297
+
298
+
216
299
 
217
300
  @property
218
301
  def universe_domain(self) -> str:
@@ -259,6 +342,49 @@ class Connector:
259
342
  )
260
343
  return connect_future.result()
261
344
 
345
+ def _get_or_create_cache(
346
+ self,
347
+ conn_name: ConnectionName,
348
+ enable_iam_auth: bool,
349
+ ) -> MonitoredCache:
350
+ assert self._client is not None, "client must be initialized before creating cache"
351
+ assert self._keys is not None, "keys must be initialized before creating cache"
352
+ assert self._resolver is not None, "resolver must be initialized before creating cache"
353
+ if (str(conn_name), enable_iam_auth) in self._cache and not self._cache[
354
+ (str(conn_name), enable_iam_auth)
355
+ ].closed:
356
+ return self._cache[(str(conn_name), enable_iam_auth)]
357
+
358
+ if self._refresh_strategy == RefreshStrategy.LAZY:
359
+ logger.debug(
360
+ f"['{conn_name}']: Refresh strategy is set to lazy refresh"
361
+ )
362
+ cache: LazyRefreshCache | RefreshAheadCache = LazyRefreshCache(
363
+ conn_name,
364
+ self._client,
365
+ self._keys,
366
+ enable_iam_auth,
367
+ )
368
+ else:
369
+ logger.debug(
370
+ f"['{conn_name}']: Refresh strategy is set to backgound refresh"
371
+ )
372
+ cache = RefreshAheadCache(
373
+ conn_name,
374
+ self._client,
375
+ self._keys,
376
+ enable_iam_auth,
377
+ )
378
+ # wrap cache as a MonitoredCache
379
+ monitored_cache = MonitoredCache(
380
+ cache,
381
+ self._failover_period,
382
+ self._resolver,
383
+ )
384
+ logger.debug(f"['{conn_name}']: Connection info added to cache")
385
+ self._cache[(str(conn_name), enable_iam_auth)] = monitored_cache
386
+ return monitored_cache
387
+
262
388
  async def connect_async(
263
389
  self, instance_connection_string: str, driver: str, **kwargs: Any
264
390
  ) -> Any:
@@ -320,48 +446,20 @@ class Connector:
320
446
  if self._resolver is None:
321
447
  self._resolver = self._resolver_cls(client=self._client)
322
448
  enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)
449
+ ip_type = kwargs.pop("ip_type", self._ip_type)
450
+ if isinstance(ip_type, str):
451
+ ip_type = IPTypes._from_str(ip_type)
323
452
 
324
453
  conn_name = await self._resolver.resolve(instance_connection_string)
325
- # Cache entry must exist and not be closed
326
- if (str(conn_name), enable_iam_auth) in self._cache and not self._cache[
327
- (str(conn_name), enable_iam_auth)
328
- ].closed:
329
- monitored_cache = self._cache[(str(conn_name), enable_iam_auth)]
330
- else:
331
- if self._refresh_strategy == RefreshStrategy.LAZY:
332
- logger.debug(
333
- f"['{conn_name}']: Refresh strategy is set to lazy refresh"
334
- )
335
- cache: LazyRefreshCache | RefreshAheadCache = LazyRefreshCache(
336
- conn_name,
337
- self._client,
338
- self._keys,
339
- enable_iam_auth,
340
- )
341
- else:
342
- logger.debug(
343
- f"['{conn_name}']: Refresh strategy is set to backgound refresh"
344
- )
345
- cache = RefreshAheadCache(
346
- conn_name,
347
- self._client,
348
- self._keys,
349
- enable_iam_auth,
350
- )
351
- # wrap cache as a MonitoredCache
352
- monitored_cache = MonitoredCache(
353
- cache,
354
- self._failover_period,
355
- self._resolver,
356
- )
357
- logger.debug(f"['{conn_name}']: Connection info added to cache")
358
- self._cache[(str(conn_name), enable_iam_auth)] = monitored_cache
454
+ if ip_type != IPTypes.SQL_DATA:
455
+ monitored_cache = self._get_or_create_cache(conn_name, enable_iam_auth)
359
456
 
360
457
  connect_func = {
361
458
  "pymysql": pymysql.connect,
362
459
  "pg8000": pg8000.connect,
363
460
  "asyncpg": asyncpg.connect,
364
461
  "pytds": pytds.connect,
462
+ "psycopg": psycopg.connect,
365
463
  }
366
464
 
367
465
  # only accept supported database drivers
@@ -369,11 +467,6 @@ class Connector:
369
467
  connector: Callable = connect_func[driver] # type: ignore
370
468
  except KeyError:
371
469
  raise KeyError(f"Driver '{driver}' is not supported.")
372
-
373
- ip_type = kwargs.pop("ip_type", self._ip_type)
374
- # if ip_type is str, convert to IPTypes enum
375
- if isinstance(ip_type, str):
376
- ip_type = IPTypes._from_str(ip_type)
377
470
  kwargs["timeout"] = kwargs.get("timeout", self._timeout)
378
471
 
379
472
  # Host and ssl options come from the certificates and metadata, so we don't
@@ -382,113 +475,210 @@ class Connector:
382
475
  kwargs.pop("ssl", None)
383
476
  kwargs.pop("port", None)
384
477
 
385
- # attempt to get connection info for Cloud SQL instance
478
+ # attempt to establish connection
386
479
  try:
387
- conn_info = await monitored_cache.connect_info()
388
- # validate driver matches intended database engine
389
- DriverMapping.validate_engine(driver, conn_info.database_version)
390
- preferred_ips = conn_info.get_preferred_ips(ip_type)
391
- except Exception:
392
- # with an error from Cloud SQL Admin API call or IP type, invalidate
393
- # the cache and re-raise the error
394
- await self._remove_cached(str(conn_name), enable_iam_auth)
395
- raise
396
-
397
- targets = []
398
- # If the connector is configured with a custom DNS name, attempt to use
399
- # that DNS name to connect to the instance. Fall back to the metadata IP
400
- # address if the DNS name does not resolve to an IP address.
401
- if conn_info.conn_name.domain_name and isinstance(self._resolver, DnsResolver):
402
- try:
403
- ips = await self._resolver.resolve_a_record(conn_info.conn_name.domain_name)
404
- if ips:
405
- targets.extend(ips)
480
+ if ip_type == IPTypes.SQL_DATA:
481
+ if driver in ASYNC_DRIVERS:
482
+ raise IncompatibleDriverError(
483
+ f"Driver '{driver}' is not supported with ip_type '{ip_type}'."
484
+ )
485
+ state = self._sql_data_conn_state.setdefault(
486
+ str(conn_name), SqlDataConnState()
487
+ )
488
+ if state.is_cooldown_active():
406
489
  logger.debug(
407
- f"['{instance_connection_string}']: Custom DNS name "
408
- f"'{conn_info.conn_name.domain_name}' resolved to '{ips}', "
409
- "using it to connect"
490
+ f"['{conn_name}']: SQL Data Service in cooldown until {state.cooldown_until}"
491
+ )
492
+ raise ResourceExhaustedError(
493
+ "cooldown active", str(conn_name), state.last_err
494
+ )
495
+
496
+ logger.debug(f"['{conn_name}']: Connecting via SQL Data Service tunnel")
497
+ if enable_iam_auth:
498
+ engine = DriverMapping[driver.upper()].value
499
+ formatted_user = format_database_user(
500
+ engine, kwargs["user"]
501
+ )
502
+ if formatted_user != kwargs["user"]:
503
+ logger.debug(
504
+ f"['{instance_connection_string}']: Truncated IAM database username from {kwargs['user']} to {formatted_user}"
505
+ )
506
+ kwargs["user"] = formatted_user
507
+
508
+ if self._sqldata_client is None:
509
+ sqldata_client = SqlDataClient(
510
+ endpoint=self._sql_data_endpoint,
511
+ credentials=self._credentials,
512
+ quota_project=self._quota_project,
513
+ timeout=self._sql_data_stream_timeout,
514
+ )
515
+ self._sqldata_client = sqldata_client
516
+ self._sqldata_clients.add(sqldata_client)
517
+ sqldata_client._on_close_callbacks.append(
518
+ lambda: self._sqldata_clients.discard(sqldata_client)
410
519
  )
411
520
  else:
521
+ sqldata_client = self._sqldata_client
522
+
523
+ def on_resource_exhausted(err: Exception) -> None:
524
+ backoff = state.record_exhausted(
525
+ err, self._resource_exhausted_cooldown_period
526
+ )
412
527
  logger.debug(
413
- f"['{instance_connection_string}']: Custom DNS name "
414
- f"'{conn_info.conn_name.domain_name}' resolved but returned no "
415
- f"entries, using '{preferred_ips}' from instance metadata"
528
+ f"['{conn_name}']: ResourceExhausted occurred, backing off for {backoff:.2f}s "
529
+ f"(attempt {state.backoff_counter})"
416
530
  )
417
- targets.extend(preferred_ips)
418
- except Exception as e: # noqa: BLE001
419
- logger.debug(
420
- f"['{instance_connection_string}']: Custom DNS name "
421
- f"'{conn_info.conn_name.domain_name}' did not resolve to an IP "
422
- f"address: {e}, using '{preferred_ips}' from instance metadata"
423
- )
424
- targets.extend(preferred_ips)
425
- else:
426
- targets.extend(preferred_ips)
427
531
 
428
- # format `user` param for automatic IAM database authn
429
- if enable_iam_auth:
430
- formatted_user = format_database_user(
431
- conn_info.database_version, kwargs["user"]
432
- )
433
- if formatted_user != kwargs["user"]:
434
- logger.debug(
435
- f"['{instance_connection_string}']: Truncated IAM database username from {kwargs['user']} to {formatted_user}"
532
+ def on_success() -> None:
533
+ state.record_success()
534
+
535
+ def on_fallback(name: str) -> None:
536
+ state.record_fallback()
537
+ self._sql_data_fallback_cache.add(name)
538
+
539
+ def is_fallback_cached(name: str) -> bool:
540
+ return not state.allowed or name in self._sql_data_fallback_cache
541
+
542
+ # Defer cache creation and connect_info call
543
+ async def get_conn_info():
544
+ cache = self._get_or_create_cache(conn_name, enable_iam_auth)
545
+ return await cache.connect_info()
546
+
547
+ sock = await sqldata_client.connect(
548
+ instance_connection_name=str(conn_name),
549
+ region=conn_name.region,
550
+ project=conn_name.project,
551
+ get_conn_info=get_conn_info,
552
+ enable_iam_auth=enable_iam_auth,
553
+ on_fallback=on_fallback,
554
+ is_fallback_cached=is_fallback_cached,
555
+ on_resource_exhausted=on_resource_exhausted,
556
+ on_success=on_success,
557
+ connect_timeout=kwargs.get("timeout", self._timeout),
436
558
  )
437
- kwargs["user"] = formatted_user
438
559
 
439
- try:
440
- last_ex = None
441
- for target_ip in targets:
442
- logger.debug(f"['{conn_info.conn_name}']: Connecting to {target_ip}:3307")
560
+ cache = None
561
+ if conn_name.domain_name:
562
+ cache = self._get_or_create_cache(conn_name, enable_iam_auth)
563
+ cache.sockets.append(sock)
564
+
565
+ connect_partial = partial(
566
+ connector,
567
+ "127.0.0.1",
568
+ sock,
569
+ **kwargs,
570
+ )
571
+ try:
572
+ return await self._loop.run_in_executor(None, connect_partial)
573
+ except Exception:
574
+ sock.close()
575
+ if conn_name.domain_name and cache:
576
+ cache._purge_closed_sockets()
577
+ raise
578
+ else:
443
579
  try:
444
- # async drivers are unblocking and can be awaited directly
445
- if driver in ASYNC_DRIVERS:
446
- conn = await connector(
580
+ conn_info = await monitored_cache.connect_info()
581
+ # validate driver matches intended database engine
582
+ DriverMapping.validate_engine(driver, conn_info.database_version)
583
+ preferred_ips = conn_info.get_preferred_ips(ip_type)
584
+ except Exception:
585
+ # with an error from Cloud SQL Admin API call or IP type, invalidate
586
+ # the cache and re-raise the error
587
+ await self._remove_cached(str(conn_name), enable_iam_auth)
588
+ raise
589
+
590
+ targets = []
591
+ if conn_info.conn_name.domain_name and isinstance(self._resolver, DnsResolver):
592
+ try:
593
+ ips = await self._resolver.resolve_a_record(conn_info.conn_name.domain_name)
594
+ if ips:
595
+ targets.extend(ips)
596
+ logger.debug(
597
+ f"['{instance_connection_string}']: Custom DNS name "
598
+ f"'{conn_info.conn_name.domain_name}' resolved to '{ips}', "
599
+ "using it to connect"
600
+ )
601
+ else:
602
+ logger.debug(
603
+ f"['{instance_connection_string}']: Custom DNS name "
604
+ f"'{conn_info.conn_name.domain_name}' resolved but returned no "
605
+ f"entries, using '{preferred_ips}' from instance metadata"
606
+ )
607
+ targets.extend(preferred_ips)
608
+ except Exception as e: # noqa: BLE001
609
+ logger.debug(
610
+ f"['{instance_connection_string}']: Custom DNS name "
611
+ f"'{conn_info.conn_name.domain_name}' did not resolve to an IP "
612
+ f"address: {e}, using '{preferred_ips}' from instance metadata"
613
+ )
614
+ targets.extend(preferred_ips)
615
+ else:
616
+ targets.extend(preferred_ips)
617
+ # format `user` param for automatic IAM database authn
618
+ if enable_iam_auth:
619
+ formatted_user = format_database_user(
620
+ conn_info.database_version, kwargs["user"]
621
+ )
622
+ if formatted_user != kwargs["user"]:
623
+ logger.debug(
624
+ f"['{instance_connection_string}']: Truncated IAM database username from {kwargs['user']} to {formatted_user}"
625
+ )
626
+ kwargs["user"] = formatted_user
627
+
628
+ last_ex = None
629
+ for target_ip in targets:
630
+ logger.debug(f"['{conn_info.conn_name}']: Connecting to {target_ip}:3307")
631
+ try:
632
+ # async drivers are unblocking and can be awaited directly
633
+ if driver in ASYNC_DRIVERS:
634
+ conn = await connector(
635
+ target_ip,
636
+ await conn_info.create_ssl_context(enable_iam_auth),
637
+ **kwargs,
638
+ )
639
+ last_ex = None
640
+ return conn
641
+
642
+ # Create socket with SSLContext for sync drivers
643
+ ctx = await conn_info.create_ssl_context(enable_iam_auth)
644
+ raw_sock = socket.create_connection((target_ip, SERVER_PROXY_PORT))
645
+ try:
646
+ sock = ctx.wrap_socket(
647
+ raw_sock,
648
+ server_hostname=target_ip,
649
+ )
650
+ except Exception:
651
+ raw_sock.close()
652
+ raise
653
+
654
+ # If this connection was opened using a domain name, then store it
655
+ # for later in case we need to forcibly close it on failover.
656
+ if conn_info.conn_name.domain_name:
657
+ monitored_cache.sockets.append(sock)
658
+ # Synchronous drivers are blocking and run using executor
659
+ connect_partial = partial(
660
+ connector,
447
661
  target_ip,
448
- await conn_info.create_ssl_context(enable_iam_auth),
662
+ sock,
449
663
  **kwargs,
450
664
  )
665
+ conn = await self._loop.run_in_executor(None, connect_partial)
451
666
  last_ex = None
452
667
  return conn
453
-
454
- # Create socket with SSLContext for sync drivers
455
- ctx = await conn_info.create_ssl_context(enable_iam_auth)
456
- raw_sock = socket.create_connection((target_ip, SERVER_PROXY_PORT))
457
- try:
458
- sock = ctx.wrap_socket(
459
- raw_sock,
460
- server_hostname=target_ip,
668
+ except Exception as e: # noqa: BLE001
669
+ logger.debug(
670
+ f"['{conn_info.conn_name}']: Connection to {target_ip} failed: {e}"
461
671
  )
462
- except Exception:
463
- raw_sock.close()
464
- raise
465
-
466
- # If this connection was opened using a domain name, then store it
467
- # for later in case we need to forcibly close it on failover.
468
- if conn_info.conn_name.domain_name:
469
- monitored_cache.sockets.append(sock)
470
- # Synchronous drivers are blocking and run using executor
471
- connect_partial = partial(
472
- connector,
473
- target_ip,
474
- sock,
475
- **kwargs,
476
- )
477
- conn = await self._loop.run_in_executor(None, connect_partial)
478
- last_ex = None
479
- return conn
480
- except Exception as e: # noqa: BLE001
481
- logger.debug(
482
- f"['{conn_info.conn_name}']: Connection to {target_ip} failed: {e}"
483
- )
484
- last_ex = e
672
+ last_ex = e
485
673
 
486
- if last_ex:
487
- raise last_ex
674
+ if last_ex:
675
+ raise last_ex
488
676
 
489
677
  except Exception:
490
678
  # with any exception, we attempt a force refresh, then throw the error
491
- await monitored_cache.force_refresh()
679
+ cached_entry = self._cache.get((str(conn_name), enable_iam_auth))
680
+ if cached_entry:
681
+ await cached_entry.force_refresh()
492
682
  raise
493
683
 
494
684
  async def _remove_cached(
@@ -536,8 +726,11 @@ class Connector:
536
726
  close_future = asyncio.run_coroutine_threadsafe(
537
727
  self.close_async(), loop=self._loop
538
728
  )
539
- # Will attempt to safely shut down tasks for 3s
540
- close_future.result(timeout=3)
729
+ try:
730
+ # Will attempt to safely shut down tasks for 3s
731
+ close_future.result(timeout=3)
732
+ except Exception as e: # noqa: BLE001
733
+ logger.error(f"Error during close_async: {e}")
541
734
  # if background thread exists for Connector, clean it up
542
735
  if self._thread:
543
736
  if self._loop.is_running():
@@ -550,9 +743,14 @@ class Connector:
550
743
  """Helper function to cancel the cache's tasks
551
744
  and close aiohttp.ClientSession."""
552
745
  self._closed = True
746
+ self._sqldata_client = None
553
747
  if self._client:
554
748
  await self._client.close()
555
- await asyncio.gather(*[cache.close() for cache in self._cache.values()])
749
+ await asyncio.gather(
750
+ *[cache.close() for cache in self._cache.values()],
751
+ *[client.close() for client in list(self._sqldata_clients)],
752
+ return_exceptions=True,
753
+ )
556
754
 
557
755
 
558
756
  async def create_async_connector(
@@ -568,6 +766,9 @@ async def create_async_connector(
568
766
  refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
569
767
  resolver: type[DefaultResolver | DnsResolver] = DefaultResolver,
570
768
  failover_period: int = 30,
769
+ sql_data_endpoint: str | None = None,
770
+ sql_data_stream_timeout: int = 7200,
771
+ resource_exhausted_cooldown_period: float = 5.0,
571
772
  ) -> Connector:
572
773
  """Helper function to create Connector object for asyncio connections.
573
774
 
@@ -577,8 +778,11 @@ async def create_async_connector(
577
778
  Args:
578
779
  ip_type (str | IPTypes): The default IP address type used to connect to
579
780
  Cloud SQL instances. Can be one of the following:
580
- IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), or
581
- IPTypes.PSC ("PSC"). Default: IPTypes.PUBLIC
781
+ IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"),
782
+ IPTypes.PSC ("PSC"), or IPTypes.SQL_DATA ("SQL_DATA").
783
+ Note: IPTypes.SQL_DATA is not currently supported with the
784
+ asyncpg driver.
785
+ Default: IPTypes.PUBLIC
582
786
 
583
787
  enable_iam_auth (bool): Enables automatic IAM database authentication
584
788
  (Postgres and MySQL) as the default authentication method for all
@@ -624,6 +828,16 @@ async def create_async_connector(
624
828
  Must be used with `resolver=DnsResolver` to have any effect.
625
829
  Default: 30
626
830
 
831
+ sql_data_endpoint (str | None): Endpoint host for SQL Data Service
832
+ calls. If not specified, derived from `sqladmin_api_endpoint`.
833
+ Default: "sqladmin.googleapis.com".
834
+
835
+ sql_data_stream_timeout (int): Timeout in seconds for the SQL Data
836
+ Service gRPC stream. Default: 7200.
837
+
838
+ resource_exhausted_cooldown_period (float): Cooldown period in seconds
839
+ after a ResourceExhausted error. Default: 5.0.
840
+
627
841
  Returns:
628
842
  A Connector instance configured with running event loop.
629
843
  """
@@ -643,4 +857,7 @@ async def create_async_connector(
643
857
  refresh_strategy=refresh_strategy,
644
858
  resolver=resolver,
645
859
  failover_period=failover_period,
860
+ sql_data_endpoint=sql_data_endpoint,
861
+ sql_data_stream_timeout=sql_data_stream_timeout,
862
+ resource_exhausted_cooldown_period=resource_exhausted_cooldown_period,
646
863
  )
@@ -41,6 +41,7 @@ class IPTypes(Enum):
41
41
  PUBLIC = "PRIMARY"
42
42
  PRIVATE = "PRIVATE"
43
43
  PSC = "PSC"
44
+ SQL_DATA = "SQL_DATA"
44
45
 
45
46
  @classmethod
46
47
  def _missing_(cls, value: object) -> None:
@@ -54,6 +55,8 @@ class IPTypes(Enum):
54
55
  """Convert IP type from a str into IPTypes."""
55
56
  if ip_type_str.upper() == "PUBLIC":
56
57
  ip_type_str = "PRIMARY"
58
+ elif ip_type_str.upper() in ("SQLDATA", "SQL_DATA"):
59
+ ip_type_str = "SQL_DATA"
57
60
  return cls(ip_type_str.upper())
58
61
 
59
62
 
@@ -62,6 +65,7 @@ class DriverMapping(Enum):
62
65
 
63
66
  ASYNCPG = "POSTGRES"
64
67
  PG8000 = "POSTGRES" # noqa: PIE796
68
+ PSYCOPG = "POSTGRES" # noqa: PIE796
65
69
  PYMYSQL = "MYSQL"
66
70
  PYTDS = "SQLSERVER"
67
71