cloud-sql-python-connector 1.22.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
@@ -35,16 +38,20 @@ from google.cloud.sql.connector import psycopg
35
38
  from google.cloud.sql.connector import pymysql
36
39
  from google.cloud.sql.connector import pytds
37
40
  from google.cloud.sql.connector.client import CloudSQLClient
41
+ from google.cloud.sql.connector.connection_name import ConnectionName
38
42
  from google.cloud.sql.connector.enums import DriverMapping
39
43
  from google.cloud.sql.connector.enums import IPTypes
40
44
  from google.cloud.sql.connector.enums import RefreshStrategy
41
45
  from google.cloud.sql.connector.exceptions import ClosedConnectorError
42
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
43
49
  from google.cloud.sql.connector.instance import RefreshAheadCache
44
50
  from google.cloud.sql.connector.lazy import LazyRefreshCache
45
51
  from google.cloud.sql.connector.monitored_cache import MonitoredCache
46
52
  from google.cloud.sql.connector.resolver import DefaultResolver
47
53
  from google.cloud.sql.connector.resolver import DnsResolver
54
+ from google.cloud.sql.connector.sqldata_client import SqlDataClient
48
55
  from google.cloud.sql.connector.utils import format_database_user
49
56
  from google.cloud.sql.connector.utils import generate_keys
50
57
 
@@ -57,6 +64,49 @@ _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
57
64
  _SQLADMIN_HOST_TEMPLATE = "sqladmin.{universe_domain}"
58
65
 
59
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
+
60
110
  class Connector:
61
111
  """Configure and create secure connections to Cloud SQL."""
62
112
 
@@ -74,14 +124,20 @@ class Connector:
74
124
  refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
75
125
  resolver: type[DefaultResolver | DnsResolver] = DefaultResolver,
76
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,
77
130
  ) -> None:
78
131
  """Initializes a Connector instance.
79
132
 
80
133
  Args:
81
134
  ip_type (str | IPTypes): The default IP address type used to connect to
82
135
  Cloud SQL instances. Can be one of the following:
83
- IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), or
84
- 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
85
141
 
86
142
  enable_iam_auth (bool): Enables automatic IAM database authentication
87
143
  (Postgres and MySQL) as the default authentication method for all
@@ -126,6 +182,16 @@ class Connector:
126
182
  attempt to check if a failover has occured for a given instance.
127
183
  Must be used with `resolver=DnsResolver` to have any effect.
128
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.
129
195
  """
130
196
  # if refresh_strategy is str, convert to RefreshStrategy enum
131
197
  if isinstance(refresh_strategy, str):
@@ -214,6 +280,22 @@ class Connector:
214
280
  "configured the universe domain explicitly, `googleapis.com` "
215
281
  "is the default."
216
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
+
217
299
 
218
300
  @property
219
301
  def universe_domain(self) -> str:
@@ -260,6 +342,49 @@ class Connector:
260
342
  )
261
343
  return connect_future.result()
262
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
+
263
388
  async def connect_async(
264
389
  self, instance_connection_string: str, driver: str, **kwargs: Any
265
390
  ) -> Any:
@@ -321,42 +446,13 @@ class Connector:
321
446
  if self._resolver is None:
322
447
  self._resolver = self._resolver_cls(client=self._client)
323
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)
324
452
 
325
453
  conn_name = await self._resolver.resolve(instance_connection_string)
326
- # Cache entry must exist and not be closed
327
- if (str(conn_name), enable_iam_auth) in self._cache and not self._cache[
328
- (str(conn_name), enable_iam_auth)
329
- ].closed:
330
- monitored_cache = self._cache[(str(conn_name), enable_iam_auth)]
331
- else:
332
- if self._refresh_strategy == RefreshStrategy.LAZY:
333
- logger.debug(
334
- f"['{conn_name}']: Refresh strategy is set to lazy refresh"
335
- )
336
- cache: LazyRefreshCache | RefreshAheadCache = LazyRefreshCache(
337
- conn_name,
338
- self._client,
339
- self._keys,
340
- enable_iam_auth,
341
- )
342
- else:
343
- logger.debug(
344
- f"['{conn_name}']: Refresh strategy is set to backgound refresh"
345
- )
346
- cache = RefreshAheadCache(
347
- conn_name,
348
- self._client,
349
- self._keys,
350
- enable_iam_auth,
351
- )
352
- # wrap cache as a MonitoredCache
353
- monitored_cache = MonitoredCache(
354
- cache,
355
- self._failover_period,
356
- self._resolver,
357
- )
358
- logger.debug(f"['{conn_name}']: Connection info added to cache")
359
- 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)
360
456
 
361
457
  connect_func = {
362
458
  "pymysql": pymysql.connect,
@@ -371,11 +467,6 @@ class Connector:
371
467
  connector: Callable = connect_func[driver] # type: ignore
372
468
  except KeyError:
373
469
  raise KeyError(f"Driver '{driver}' is not supported.")
374
-
375
- ip_type = kwargs.pop("ip_type", self._ip_type)
376
- # if ip_type is str, convert to IPTypes enum
377
- if isinstance(ip_type, str):
378
- ip_type = IPTypes._from_str(ip_type)
379
470
  kwargs["timeout"] = kwargs.get("timeout", self._timeout)
380
471
 
381
472
  # Host and ssl options come from the certificates and metadata, so we don't
@@ -384,113 +475,210 @@ class Connector:
384
475
  kwargs.pop("ssl", None)
385
476
  kwargs.pop("port", None)
386
477
 
387
- # attempt to get connection info for Cloud SQL instance
478
+ # attempt to establish connection
388
479
  try:
389
- conn_info = await monitored_cache.connect_info()
390
- # validate driver matches intended database engine
391
- DriverMapping.validate_engine(driver, conn_info.database_version)
392
- preferred_ips = conn_info.get_preferred_ips(ip_type)
393
- except Exception:
394
- # with an error from Cloud SQL Admin API call or IP type, invalidate
395
- # the cache and re-raise the error
396
- await self._remove_cached(str(conn_name), enable_iam_auth)
397
- raise
398
-
399
- targets = []
400
- # If the connector is configured with a custom DNS name, attempt to use
401
- # that DNS name to connect to the instance. Fall back to the metadata IP
402
- # address if the DNS name does not resolve to an IP address.
403
- if conn_info.conn_name.domain_name and isinstance(self._resolver, DnsResolver):
404
- try:
405
- ips = await self._resolver.resolve_a_record(conn_info.conn_name.domain_name)
406
- if ips:
407
- 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():
408
489
  logger.debug(
409
- f"['{instance_connection_string}']: Custom DNS name "
410
- f"'{conn_info.conn_name.domain_name}' resolved to '{ips}', "
411
- "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)
412
519
  )
413
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
+ )
414
527
  logger.debug(
415
- f"['{instance_connection_string}']: Custom DNS name "
416
- f"'{conn_info.conn_name.domain_name}' resolved but returned no "
417
- 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})"
418
530
  )
419
- targets.extend(preferred_ips)
420
- except Exception as e: # noqa: BLE001
421
- logger.debug(
422
- f"['{instance_connection_string}']: Custom DNS name "
423
- f"'{conn_info.conn_name.domain_name}' did not resolve to an IP "
424
- f"address: {e}, using '{preferred_ips}' from instance metadata"
425
- )
426
- targets.extend(preferred_ips)
427
- else:
428
- targets.extend(preferred_ips)
429
531
 
430
- # format `user` param for automatic IAM database authn
431
- if enable_iam_auth:
432
- formatted_user = format_database_user(
433
- conn_info.database_version, kwargs["user"]
434
- )
435
- if formatted_user != kwargs["user"]:
436
- logger.debug(
437
- 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),
438
558
  )
439
- kwargs["user"] = formatted_user
440
559
 
441
- try:
442
- last_ex = None
443
- for target_ip in targets:
444
- 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:
445
579
  try:
446
- # async drivers are unblocking and can be awaited directly
447
- if driver in ASYNC_DRIVERS:
448
- 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,
449
661
  target_ip,
450
- await conn_info.create_ssl_context(enable_iam_auth),
662
+ sock,
451
663
  **kwargs,
452
664
  )
665
+ conn = await self._loop.run_in_executor(None, connect_partial)
453
666
  last_ex = None
454
667
  return conn
455
-
456
- # Create socket with SSLContext for sync drivers
457
- ctx = await conn_info.create_ssl_context(enable_iam_auth)
458
- raw_sock = socket.create_connection((target_ip, SERVER_PROXY_PORT))
459
- try:
460
- sock = ctx.wrap_socket(
461
- raw_sock,
462
- 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}"
463
671
  )
464
- except Exception:
465
- raw_sock.close()
466
- raise
467
-
468
- # If this connection was opened using a domain name, then store it
469
- # for later in case we need to forcibly close it on failover.
470
- if conn_info.conn_name.domain_name:
471
- monitored_cache.sockets.append(sock)
472
- # Synchronous drivers are blocking and run using executor
473
- connect_partial = partial(
474
- connector,
475
- target_ip,
476
- sock,
477
- **kwargs,
478
- )
479
- conn = await self._loop.run_in_executor(None, connect_partial)
480
- last_ex = None
481
- return conn
482
- except Exception as e: # noqa: BLE001
483
- logger.debug(
484
- f"['{conn_info.conn_name}']: Connection to {target_ip} failed: {e}"
485
- )
486
- last_ex = e
672
+ last_ex = e
487
673
 
488
- if last_ex:
489
- raise last_ex
674
+ if last_ex:
675
+ raise last_ex
490
676
 
491
677
  except Exception:
492
678
  # with any exception, we attempt a force refresh, then throw the error
493
- 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()
494
682
  raise
495
683
 
496
684
  async def _remove_cached(
@@ -538,8 +726,11 @@ class Connector:
538
726
  close_future = asyncio.run_coroutine_threadsafe(
539
727
  self.close_async(), loop=self._loop
540
728
  )
541
- # Will attempt to safely shut down tasks for 3s
542
- 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}")
543
734
  # if background thread exists for Connector, clean it up
544
735
  if self._thread:
545
736
  if self._loop.is_running():
@@ -552,9 +743,14 @@ class Connector:
552
743
  """Helper function to cancel the cache's tasks
553
744
  and close aiohttp.ClientSession."""
554
745
  self._closed = True
746
+ self._sqldata_client = None
555
747
  if self._client:
556
748
  await self._client.close()
557
- 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
+ )
558
754
 
559
755
 
560
756
  async def create_async_connector(
@@ -570,6 +766,9 @@ async def create_async_connector(
570
766
  refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
571
767
  resolver: type[DefaultResolver | DnsResolver] = DefaultResolver,
572
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,
573
772
  ) -> Connector:
574
773
  """Helper function to create Connector object for asyncio connections.
575
774
 
@@ -579,8 +778,11 @@ async def create_async_connector(
579
778
  Args:
580
779
  ip_type (str | IPTypes): The default IP address type used to connect to
581
780
  Cloud SQL instances. Can be one of the following:
582
- IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), or
583
- 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
584
786
 
585
787
  enable_iam_auth (bool): Enables automatic IAM database authentication
586
788
  (Postgres and MySQL) as the default authentication method for all
@@ -626,6 +828,16 @@ async def create_async_connector(
626
828
  Must be used with `resolver=DnsResolver` to have any effect.
627
829
  Default: 30
628
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
+
629
841
  Returns:
630
842
  A Connector instance configured with running event loop.
631
843
  """
@@ -645,4 +857,7 @@ async def create_async_connector(
645
857
  refresh_strategy=refresh_strategy,
646
858
  resolver=resolver,
647
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,
648
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