diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index feb74526..2cfbd004 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -86,6 +86,14 @@ jobs: POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME POSTGRES_MCP_CONNECTION_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_MCP_CONNECTION_NAME POSTGRES_MCP_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_MCP_PASS + POSTGRES_AIDE_CONNECTION_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_AIDE_CONNECTION_NAME + POSTGRES_AIDE_USER:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_AIDE_USER + POSTGRES_AIDE_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_AIDE_PASS + POSTGRES_AIDE_DB:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_AIDE_DB + POSTGRES_FALLBACK_CONNECTION_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_FALLBACK_CONNECTION_NAME + POSTGRES_FALLBACK_USER:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_FALLBACK_USER + POSTGRES_FALLBACK_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_FALLBACK_PASS + POSTGRES_FALLBACK_DB:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_FALLBACK_DB SQLSERVER_CONNECTION_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/SQLSERVER_CONNECTION_NAME SQLSERVER_USER:${{ vars.GOOGLE_CLOUD_PROJECT }}/SQLSERVER_USER SQLSERVER_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/SQLSERVER_PASS @@ -112,6 +120,14 @@ jobs: POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME: "${{ steps.secrets.outputs.POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME }}" POSTGRES_MCP_CONNECTION_NAME: "${{ steps.secrets.outputs.POSTGRES_MCP_CONNECTION_NAME }}" POSTGRES_MCP_PASS: "${{ steps.secrets.outputs.POSTGRES_MCP_PASS }}" + POSTGRES_AIDE_CONNECTION_NAME: "${{ steps.secrets.outputs.POSTGRES_AIDE_CONNECTION_NAME }}" + POSTGRES_AIDE_USER: "${{ steps.secrets.outputs.POSTGRES_AIDE_USER }}" + POSTGRES_AIDE_PASS: "${{ steps.secrets.outputs.POSTGRES_AIDE_PASS }}" + POSTGRES_AIDE_DB: "${{ steps.secrets.outputs.POSTGRES_AIDE_DB }}" + POSTGRES_FALLBACK_CONNECTION_NAME: "${{ steps.secrets.outputs.POSTGRES_FALLBACK_CONNECTION_NAME }}" + POSTGRES_FALLBACK_USER: "${{ steps.secrets.outputs.POSTGRES_FALLBACK_USER }}" + POSTGRES_FALLBACK_PASS: "${{ steps.secrets.outputs.POSTGRES_FALLBACK_PASS }}" + POSTGRES_FALLBACK_DB: "${{ steps.secrets.outputs.POSTGRES_FALLBACK_DB }}" SQLSERVER_CONNECTION_NAME: "${{ steps.secrets.outputs.SQLSERVER_CONNECTION_NAME }}" SQLSERVER_USER: "${{ steps.secrets.outputs.SQLSERVER_USER }}" SQLSERVER_PASS: "${{ steps.secrets.outputs.SQLSERVER_PASS }}" diff --git a/.gitignore b/.gitignore index 6fed8e21..a8ce0684 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ dist/ sponge_log.xml .envrc *.iml +build/ .mypy_cache/ .nox/ .pytest_cache/ diff --git a/build.sh b/build.sh index 3c76d992..2987465e 100755 --- a/build.sh +++ b/build.sh @@ -120,6 +120,14 @@ function write_e2e_env(){ POSTGRES_CUSTOMER_CAS_INVALID_DOMAIN_NAME=POSTGRES_CUSTOMER_CAS_INVALID_DOMAIN_NAME POSTGRES_MCP_CONNECTION_NAME=POSTGRES_MCP_CONNECTION_NAME POSTGRES_MCP_PASS=POSTGRES_MCP_PASS + POSTGRES_AIDE_CONNECTION_NAME=POSTGRES_AIDE_CONNECTION_NAME + POSTGRES_AIDE_USER=POSTGRES_AIDE_USER + POSTGRES_AIDE_PASS=POSTGRES_AIDE_PASS + POSTGRES_AIDE_DB=POSTGRES_AIDE_DB + POSTGRES_FALLBACK_CONNECTION_NAME=POSTGRES_FALLBACK_CONNECTION_NAME + POSTGRES_FALLBACK_USER=POSTGRES_FALLBACK_USER + POSTGRES_FALLBACK_PASS=POSTGRES_FALLBACK_PASS + POSTGRES_FALLBACK_DB=POSTGRES_FALLBACK_DB SQLSERVER_CONNECTION_NAME=SQLSERVER_CONNECTION_NAME SQLSERVER_USER=SQLSERVER_USER SQLSERVER_PASS=SQLSERVER_PASS diff --git a/google/cloud/sql/connector/asyncpg.py b/google/cloud/sql/connector/asyncpg.py index 2fbc3027..73f99b16 100644 --- a/google/cloud/sql/connector/asyncpg.py +++ b/google/cloud/sql/connector/asyncpg.py @@ -14,6 +14,8 @@ limitations under the License. """ +from __future__ import annotations + import ssl from typing import Any, TYPE_CHECKING @@ -24,16 +26,15 @@ async def connect( - ip_address: str, ctx: ssl.SSLContext, **kwargs: Any -) -> "asyncpg.Connection": + ip_address: str, ctx: ssl.SSLContext | None, **kwargs: Any +) -> asyncpg.Connection: """Helper function to create an asyncpg DB-API connection object. Args: ip_address (str): A string containing an IP address for the Cloud SQL instance. ctx (ssl.SSLContext): An SSLContext object created from the Cloud SQL - server CA cert and ephemeral cert. - server CA cert and ephemeral cert. + server CA cert and ephemeral cert. Pass None to disable SSL. kwargs: Keyword arguments for establishing asyncpg connection object to Cloud SQL instance. @@ -53,14 +54,18 @@ async def connect( user = kwargs.pop("user") db = kwargs.pop("db") passwd = kwargs.pop("password", None) + port = kwargs.pop("port", SERVER_PROXY_PORT) - return await asyncpg.connect( - user=user, - database=db, - password=passwd, - host=ip_address, - port=SERVER_PROXY_PORT, - ssl=ctx, - direct_tls=True, + connect_args = { + "user": user, + "database": db, + "password": passwd, + "host": ip_address, + "port": port, **kwargs, - ) + } + if ctx is not None: + connect_args["ssl"] = ctx + connect_args["direct_tls"] = True + + return await asyncpg.connect(**connect_args) diff --git a/google/cloud/sql/connector/client.py b/google/cloud/sql/connector/client.py index ebf823e2..b0f83b6d 100644 --- a/google/cloud/sql/connector/client.py +++ b/google/cloud/sql/connector/client.py @@ -173,9 +173,13 @@ async def _get_metadata( if psc_dns_names: ip_addresses["PSC"] = psc_dns_names + server_ca_cert = None + if "serverCaCert" in ret_dict and "cert" in ret_dict["serverCaCert"]: + server_ca_cert = ret_dict["serverCaCert"]["cert"] + return { "ip_addresses": ip_addresses, - "server_ca_cert": ret_dict["serverCaCert"]["cert"], + "server_ca_cert": server_ca_cert, "database_version": ret_dict["databaseVersion"], } @@ -271,7 +275,15 @@ async def _get_ephemeral( finally: resp.raise_for_status() - ephemeral_cert: str = ret_dict["ephemeralCert"]["cert"] + try: + ephemeral_cert: str = ret_dict["ephemeralCert"]["cert"] + except KeyError as e: + logger.error( + "KeyError in _get_ephemeral parsing generateEphemeralCert: %s. Response dict: %s", + e, + ret_dict, + ) + raise # decode cert to read expiration x509 = load_pem_x509_certificate( diff --git a/google/cloud/sql/connector/connection_info.py b/google/cloud/sql/connector/connection_info.py index 49404f49..f0b8f1dd 100644 --- a/google/cloud/sql/connector/connection_info.py +++ b/google/cloud/sql/connector/connection_info.py @@ -21,6 +21,7 @@ from typing import Any, TYPE_CHECKING from google.cloud.sql.connector.connection_name import ConnectionName +from google.cloud.sql.connector.exceptions import CloudSQLConnectionError from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError from google.cloud.sql.connector.exceptions import TLSVersionError from google.cloud.sql.connector.utils import AsyncTemporaryDirectory @@ -62,7 +63,7 @@ class ConnectionInfo: conn_name: ConnectionName client_cert: str - server_ca_cert: str + server_ca_cert: str | None private_key: bytes ip_addrs: dict[str, Any] database_version: str @@ -78,6 +79,12 @@ async def create_ssl_context(self, enable_iam_auth: bool = False) -> ssl.SSLCont # if SSL context is cached, use it if self.context is not None: return self.context + + if self.server_ca_cert is None: + raise CloudSQLConnectionError( + "Cannot create SSL context: server CA certificate is missing." + ) + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) # update ssl.PROTOCOL_TLS_CLIENT default diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py index 3a1df0ea..2e06d069 100644 --- a/google/cloud/sql/connector/connector.py +++ b/google/cloud/sql/connector/connector.py @@ -20,8 +20,10 @@ from functools import partial import logging import os +import random import socket from threading import Thread +import time from types import TracebackType from typing import Any, Callable @@ -34,16 +36,20 @@ from google.cloud.sql.connector import pymysql from google.cloud.sql.connector import pytds from google.cloud.sql.connector.client import CloudSQLClient +from google.cloud.sql.connector.connection_name import ConnectionName from google.cloud.sql.connector.enums import DriverMapping from google.cloud.sql.connector.enums import IPTypes from google.cloud.sql.connector.enums import RefreshStrategy from google.cloud.sql.connector.exceptions import ClosedConnectorError from google.cloud.sql.connector.exceptions import ConnectorLoopError +from google.cloud.sql.connector.exceptions import ResourceExhaustedError from google.cloud.sql.connector.instance import RefreshAheadCache from google.cloud.sql.connector.lazy import LazyRefreshCache from google.cloud.sql.connector.monitored_cache import MonitoredCache from google.cloud.sql.connector.resolver import DefaultResolver from google.cloud.sql.connector.resolver import DnsResolver +from google.cloud.sql.connector.sqldata_client import FallbackSocket +from google.cloud.sql.connector.sqldata_client import SqlDataClient from google.cloud.sql.connector.utils import format_database_user from google.cloud.sql.connector.utils import generate_keys @@ -56,6 +62,49 @@ _SQLADMIN_HOST_TEMPLATE = "sqladmin.{universe_domain}" +class SqlDataConnState: + """Tracks connection state, fallback status, and resource exhaustion cooldown for SQL Data Service.""" + + def __init__(self) -> None: + self.allowed: bool = True + self.cooldown_until: float | None = None + self.backoff_counter: int = 0 + self.last_err: Exception | None = None + + def is_cooldown_active(self) -> bool: + """Returns True if the instance connection is in cooldown due to resource exhaustion.""" + return bool( + self.allowed + and self.cooldown_until + and time.time() < self.cooldown_until + ) + + def record_exhausted(self, err: Exception, base_cooldown: float) -> float: + """Records a resource exhaustion error, increments backoff counter, and returns the cooldown duration.""" + if self.backoff_counter < 5: + self.backoff_counter += 1 + backoff = _cooldown_backoff(base_cooldown, self.backoff_counter) + self.cooldown_until = time.time() + backoff + self.last_err = err + return backoff + + def record_success(self) -> None: + """Resets cooldown and backoff state on successful communication.""" + self.backoff_counter = 0 + self.cooldown_until = None + self.last_err = None + + def record_fallback(self) -> None: + """Marks SQL Data Service as not allowed for this instance.""" + self.allowed = False + + +def _cooldown_backoff(base_cooldown: float, attempt: int) -> float: + multi = 1.618 + exp = float(attempt - 1) + random.random() + return base_cooldown * (multi**exp) + + class Connector: """Configure and create secure connections to Cloud SQL.""" @@ -73,14 +122,18 @@ def __init__( refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, failover_period: int = 30, + sql_data_endpoint: str = "sqladmin.googleapis.com", + sql_data_stream_timeout: int = 7200, + resource_exhausted_cooldown_period: float = 5.0, ) -> None: """Initializes a Connector instance. Args: ip_type (str | IPTypes): The default IP address type used to connect to Cloud SQL instances. Can be one of the following: - IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), or - IPTypes.PSC ("PSC"). Default: IPTypes.PUBLIC + IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), + IPTypes.PSC ("PSC"), or IPTypes.SQL_DATA ("SQL_DATA"). + Default: IPTypes.PUBLIC enable_iam_auth (bool): Enables automatic IAM database authentication (Postgres and MySQL) as the default authentication method for all @@ -125,6 +178,15 @@ def __init__( attempt to check if a failover has occured for a given instance. Must be used with `resolver=DnsResolver` to have any effect. Default: 30 + + sql_data_endpoint (str): Endpoint host for SQL Data Service calls. + Default: "sqladmin.googleapis.com". + + sql_data_stream_timeout (int): Timeout in seconds for the SQL Data + Service gRPC stream. Default: 7200. + + resource_exhausted_cooldown_period (float): Cooldown period in seconds + after a ResourceExhausted error. Default: 5.0. """ # if refresh_strategy is str, convert to RefreshStrategy enum if isinstance(refresh_strategy, str): @@ -213,6 +275,16 @@ def __init__( "configured the universe domain explicitly, `googleapis.com` " "is the default." ) + self._sql_data_endpoint = sql_data_endpoint + self._sql_data_stream_timeout = sql_data_stream_timeout + self._resource_exhausted_cooldown_period = ( + resource_exhausted_cooldown_period + ) + self._sql_data_fallback_cache: set[str] = set() + self._sql_data_conn_state: dict[str, SqlDataConnState] = {} + self._sqldata_clients: set[SqlDataClient] = set() + + @property def universe_domain(self) -> str: @@ -259,6 +331,49 @@ def connect( ) return connect_future.result() + def _get_or_create_cache( + self, + conn_name: ConnectionName, + enable_iam_auth: bool, + ) -> MonitoredCache: + assert self._client is not None, "client must be initialized before creating cache" + assert self._keys is not None, "keys must be initialized before creating cache" + assert self._resolver is not None, "resolver must be initialized before creating cache" + if (str(conn_name), enable_iam_auth) in self._cache and not self._cache[ + (str(conn_name), enable_iam_auth) + ].closed: + return self._cache[(str(conn_name), enable_iam_auth)] + + if self._refresh_strategy == RefreshStrategy.LAZY: + logger.debug( + f"['{conn_name}']: Refresh strategy is set to lazy refresh" + ) + cache: LazyRefreshCache | RefreshAheadCache = LazyRefreshCache( + conn_name, + self._client, + self._keys, + enable_iam_auth, + ) + else: + logger.debug( + f"['{conn_name}']: Refresh strategy is set to backgound refresh" + ) + cache = RefreshAheadCache( + conn_name, + self._client, + self._keys, + enable_iam_auth, + ) + # wrap cache as a MonitoredCache + monitored_cache = MonitoredCache( + cache, + self._failover_period, + self._resolver, + ) + logger.debug(f"['{conn_name}']: Connection info added to cache") + self._cache[(str(conn_name), enable_iam_auth)] = monitored_cache + return monitored_cache + async def connect_async( self, instance_connection_string: str, driver: str, **kwargs: Any ) -> Any: @@ -320,42 +435,13 @@ async def connect_async( if self._resolver is None: self._resolver = self._resolver_cls(client=self._client) enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth) + ip_type = kwargs.pop("ip_type", self._ip_type) + if isinstance(ip_type, str): + ip_type = IPTypes._from_str(ip_type) conn_name = await self._resolver.resolve(instance_connection_string) - # Cache entry must exist and not be closed - if (str(conn_name), enable_iam_auth) in self._cache and not self._cache[ - (str(conn_name), enable_iam_auth) - ].closed: - monitored_cache = self._cache[(str(conn_name), enable_iam_auth)] - else: - if self._refresh_strategy == RefreshStrategy.LAZY: - logger.debug( - f"['{conn_name}']: Refresh strategy is set to lazy refresh" - ) - cache: LazyRefreshCache | RefreshAheadCache = LazyRefreshCache( - conn_name, - self._client, - self._keys, - enable_iam_auth, - ) - else: - logger.debug( - f"['{conn_name}']: Refresh strategy is set to backgound refresh" - ) - cache = RefreshAheadCache( - conn_name, - self._client, - self._keys, - enable_iam_auth, - ) - # wrap cache as a MonitoredCache - monitored_cache = MonitoredCache( - cache, - self._failover_period, - self._resolver, - ) - logger.debug(f"['{conn_name}']: Connection info added to cache") - self._cache[(str(conn_name), enable_iam_auth)] = monitored_cache + if ip_type != IPTypes.SQL_DATA: + monitored_cache = self._get_or_create_cache(conn_name, enable_iam_auth) connect_func = { "pymysql": pymysql.connect, @@ -369,11 +455,6 @@ async def connect_async( connector: Callable = connect_func[driver] # type: ignore except KeyError: raise KeyError(f"Driver '{driver}' is not supported.") - - ip_type = kwargs.pop("ip_type", self._ip_type) - # if ip_type is str, convert to IPTypes enum - if isinstance(ip_type, str): - ip_type = IPTypes._from_str(ip_type) kwargs["timeout"] = kwargs.get("timeout", self._timeout) # Host and ssl options come from the certificates and metadata, so we don't @@ -382,113 +463,215 @@ async def connect_async( kwargs.pop("ssl", None) kwargs.pop("port", None) - # attempt to get connection info for Cloud SQL instance + # attempt to establish connection try: - conn_info = await monitored_cache.connect_info() - # validate driver matches intended database engine - DriverMapping.validate_engine(driver, conn_info.database_version) - preferred_ips = conn_info.get_preferred_ips(ip_type) - except Exception: - # with an error from Cloud SQL Admin API call or IP type, invalidate - # the cache and re-raise the error - await self._remove_cached(str(conn_name), enable_iam_auth) - raise - - targets = [] - # If the connector is configured with a custom DNS name, attempt to use - # that DNS name to connect to the instance. Fall back to the metadata IP - # address if the DNS name does not resolve to an IP address. - if conn_info.conn_name.domain_name and isinstance(self._resolver, DnsResolver): - try: - ips = await self._resolver.resolve_a_record(conn_info.conn_name.domain_name) - if ips: - targets.extend(ips) + if ip_type == IPTypes.SQL_DATA: + state = self._sql_data_conn_state.setdefault( + str(conn_name), SqlDataConnState() + ) + if state.is_cooldown_active(): logger.debug( - f"['{instance_connection_string}']: Custom DNS name " - f"'{conn_info.conn_name.domain_name}' resolved to '{ips}', " - "using it to connect" + f"['{conn_name}']: SQL Data Service in cooldown until {state.cooldown_until}" ) - else: - logger.debug( - f"['{instance_connection_string}']: Custom DNS name " - f"'{conn_info.conn_name.domain_name}' resolved but returned no " - f"entries, using '{preferred_ips}' from instance metadata" + raise ResourceExhaustedError( + "cooldown active", str(conn_name), state.last_err ) - targets.extend(preferred_ips) - except Exception as e: # noqa: BLE001 - logger.debug( - f"['{instance_connection_string}']: Custom DNS name " - f"'{conn_info.conn_name.domain_name}' did not resolve to an IP " - f"address: {e}, using '{preferred_ips}' from instance metadata" + + logger.debug(f"['{conn_name}']: Connecting via SQL Data Service tunnel") + if enable_iam_auth: + engine = DriverMapping[driver.upper()].value + formatted_user = format_database_user( + engine, kwargs["user"] + ) + if formatted_user != kwargs["user"]: + logger.debug( + f"['{instance_connection_string}']: Truncated IAM database username from {kwargs['user']} to {formatted_user}" + ) + kwargs["user"] = formatted_user + + sqldata_client = SqlDataClient( + endpoint=self._sql_data_endpoint, + credentials=self._credentials, + quota_project=self._quota_project, + timeout=self._sql_data_stream_timeout, + ) + self._sqldata_clients.add(sqldata_client) + sqldata_client._on_close_callbacks.append( + lambda: self._sqldata_clients.discard(sqldata_client) ) - targets.extend(preferred_ips) - else: - targets.extend(preferred_ips) - # format `user` param for automatic IAM database authn - if enable_iam_auth: - formatted_user = format_database_user( - conn_info.database_version, kwargs["user"] - ) - if formatted_user != kwargs["user"]: - logger.debug( - f"['{instance_connection_string}']: Truncated IAM database username from {kwargs['user']} to {formatted_user}" + def on_resource_exhausted(err: Exception) -> None: + backoff = state.record_exhausted( + err, self._resource_exhausted_cooldown_period + ) + logger.debug( + f"['{conn_name}']: ResourceExhausted occurred, backing off for {backoff:.2f}s " + f"(attempt {state.backoff_counter})" + ) + + def on_success() -> None: + state.record_success() + + def on_fallback(name: str) -> None: + state.record_fallback() + self._sql_data_fallback_cache.add(name) + + def is_fallback_cached(name: str) -> bool: + return not state.allowed or name in self._sql_data_fallback_cache + + # Defer cache creation and connect_info call + async def get_conn_info(): + cache = self._get_or_create_cache(conn_name, enable_iam_auth) + return await cache.connect_info() + + tunnel_port = await sqldata_client.connect_tunnel( + instance_connection_name=str(conn_name), + region=conn_name.region, + project=conn_name.project, + get_conn_info=get_conn_info, + enable_iam_auth=enable_iam_auth, + on_fallback=on_fallback, + is_fallback_cached=is_fallback_cached, + on_resource_exhausted=on_resource_exhausted, + on_success=on_success, + connect_timeout=kwargs.get("timeout", self._timeout), ) - kwargs["user"] = formatted_user - try: - last_ex = None - for target_ip in targets: - logger.debug(f"['{conn_info.conn_name}']: Connecting to {target_ip}:3307") - try: - # async drivers are unblocking and can be awaited directly - if driver in ASYNC_DRIVERS: - conn = await connector( - target_ip, - await conn_info.create_ssl_context(enable_iam_auth), - **kwargs, - ) - last_ex = None - return conn + if driver in ASYNC_DRIVERS: + return await connector( + "127.0.0.1", + None, + port=tunnel_port, + **kwargs, + ) + else: + raw_sock = socket.create_connection(("127.0.0.1", tunnel_port)) + fd = raw_sock.detach() + fallback_sock = FallbackSocket(fileno=fd) - # Create socket with SSLContext for sync drivers - ctx = await conn_info.create_ssl_context(enable_iam_auth) - raw_sock = socket.create_connection((target_ip, SERVER_PROXY_PORT)) - try: - sock = ctx.wrap_socket( - raw_sock, - server_hostname=target_ip, - ) - except Exception: - raw_sock.close() - raise + cache = None + if conn_name.domain_name: + cache = self._get_or_create_cache(conn_name, enable_iam_auth) + cache.sockets.append(fallback_sock) - # If this connection was opened using a domain name, then store it - # for later in case we need to forcibly close it on failover. - if conn_info.conn_name.domain_name: - monitored_cache.sockets.append(sock) - # Synchronous drivers are blocking and run using executor connect_partial = partial( connector, - target_ip, - sock, + "127.0.0.1", + fallback_sock, **kwargs, ) - conn = await self._loop.run_in_executor(None, connect_partial) - last_ex = None - return conn - except Exception as e: # noqa: BLE001 - logger.debug( - f"['{conn_info.conn_name}']: Connection to {target_ip} failed: {e}" + try: + return await self._loop.run_in_executor(None, connect_partial) + except Exception: + fallback_sock.close() + if conn_name.domain_name and cache: + cache._purge_closed_sockets() + raise + else: + try: + conn_info = await monitored_cache.connect_info() + # validate driver matches intended database engine + DriverMapping.validate_engine(driver, conn_info.database_version) + preferred_ips = conn_info.get_preferred_ips(ip_type) + except Exception: + # with an error from Cloud SQL Admin API call or IP type, invalidate + # the cache and re-raise the error + await self._remove_cached(str(conn_name), enable_iam_auth) + raise + + targets = [] + if conn_info.conn_name.domain_name and isinstance(self._resolver, DnsResolver): + try: + ips = await self._resolver.resolve_a_record(conn_info.conn_name.domain_name) + if ips: + targets.extend(ips) + logger.debug( + f"['{instance_connection_string}']: Custom DNS name " + f"'{conn_info.conn_name.domain_name}' resolved to '{ips}', " + "using it to connect" + ) + else: + logger.debug( + f"['{instance_connection_string}']: Custom DNS name " + f"'{conn_info.conn_name.domain_name}' resolved but returned no " + f"entries, using '{preferred_ips}' from instance metadata" + ) + targets.extend(preferred_ips) + except Exception as e: # noqa: BLE001 + logger.debug( + f"['{instance_connection_string}']: Custom DNS name " + f"'{conn_info.conn_name.domain_name}' did not resolve to an IP " + f"address: {e}, using '{preferred_ips}' from instance metadata" + ) + targets.extend(preferred_ips) + else: + targets.extend(preferred_ips) + + # format `user` param for automatic IAM database authn + if enable_iam_auth: + formatted_user = format_database_user( + conn_info.database_version, kwargs["user"] ) - last_ex = e + if formatted_user != kwargs["user"]: + logger.debug( + f"['{instance_connection_string}']: Truncated IAM database username from {kwargs['user']} to {formatted_user}" + ) + kwargs["user"] = formatted_user + + last_ex = None + for target_ip in targets: + logger.debug(f"['{conn_info.conn_name}']: Connecting to {target_ip}:3307") + try: + # async drivers are unblocking and can be awaited directly + if driver in ASYNC_DRIVERS: + conn = await connector( + target_ip, + await conn_info.create_ssl_context(enable_iam_auth), + **kwargs, + ) + last_ex = None + return conn + + # Create socket with SSLContext for sync drivers + ctx = await conn_info.create_ssl_context(enable_iam_auth) + raw_sock = socket.create_connection((target_ip, SERVER_PROXY_PORT)) + try: + sock = ctx.wrap_socket( + raw_sock, + server_hostname=target_ip, + ) + except Exception: + raw_sock.close() + raise + + # If this connection was opened using a domain name, then store it + # for later in case we need to forcibly close it on failover. + if conn_info.conn_name.domain_name: + monitored_cache.sockets.append(sock) + # Synchronous drivers are blocking and run using executor + connect_partial = partial( + connector, + target_ip, + sock, + **kwargs, + ) + conn = await self._loop.run_in_executor(None, connect_partial) + last_ex = None + return conn + except Exception as e: # noqa: BLE001 + logger.debug( + f"['{conn_info.conn_name}']: Connection to {target_ip} failed: {e}" + ) + last_ex = e - if last_ex: - raise last_ex + if last_ex: + raise last_ex except Exception: # with any exception, we attempt a force refresh, then throw the error - await monitored_cache.force_refresh() + cached_entry = self._cache.get((str(conn_name), enable_iam_auth)) + if cached_entry: + await cached_entry.force_refresh() raise async def _remove_cached( @@ -536,8 +719,11 @@ def close(self) -> None: close_future = asyncio.run_coroutine_threadsafe( self.close_async(), loop=self._loop ) - # Will attempt to safely shut down tasks for 3s - close_future.result(timeout=3) + try: + # Will attempt to safely shut down tasks for 3s + close_future.result(timeout=3) + except Exception as e: # noqa: BLE001 + logger.error(f"Error during close_async: {e}") # if background thread exists for Connector, clean it up if self._thread: if self._loop.is_running(): @@ -552,7 +738,11 @@ async def close_async(self) -> None: self._closed = True if self._client: await self._client.close() - await asyncio.gather(*[cache.close() for cache in self._cache.values()]) + await asyncio.gather( + *[cache.close() for cache in self._cache.values()], + *[client.close() for client in list(self._sqldata_clients)], + return_exceptions=True, + ) async def create_async_connector( @@ -568,6 +758,9 @@ async def create_async_connector( refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, resolver: type[DefaultResolver | DnsResolver] = DefaultResolver, failover_period: int = 30, + sql_data_endpoint: str = "sqladmin.googleapis.com", + sql_data_stream_timeout: int = 7200, + resource_exhausted_cooldown_period: float = 5.0, ) -> Connector: """Helper function to create Connector object for asyncio connections. @@ -577,8 +770,9 @@ async def create_async_connector( Args: ip_type (str | IPTypes): The default IP address type used to connect to Cloud SQL instances. Can be one of the following: - IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), or - IPTypes.PSC ("PSC"). Default: IPTypes.PUBLIC + IPTypes.PUBLIC ("PUBLIC"), IPTypes.PRIVATE ("PRIVATE"), + IPTypes.PSC ("PSC"), or IPTypes.SQL_DATA ("SQL_DATA"). + Default: IPTypes.PUBLIC enable_iam_auth (bool): Enables automatic IAM database authentication (Postgres and MySQL) as the default authentication method for all @@ -624,6 +818,15 @@ async def create_async_connector( Must be used with `resolver=DnsResolver` to have any effect. Default: 30 + sql_data_endpoint (str): Endpoint host for SQL Data Service calls. + Default: "sqladmin.googleapis.com". + + sql_data_stream_timeout (int): Timeout in seconds for the SQL Data + Service gRPC stream. Default: 7200. + + resource_exhausted_cooldown_period (float): Cooldown period in seconds + after a ResourceExhausted error. Default: 5.0. + Returns: A Connector instance configured with running event loop. """ @@ -643,4 +846,7 @@ async def create_async_connector( refresh_strategy=refresh_strategy, resolver=resolver, failover_period=failover_period, + sql_data_endpoint=sql_data_endpoint, + sql_data_stream_timeout=sql_data_stream_timeout, + resource_exhausted_cooldown_period=resource_exhausted_cooldown_period, ) diff --git a/google/cloud/sql/connector/enums.py b/google/cloud/sql/connector/enums.py index 88b5bf47..3a8af8a4 100644 --- a/google/cloud/sql/connector/enums.py +++ b/google/cloud/sql/connector/enums.py @@ -41,6 +41,7 @@ class IPTypes(Enum): PUBLIC = "PRIMARY" PRIVATE = "PRIVATE" PSC = "PSC" + SQL_DATA = "SQL_DATA" @classmethod def _missing_(cls, value: object) -> None: @@ -54,6 +55,8 @@ def _from_str(cls, ip_type_str: str) -> IPTypes: """Convert IP type from a str into IPTypes.""" if ip_type_str.upper() == "PUBLIC": ip_type_str = "PRIMARY" + elif ip_type_str.upper() in ("SQLDATA", "SQL_DATA"): + ip_type_str = "SQL_DATA" return cls(ip_type_str.upper()) diff --git a/google/cloud/sql/connector/exceptions.py b/google/cloud/sql/connector/exceptions.py index 4c3d1acb..49d710d7 100644 --- a/google/cloud/sql/connector/exceptions.py +++ b/google/cloud/sql/connector/exceptions.py @@ -14,6 +14,8 @@ limitations under the License. """ +from __future__ import annotations + class ConnectorLoopError(Exception): """ @@ -91,3 +93,34 @@ class ClosedConnectorError(Exception): Exception to be raised when a Connector is closed and connect method is called on it. """ + + +class CloudSQLConnectionError(Exception): + """ + Exception to be raised when a connection cannot be established to a Cloud SQL instance. + """ + + +class ResourceExhaustedError(CloudSQLConnectionError): + """ + Exception to be raised when a connection cannot be established because + the SQL Data Service is busy / in cooldown due to resource exhaustion. + """ + + def __init__( + self, + message: str, + connection_name: str | None = None, + raw_error: Exception | None = None, + ) -> None: + self.message = message + self.connection_name = connection_name + self.raw_error = raw_error + super().__init__( + f"[{connection_name}] {message}: {raw_error}" + if connection_name and raw_error + else f"[{connection_name}] {message}" + if connection_name + else message + ) + diff --git a/google/cloud/sql/connector/monitored_cache.py b/google/cloud/sql/connector/monitored_cache.py index d2ba2290..05d04b0f 100644 --- a/google/cloud/sql/connector/monitored_cache.py +++ b/google/cloud/sql/connector/monitored_cache.py @@ -16,7 +16,7 @@ import asyncio import logging -import ssl +import socket from typing import Any, Callable import aiohttp @@ -44,7 +44,7 @@ def __init__( self.resolver = resolver self.cache = cache self.domain_name_ticker: asyncio.Task | None = None - self.sockets: list[ssl.SSLSocket] = [] + self.sockets: list[socket.socket] = [] # If domain name is configured for instance and failover period is set, # poll for DNS record changes. @@ -77,11 +77,11 @@ def _purge_closed_sockets(self) -> None: list of sockets. """ open_sockets = [] - for socket in self.sockets: + for sock in self.sockets: # Check fileno for if socket is closed. Will return # -1 on failure, which will be used to signal socket closed. - if socket.fileno() != -1: - open_sockets.append(socket) + if sock.fileno() != -1: + open_sockets.append(sock) self.sockets = open_sockets async def _check_domain_name(self) -> None: @@ -159,11 +159,11 @@ async def close(self) -> None: await self.cache.close() # Close any still open sockets - for socket in self.sockets: + for sock in self.sockets: # Check fileno for if socket is closed. Will return # -1 on failure, which will be used to signal socket closed. - if socket.fileno() != -1: - socket.close() + if sock.fileno() != -1: + sock.close() async def ticker(interval: int, function: Callable, *args: Any, **kwargs: Any) -> None: diff --git a/google/cloud/sql/connector/proto/google/rpc/code.proto b/google/cloud/sql/connector/proto/google/rpc/code.proto new file mode 100644 index 00000000..8fef4117 --- /dev/null +++ b/google/cloud/sql/connector/proto/google/rpc/code.proto @@ -0,0 +1,186 @@ +// Copyright 2017 Google Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto3"; + +package google.rpc; + +option go_package = "google.golang.org/genproto/googleapis/rpc/code;code"; +option java_multiple_files = true; +option java_outer_classname = "CodeProto"; +option java_package = "com.google.rpc"; +option objc_class_prefix = "RPC"; + + +// The canonical error codes for Google APIs. +// +// +// Sometimes multiple error codes may apply. Services should return +// the most specific error code that applies. For example, prefer +// `OUT_OF_RANGE` over `FAILED_PRECONDITION` if both codes apply. +// Similarly prefer `NOT_FOUND` or `ALREADY_EXISTS` over `FAILED_PRECONDITION`. +enum Code { + // Not an error; returned on success + // + // HTTP Mapping: 200 OK + OK = 0; + + // The operation was cancelled, typically by the caller. + // + // HTTP Mapping: 499 Client Closed Request + CANCELLED = 1; + + // Unknown error. For example, this error may be returned when + // a `Status` value received from another address space belongs to + // an error space that is not known in this address space. Also + // errors raised by APIs that do not return enough error information + // may be converted to this error. + // + // HTTP Mapping: 500 Internal Server Error + UNKNOWN = 2; + + // The client specified an invalid argument. Note that this differs + // from `FAILED_PRECONDITION`. `INVALID_ARGUMENT` indicates arguments + // that are problematic regardless of the state of the system + // (e.g., a malformed file name). + // + // HTTP Mapping: 400 Bad Request + INVALID_ARGUMENT = 3; + + // The deadline expired before the operation could complete. For operations + // that change the state of the system, this error may be returned + // even if the operation has completed successfully. For example, a + // successful response from a server could have been delayed long + // enough for the deadline to expire. + // + // HTTP Mapping: 504 Gateway Timeout + DEADLINE_EXCEEDED = 4; + + // Some requested entity (e.g., file or directory) was not found. + // + // Note to server developers: if a request is denied for an entire class + // of users, such as gradual feature rollout or undocumented whitelist, + // `NOT_FOUND` may be used. If a request is denied for some users within + // a class of users, such as user-based access control, `PERMISSION_DENIED` + // must be used. + // + // HTTP Mapping: 404 Not Found + NOT_FOUND = 5; + + // The entity that a client attempted to create (e.g., file or directory) + // already exists. + // + // HTTP Mapping: 409 Conflict + ALREADY_EXISTS = 6; + + // The caller does not have permission to execute the specified + // operation. `PERMISSION_DENIED` must not be used for rejections + // caused by exhausting some resource (use `RESOURCE_EXHAUSTED` + // instead for those errors). `PERMISSION_DENIED` must not be + // used if the caller can not be identified (use `UNAUTHENTICATED` + // instead for those errors). This error code does not imply the + // request is valid or the requested entity exists or satisfies + // other pre-conditions. + // + // HTTP Mapping: 403 Forbidden + PERMISSION_DENIED = 7; + + // The request does not have valid authentication credentials for the + // operation. + // + // HTTP Mapping: 401 Unauthorized + UNAUTHENTICATED = 16; + + // Some resource has been exhausted, perhaps a per-user quota, or + // perhaps the entire file system is out of space. + // + // HTTP Mapping: 429 Too Many Requests + RESOURCE_EXHAUSTED = 8; + + // The operation was rejected because the system is not in a state + // required for the operation's execution. For example, the directory + // to be deleted is non-empty, an rmdir operation is applied to + // a non-directory, etc. + // + // Service implementors can use the following guidelines to decide + // between `FAILED_PRECONDITION`, `ABORTED`, and `UNAVAILABLE`: + // (a) Use `UNAVAILABLE` if the client can retry just the failing call. + // (b) Use `ABORTED` if the client should retry at a higher level + // (e.g., when a client-specified test-and-set fails, indicating the + // client should restart a read-modify-write sequence). + // (c) Use `FAILED_PRECONDITION` if the client should not retry until + // the system state has been explicitly fixed. E.g., if an "rmdir" + // fails because the directory is non-empty, `FAILED_PRECONDITION` + // should be returned since the client should not retry unless + // the files are deleted from the directory. + // + // HTTP Mapping: 400 Bad Request + FAILED_PRECONDITION = 9; + + // The operation was aborted, typically due to a concurrency issue such as + // a sequencer check failure or transaction abort. + // + // See the guidelines above for deciding between `FAILED_PRECONDITION`, + // `ABORTED`, and `UNAVAILABLE`. + // + // HTTP Mapping: 409 Conflict + ABORTED = 10; + + // The operation was attempted past the valid range. E.g., seeking or + // reading past end-of-file. + // + // Unlike `INVALID_ARGUMENT`, this error indicates a problem that may + // be fixed if the system state changes. For example, a 32-bit file + // system will generate `INVALID_ARGUMENT` if asked to read at an + // offset that is not in the range [0,2^32-1], but it will generate + // `OUT_OF_RANGE` if asked to read from an offset past the current + // file size. + // + // There is a fair bit of overlap between `FAILED_PRECONDITION` and + // `OUT_OF_RANGE`. We recommend using `OUT_OF_RANGE` (the more specific + // error) when it applies so that callers who are iterating through + // a space can easily look for an `OUT_OF_RANGE` error to detect when + // they are done. + // + // HTTP Mapping: 400 Bad Request + OUT_OF_RANGE = 11; + + // The operation is not implemented or is not supported/enabled in this + // service. + // + // HTTP Mapping: 501 Not Implemented + UNIMPLEMENTED = 12; + + // Internal errors. This means that some invariants expected by the + // underlying system have been broken. This error code is reserved + // for serious errors. + // + // HTTP Mapping: 500 Internal Server Error + INTERNAL = 13; + + // The service is currently unavailable. This is most likely a + // transient condition, which can be corrected by retrying with + // a backoff. + // + // See the guidelines above for deciding between `FAILED_PRECONDITION`, + // `ABORTED`, and `UNAVAILABLE`. + // + // HTTP Mapping: 503 Service Unavailable + UNAVAILABLE = 14; + + // Unrecoverable data loss or corruption. + // + // HTTP Mapping: 500 Internal Server Error + DATA_LOSS = 15; +} diff --git a/google/cloud/sql/connector/proto/google/rpc/error_details.proto b/google/cloud/sql/connector/proto/google/rpc/error_details.proto new file mode 100644 index 00000000..f24ae009 --- /dev/null +++ b/google/cloud/sql/connector/proto/google/rpc/error_details.proto @@ -0,0 +1,200 @@ +// Copyright 2017 Google Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto3"; + +package google.rpc; + +import "google/protobuf/duration.proto"; + +option go_package = "google.golang.org/genproto/googleapis/rpc/errdetails;errdetails"; +option java_multiple_files = true; +option java_outer_classname = "ErrorDetailsProto"; +option java_package = "com.google.rpc"; +option objc_class_prefix = "RPC"; + + +// Describes when the clients can retry a failed request. Clients could ignore +// the recommendation here or retry when this information is missing from error +// responses. +// +// It's always recommended that clients should use exponential backoff when +// retrying. +// +// Clients should wait until `retry_delay` amount of time has passed since +// receiving the error response before retrying. If retrying requests also +// fail, clients should use an exponential backoff scheme to gradually increase +// the delay between retries based on `retry_delay`, until either a maximum +// number of retires have been reached or a maximum retry delay cap has been +// reached. +message RetryInfo { + // Clients should wait at least this long between retrying the same request. + google.protobuf.Duration retry_delay = 1; +} + +// Describes additional debugging info. +message DebugInfo { + // The stack trace entries indicating where the error occurred. + repeated string stack_entries = 1; + + // Additional debugging information provided by the server. + string detail = 2; +} + +// Describes how a quota check failed. +// +// For example if a daily limit was exceeded for the calling project, +// a service could respond with a QuotaFailure detail containing the project +// id and the description of the quota limit that was exceeded. If the +// calling project hasn't enabled the service in the developer console, then +// a service could respond with the project id and set `service_disabled` +// to true. +// +// Also see RetryDetail and Help types for other details about handling a +// quota failure. +message QuotaFailure { + // A message type used to describe a single quota violation. For example, a + // daily quota or a custom quota that was exceeded. + message Violation { + // The subject on which the quota check failed. + // For example, "clientip:" or "project:". + string subject = 1; + + // A description of how the quota check failed. Clients can use this + // description to find more about the quota configuration in the service's + // public documentation, or find the relevant quota limit to adjust through + // developer console. + // + // For example: "Service disabled" or "Daily Limit for read operations + // exceeded". + string description = 2; + } + + // Describes all quota violations. + repeated Violation violations = 1; +} + +// Describes what preconditions have failed. +// +// For example, if an RPC failed because it required the Terms of Service to be +// acknowledged, it could list the terms of service violation in the +// PreconditionFailure message. +message PreconditionFailure { + // A message type used to describe a single precondition failure. + message Violation { + // The type of PreconditionFailure. We recommend using a service-specific + // enum type to define the supported precondition violation types. For + // example, "TOS" for "Terms of Service violation". + string type = 1; + + // The subject, relative to the type, that failed. + // For example, "google.com/cloud" relative to the "TOS" type would + // indicate which terms of service is being referenced. + string subject = 2; + + // A description of how the precondition failed. Developers can use this + // description to understand how to fix the failure. + // + // For example: "Terms of service not accepted". + string description = 3; + } + + // Describes all precondition violations. + repeated Violation violations = 1; +} + +// Describes violations in a client request. This error type focuses on the +// syntactic aspects of the request. +message BadRequest { + // A message type used to describe a single bad request field. + message FieldViolation { + // A path leading to a field in the request body. The value will be a + // sequence of dot-separated identifiers that identify a protocol buffer + // field. E.g., "field_violations.field" would identify this field. + string field = 1; + + // A description of why the request element is bad. + string description = 2; + } + + // Describes all violations in a client request. + repeated FieldViolation field_violations = 1; +} + +// Contains metadata about the request that clients can attach when filing a bug +// or providing other forms of feedback. +message RequestInfo { + // An opaque string that should only be interpreted by the service generating + // it. For example, it can be used to identify requests in the service's logs. + string request_id = 1; + + // Any data that was used to serve this request. For example, an encrypted + // stack trace that can be sent back to the service provider for debugging. + string serving_data = 2; +} + +// Describes the resource that is being accessed. +message ResourceInfo { + // A name for the type of resource being accessed, e.g. "sql table", + // "cloud storage bucket", "file", "Google calendar"; or the type URL + // of the resource: e.g. "type.googleapis.com/google.pubsub.v1.Topic". + string resource_type = 1; + + // The name of the resource being accessed. For example, a shared calendar + // name: "example.com_4fghdhgsrgh@group.calendar.google.com", if the current + // error is [google.rpc.Code.PERMISSION_DENIED][google.rpc.Code.PERMISSION_DENIED]. + string resource_name = 2; + + // The owner of the resource (optional). + // For example, "user:" or "project:". + string owner = 3; + + // Describes what error is encountered when accessing this resource. + // For example, updating a cloud project may require the `writer` permission + // on the developer console project. + string description = 4; +} + +// Provides links to documentation or for performing an out of band action. +// +// For example, if a quota check failed with an error indicating the calling +// project hasn't enabled the accessed service, this can contain a URL pointing +// directly to the right place in the developer console to flip the bit. +message Help { + // Describes a URL link. + message Link { + // Describes what the link offers. + string description = 1; + + // The URL of the link. + string url = 2; + } + + // URL(s) pointing to additional information on handling the current error. + repeated Link links = 1; +} + +// Provides a localized error message that is safe to return to the user +// which can be attached to an RPC error. +message LocalizedMessage { + // The locale used following the specification defined at + // http://www.rfc-editor.org/rfc/bcp/bcp47.txt. + // Examples are: "en-US", "fr-CH", "es-MX" + string locale = 1; + + // The localized error message in the above locale. + string message = 2; +} diff --git a/google/cloud/sql/connector/proto/google/rpc/status.proto b/google/cloud/sql/connector/proto/google/rpc/status.proto new file mode 100644 index 00000000..0839ee96 --- /dev/null +++ b/google/cloud/sql/connector/proto/google/rpc/status.proto @@ -0,0 +1,92 @@ +// Copyright 2017 Google Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto3"; + +package google.rpc; + +import "google/protobuf/any.proto"; + +option go_package = "google.golang.org/genproto/googleapis/rpc/status;status"; +option java_multiple_files = true; +option java_outer_classname = "StatusProto"; +option java_package = "com.google.rpc"; +option objc_class_prefix = "RPC"; + + +// The `Status` type defines a logical error model that is suitable for different +// programming environments, including REST APIs and RPC APIs. It is used by +// [gRPC](https://github.com/grpc). The error model is designed to be: +// +// - Simple to use and understand for most users +// - Flexible enough to meet unexpected needs +// +// # Overview +// +// The `Status` message contains three pieces of data: error code, error message, +// and error details. The error code should be an enum value of +// [google.rpc.Code][google.rpc.Code], but it may accept additional error codes if needed. The +// error message should be a developer-facing English message that helps +// developers *understand* and *resolve* the error. If a localized user-facing +// error message is needed, put the localized message in the error details or +// localize it in the client. The optional error details may contain arbitrary +// information about the error. There is a predefined set of error detail types +// in the package `google.rpc` that can be used for common error conditions. +// +// # Language mapping +// +// The `Status` message is the logical representation of the error model, but it +// is not necessarily the actual wire format. When the `Status` message is +// exposed in different client libraries and different wire protocols, it can be +// mapped differently. For example, it will likely be mapped to some exceptions +// in Java, but more likely mapped to some error codes in C. +// +// # Other uses +// +// The error model and the `Status` message can be used in a variety of +// environments, either with or without APIs, to provide a +// consistent developer experience across different environments. +// +// Example uses of this error model include: +// +// - Partial errors. If a service needs to return partial errors to the client, +// it may embed the `Status` in the normal response to indicate the partial +// errors. +// +// - Workflow errors. A typical workflow has multiple steps. Each step may +// have a `Status` message for error reporting. +// +// - Batch operations. If a client uses batch request and batch response, the +// `Status` message should be used directly inside batch response, one for +// each error sub-response. +// +// - Asynchronous operations. If an API call embeds asynchronous operation +// results in its response, the status of those operations should be +// represented directly using the `Status` message. +// +// - Logging. If some API errors are stored in logs, the message `Status` could +// be used directly after any stripping needed for security/privacy reasons. +message Status { + // The status code, which should be an enum value of [google.rpc.Code][google.rpc.Code]. + int32 code = 1; + + // A developer-facing error message, which should be in English. Any + // user-facing error message should be localized and sent in the + // [google.rpc.Status.details][google.rpc.Status.details] field, or localized by the client. + string message = 2; + + // A list of messages that carry the error details. There is a common set of + // message types for APIs to use. + repeated google.protobuf.Any details = 3; +} diff --git a/google/cloud/sql/connector/proto/sql_data_service.proto b/google/cloud/sql/connector/proto/sql_data_service.proto new file mode 100644 index 00000000..98d688cd --- /dev/null +++ b/google/cloud/sql/connector/proto/sql_data_service.proto @@ -0,0 +1,264 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto3"; + +package google.cloud.sql.v1beta4; + +option go_package = "internal/sqldata"; +option java_package = "com.google.cloud.sql.v1beta4"; +option java_outer_classname = "CloudSqlDataProto"; +option java_multiple_files = true; + +import "google/rpc/status.proto"; + +// Service for streaming data to and from Cloud SQL instances. +service SqlDataService { + // `StreamSqlData` establishes a bidirectional stream to a Cloud SQL instance, + // and then streams data to and from the instance. + // + // The first message from the client MUST be a `StreamSqlDataRequest` request + // with configuration settings, including required values for the + // `connection_settings` field. Subsequent messages from the client may + // contain the `payload` field. + // + // Messages from the server may contain the `payload` field. + // + // The `payload` fields of the request and response streams contain the raw + // data of the database's native wire protocol (e.g., PostgreSQL wire + // protocol). The database client is responsible for generating and parsing + // this data. + // + // Any errors on initial connection (e.g., connection failure, authorization + // issues, network problems) will result in the stream being terminated with + // an appropriate RPC status exception. + // + // After a successful connection is made, if an error occurs, then the server + // terminates connection and returns the appropriate RPC status exception. + rpc StreamSqlData(stream StreamSqlDataRequest) + returns (stream StreamSqlDataResponse) {} +} + +// Message sent from the client to `SqlDataService`. +message StreamSqlDataRequest { + // Deprecated: Use `StartSession.location_id` or `ContinueSession.location_id` + // instead. `location_id` is used to route the request to a specific region. + // Use the same region which was used to create the instance. Use the format + // `locations/{location}`, for example: `locations/us-central1`. + string location_id = 1; + + // Deprecated: Use the `message` oneof instead. The type of message sent + // within the stream. + oneof message_type { + // Deprecated: Use `start_session` or `continue_session` instead. + // Parameters for establishing the connection. MUST be sent as the first + // message on the stream. + ConnectionSettings connection_settings = 2; + + // Deprecated: Use `DataPacket` instead. + // Data to be forwarded to the database. + ClientPayload payload = 3; + } + + // Acknowledges data received by the client. + Ack ack = 4; + + // The message to the server. + oneof message { + // Starts a new session. When starting a new session, this is the first + // message the client sends. + StartSession start_session = 5; + // Continues an existing session. When starting a new session, this is the + // first message the client sends. + ContinueSession continue_session = 6; + // Database data. + DataPacket data = 7; + // Terminates the session. This closes the connection to the database. + TerminateSession terminate_session = 8; + } +} + +// Deprecated: New schema structure. Initial connection parameters. +message ConnectionSettings { + option deprecated = true; + + // The target of the connection. + oneof target { + // The identifier of the Cloud SQL instance. + InstanceId instance_id = 1; + } +} + +// Start a new session. The client must send this as the first message to the +// server to start a new session. The client may immediately send Data messages +// without waiting for a reply from the server. +message StartSession { + //`location_id` is used to route the + // request to a specific region. Use the same region which was used to create + // the instance. Use the format `locations/{location}`, for example: + // `locations/us-central1`. + string location_id = 1; + // The Cloud SQL instance resource name, for example: + // projects/example-project/instances/example-instance + string instance_id = 2; + // The session id, chosen by the client. This should be an unguessable string. + // If the client does not intend to reconnect to this session, the client may + // leave session_id unset. + string session_id = 3; +} + +// Reconnects to an existing session. The client must send this as the first +// message to the server to reconnect to an existing session. The client may +// immediately send Data messages without waiting for a reply from the server. +message ContinueSession { + //`location_id` is used to route the + // request to a specific region. Use the same region which was used to create + // the instance. Use the format `locations/{location}`, for example: + // `locations/us-central1`. + string location_id = 1; + + // The Cloud SQL instance resource name, for example: + // projects/example-project/instances/example-instance + string instance_id = 2; + + // The id of the session to reconnect. + string session_id = 3; +} + +// Deprecated: New schema structure. The identifier of the Cloud SQL instance. +message InstanceId { + option deprecated = true; + + // Full resource name of the Cloud SQL instance, in the form: + // `projects/{project}/instances/{instance}`, for example: + // `projects/foo-project/instances/bar-instance`. + string instance = 1; +} + +// Deprecated: New schema structure. Wrapper for data being sent to the +// database. +message ClientPayload { + option deprecated = true; + + // Raw data to be sent to the database. See the documentation for + // `StreamSqlData` for details on the expected wire format. + bytes data = 1; +} + +// Message sent from SqlDataService back to the client. +message StreamSqlDataResponse { + // Deprecated: New schema structure. The type of the message received from + // `SqlDataService`. + oneof type { + // Raw data received from the database. + ServerPayload payload = 1; + } + + // Acknowledges data received by the server. + Ack ack = 2; + // A message from the server to the client. + oneof message { + // The first message from the server to the client, containing metadata + // about this session. + SessionMetadata session_metadata = 3; + // Data from the database. + DataPacket data = 4; + // Terminates the session. This indicates that the database connection + // is closed. When the client receives this message, it should not + // attempt to reconnect. + TerminateSession terminate_session = 5; + } +} + +// Deprecated: New schema structure. Wrapper for data being received from the +// database. +message ServerPayload { + option deprecated = true; + + // Raw data received from the database. See the documentation for + // `StreamSqlData` for details on the expected wire format. + bytes data = 1; +} +// Metadata from the server to the client about the session. The server will +// always send this as the first message +message SessionMetadata { + // The features supported by the server for this session. This field is used + // by the client to determine which features are available on the server. + // The features supported by the server for this session. + repeated SqlDataFeature supported_features = 1; +} + +// Contains data being sent or received by the database. +message DataPacket { + // The absolute byte offset of the first byte in this payload. + // 0 for new connections or resumed connections that hasn't acked any bytes + // from server. Non-zero for resumed connections + int64 first_byte_offset = 1; + // Raw data being sent or received by the database. + bytes data = 2; +} +// Acknowledges data received by the client or server. +message Ack { + // The absolute number of bytes processed in the session. + int64 received_offset = 1; +} +// Indicates that the session is permanently ended. +message TerminateSession { + // The session termination status. + google.rpc.Status status = 1; +} + +// Error reasons for `StreamSqlData`. +// Typically used with standard error codes, with the error info/reason field +// set to the string representation of the enum value. +enum StreamSqlDataErrorReason { + // Indicates that the error reason is unknown. + STREAM_SQL_DATA_ERROR_REASON_UNKNOWN = 0; + + // Indicates that the operation is not supported for given instance type. + // Used with status code `google.rpc.Code.FAILED_PRECONDITION`. + STREAM_SQL_DATA_ERROR_REASON_UNSUPPORTED_INSTANCE_TYPE = 1; + + // Indicates that reconnect failed and should not be retried. + // Used with status code `google.rpc.Code.INTERNAL`. + STREAM_SQL_DATA_ERROR_REASON_RECONNECT_FAILED = 3; + + // Indicates that the database client closed its connection normally. + // Used with status code `google.rpc.Code.CANCELED`. + STREAM_SQL_DATA_ERROR_REASON_DB_CLIENT_CLOSED = 4; + + // Indicates that the database server closed its connection normally. + // Used with status code `google.rpc.Code.CANCELED`. + STREAM_SQL_DATA_ERROR_REASON_DB_SERVER_CLOSED = 5; + + // Indicates that the peer sent an ACK message that was not within an + // acceptable range. Used with the status code + // `google.rpc.Code.FAILED_PRECONDITION`. + STREAM_SQL_DATA_ERROR_REASON_INVALID_ACK = 6; + + // Indicates that the SqlDataService lost its connection to the + // database instance. This is a retryable error. + // Used with status code `google.rpc.Code.ABORTED`. + STREAM_SQL_DATA_ERROR_REASON_DISCONNECTED = 7; +} + +// The session features. The server must send the supported features in its +// first message to the client. +enum SqlDataFeature { + // The feature is not specified. This value should not be used. + SQL_DATA_FEATURE_UNSPECIFIED = 0; + // The server supports reconnecting to the session. If this feature is not + // present, the client should not attempt to reconnect to the session. + SQL_DATA_FEATURE_RECONNECT = 1; +} diff --git a/google/cloud/sql/connector/proto/sql_data_service_pb2.py b/google/cloud/sql/connector/proto/sql_data_service_pb2.py new file mode 100644 index 00000000..fb25be21 --- /dev/null +++ b/google/cloud/sql/connector/proto/sql_data_service_pb2.py @@ -0,0 +1,88 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: google/cloud/sql/connector/proto/sql_data_service.proto +# Protobuf Python Version: 6.33.5 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder + +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 33, + 5, + '', + 'google/cloud/sql/connector/proto/sql_data_service.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n7google/cloud/sql/connector/proto/sql_data_service.proto\x12\x18google.cloud.sql.v1beta4\x1a\x17google/rpc/status.proto\"\x82\x04\n\x14StreamSqlDataRequest\x12\x13\n\x0blocation_id\x18\x01 \x01(\t\x12K\n\x13\x63onnection_settings\x18\x02 \x01(\x0b\x32,.google.cloud.sql.v1beta4.ConnectionSettingsH\x00\x12:\n\x07payload\x18\x03 \x01(\x0b\x32\'.google.cloud.sql.v1beta4.ClientPayloadH\x00\x12*\n\x03\x61\x63k\x18\x04 \x01(\x0b\x32\x1d.google.cloud.sql.v1beta4.Ack\x12?\n\rstart_session\x18\x05 \x01(\x0b\x32&.google.cloud.sql.v1beta4.StartSessionH\x01\x12\x45\n\x10\x63ontinue_session\x18\x06 \x01(\x0b\x32).google.cloud.sql.v1beta4.ContinueSessionH\x01\x12\x34\n\x04\x64\x61ta\x18\x07 \x01(\x0b\x32$.google.cloud.sql.v1beta4.DataPacketH\x01\x12G\n\x11terminate_session\x18\x08 \x01(\x0b\x32*.google.cloud.sql.v1beta4.TerminateSessionH\x01\x42\x0e\n\x0cmessage_typeB\t\n\x07message\"_\n\x12\x43onnectionSettings\x12;\n\x0binstance_id\x18\x01 \x01(\x0b\x32$.google.cloud.sql.v1beta4.InstanceIdH\x00:\x02\x18\x01\x42\x08\n\x06target\"L\n\x0cStartSession\x12\x13\n\x0blocation_id\x18\x01 \x01(\t\x12\x13\n\x0binstance_id\x18\x02 \x01(\t\x12\x12\n\nsession_id\x18\x03 \x01(\t\"O\n\x0f\x43ontinueSession\x12\x13\n\x0blocation_id\x18\x01 \x01(\t\x12\x13\n\x0binstance_id\x18\x02 \x01(\t\x12\x12\n\nsession_id\x18\x03 \x01(\t\"\"\n\nInstanceId\x12\x10\n\x08instance\x18\x01 \x01(\t:\x02\x18\x01\"!\n\rClientPayload\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c:\x02\x18\x01\"\xd8\x02\n\x15StreamSqlDataResponse\x12:\n\x07payload\x18\x01 \x01(\x0b\x32\'.google.cloud.sql.v1beta4.ServerPayloadH\x00\x12*\n\x03\x61\x63k\x18\x02 \x01(\x0b\x32\x1d.google.cloud.sql.v1beta4.Ack\x12\x45\n\x10session_metadata\x18\x03 \x01(\x0b\x32).google.cloud.sql.v1beta4.SessionMetadataH\x01\x12\x34\n\x04\x64\x61ta\x18\x04 \x01(\x0b\x32$.google.cloud.sql.v1beta4.DataPacketH\x01\x12G\n\x11terminate_session\x18\x05 \x01(\x0b\x32*.google.cloud.sql.v1beta4.TerminateSessionH\x01\x42\x06\n\x04typeB\t\n\x07message\"!\n\rServerPayload\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c:\x02\x18\x01\"W\n\x0fSessionMetadata\x12\x44\n\x12supported_features\x18\x01 \x03(\x0e\x32(.google.cloud.sql.v1beta4.SqlDataFeature\"5\n\nDataPacket\x12\x19\n\x11\x66irst_byte_offset\x18\x01 \x01(\x03\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"\x1e\n\x03\x41\x63k\x12\x17\n\x0freceived_offset\x18\x01 \x01(\x03\"6\n\x10TerminateSession\x12\"\n\x06status\x18\x01 \x01(\x0b\x32\x12.google.rpc.Status*\xf6\x02\n\x18StreamSqlDataErrorReason\x12(\n$STREAM_SQL_DATA_ERROR_REASON_UNKNOWN\x10\x00\x12:\n6STREAM_SQL_DATA_ERROR_REASON_UNSUPPORTED_INSTANCE_TYPE\x10\x01\x12\x31\n-STREAM_SQL_DATA_ERROR_REASON_RECONNECT_FAILED\x10\x03\x12\x31\n-STREAM_SQL_DATA_ERROR_REASON_DB_CLIENT_CLOSED\x10\x04\x12\x31\n-STREAM_SQL_DATA_ERROR_REASON_DB_SERVER_CLOSED\x10\x05\x12,\n(STREAM_SQL_DATA_ERROR_REASON_INVALID_ACK\x10\x06\x12-\n)STREAM_SQL_DATA_ERROR_REASON_DISCONNECTED\x10\x07*R\n\x0eSqlDataFeature\x12 \n\x1cSQL_DATA_FEATURE_UNSPECIFIED\x10\x00\x12\x1e\n\x1aSQL_DATA_FEATURE_RECONNECT\x10\x01\x32\x88\x01\n\x0eSqlDataService\x12v\n\rStreamSqlData\x12..google.cloud.sql.v1beta4.StreamSqlDataRequest\x1a/.google.cloud.sql.v1beta4.StreamSqlDataResponse\"\x00(\x01\x30\x01\x42\x45\n\x1c\x63om.google.cloud.sql.v1beta4B\x11\x43loudSqlDataProtoP\x01Z\x10internal/sqldatab\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.cloud.sql.connector.proto.sql_data_service_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'\n\034com.google.cloud.sql.v1beta4B\021CloudSqlDataProtoP\001Z\020internal/sqldata' + _globals['_CONNECTIONSETTINGS']._loaded_options = None + _globals['_CONNECTIONSETTINGS']._serialized_options = b'\030\001' + _globals['_INSTANCEID']._loaded_options = None + _globals['_INSTANCEID']._serialized_options = b'\030\001' + _globals['_CLIENTPAYLOAD']._loaded_options = None + _globals['_CLIENTPAYLOAD']._serialized_options = b'\030\001' + _globals['_SERVERPAYLOAD']._loaded_options = None + _globals['_SERVERPAYLOAD']._serialized_options = b'\030\001' + _globals['_STREAMSQLDATAERRORREASON']._serialized_start=1569 + _globals['_STREAMSQLDATAERRORREASON']._serialized_end=1943 + _globals['_SQLDATAFEATURE']._serialized_start=1945 + _globals['_SQLDATAFEATURE']._serialized_end=2027 + _globals['_STREAMSQLDATAREQUEST']._serialized_start=111 + _globals['_STREAMSQLDATAREQUEST']._serialized_end=625 + _globals['_CONNECTIONSETTINGS']._serialized_start=627 + _globals['_CONNECTIONSETTINGS']._serialized_end=722 + _globals['_STARTSESSION']._serialized_start=724 + _globals['_STARTSESSION']._serialized_end=800 + _globals['_CONTINUESESSION']._serialized_start=802 + _globals['_CONTINUESESSION']._serialized_end=881 + _globals['_INSTANCEID']._serialized_start=883 + _globals['_INSTANCEID']._serialized_end=917 + _globals['_CLIENTPAYLOAD']._serialized_start=919 + _globals['_CLIENTPAYLOAD']._serialized_end=952 + _globals['_STREAMSQLDATARESPONSE']._serialized_start=955 + _globals['_STREAMSQLDATARESPONSE']._serialized_end=1299 + _globals['_SERVERPAYLOAD']._serialized_start=1301 + _globals['_SERVERPAYLOAD']._serialized_end=1334 + _globals['_SESSIONMETADATA']._serialized_start=1336 + _globals['_SESSIONMETADATA']._serialized_end=1423 + _globals['_DATAPACKET']._serialized_start=1425 + _globals['_DATAPACKET']._serialized_end=1478 + _globals['_ACK']._serialized_start=1480 + _globals['_ACK']._serialized_end=1510 + _globals['_TERMINATESESSION']._serialized_start=1512 + _globals['_TERMINATESESSION']._serialized_end=1566 + _globals['_SQLDATASERVICE']._serialized_start=2030 + _globals['_SQLDATASERVICE']._serialized_end=2166 +# @@protoc_insertion_point(module_scope) diff --git a/google/cloud/sql/connector/proto/sql_data_service_pb2_grpc.py b/google/cloud/sql/connector/proto/sql_data_service_pb2_grpc.py new file mode 100644 index 00000000..42240ecb --- /dev/null +++ b/google/cloud/sql/connector/proto/sql_data_service_pb2_grpc.py @@ -0,0 +1,137 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" + +import grpc + +from google.cloud.sql.connector.proto import ( + sql_data_service_pb2 as google_dot_cloud_dot_sql_dot_connector_dot_proto_dot_sql__data__service__pb2, +) + +GRPC_GENERATED_VERSION = '1.81.0' +GRPC_VERSION = grpc.__version__ +_version_not_supported = False + +try: + from grpc._utilities import first_version_is_lower + _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) +except ImportError: + _version_not_supported = True + +if _version_not_supported: + raise RuntimeError( + f'The grpc package installed is at version {GRPC_VERSION},' + + ' but the generated code in google/cloud/sql/connector/proto/sql_data_service_pb2_grpc.py depends on' + + f' grpcio>={GRPC_GENERATED_VERSION}.' + + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' + ) + + +class SqlDataServiceStub: + """Service for streaming data to and from Cloud SQL instances. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.StreamSqlData = channel.stream_stream( + '/google.cloud.sql.v1beta4.SqlDataService/StreamSqlData', + request_serializer=google_dot_cloud_dot_sql_dot_connector_dot_proto_dot_sql__data__service__pb2.StreamSqlDataRequest.SerializeToString, + response_deserializer=google_dot_cloud_dot_sql_dot_connector_dot_proto_dot_sql__data__service__pb2.StreamSqlDataResponse.FromString, + _registered_method=True) + + +class SqlDataServiceServicer: + """Service for streaming data to and from Cloud SQL instances. + """ + + def StreamSqlData(self, request_iterator, context): + """`StreamSqlData` establishes a bidirectional stream to a Cloud SQL instance, + and then streams data to and from the instance. + + The first message from the client MUST be a `StreamSqlDataRequest` request + with configuration settings, including required values for the + `connection_settings` field. Subsequent messages from the client may + contain the `payload` field. + + Messages from the server may contain the `payload` field. + + The `payload` fields of the request and response streams contain the raw + data of the database's native wire protocol (e.g., PostgreSQL wire + protocol). The database client is responsible for generating and parsing + this data. + + Any errors on initial connection (e.g., connection failure, authorization + issues, network problems) will result in the stream being terminated with + an appropriate RPC status exception. + + After a successful connection is made, if an error occurs, then the server + terminates connection and returns the appropriate RPC status exception. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_SqlDataServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'StreamSqlData': grpc.stream_stream_rpc_method_handler( + servicer.StreamSqlData, + request_deserializer=google_dot_cloud_dot_sql_dot_connector_dot_proto_dot_sql__data__service__pb2.StreamSqlDataRequest.FromString, + response_serializer=google_dot_cloud_dot_sql_dot_connector_dot_proto_dot_sql__data__service__pb2.StreamSqlDataResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'google.cloud.sql.v1beta4.SqlDataService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('google.cloud.sql.v1beta4.SqlDataService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class SqlDataService: + """Service for streaming data to and from Cloud SQL instances. + """ + + @staticmethod + def StreamSqlData(request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.stream_stream( + request_iterator, + target, + '/google.cloud.sql.v1beta4.SqlDataService/StreamSqlData', + google_dot_cloud_dot_sql_dot_connector_dot_proto_dot_sql__data__service__pb2.StreamSqlDataRequest.SerializeToString, + google_dot_cloud_dot_sql_dot_connector_dot_proto_dot_sql__data__service__pb2.StreamSqlDataResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/google/cloud/sql/connector/sqldata_client.py b/google/cloud/sql/connector/sqldata_client.py new file mode 100644 index 00000000..a8aeeb87 --- /dev/null +++ b/google/cloud/sql/connector/sqldata_client.py @@ -0,0 +1,511 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import logging +import socket +from typing import Any, Callable + +from google.auth.credentials import Credentials +from google.auth.transport.grpc import AuthMetadataPlugin +from google.auth.transport.requests import Request +import grpc + +from google.cloud.sql.connector.enums import IPTypes +from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError + +import google.rpc.status_pb2 # noqa: F401 # isort: skip +from google.cloud.sql.connector.proto import sql_data_service_pb2 # type: ignore +from google.cloud.sql.connector.proto import sql_data_service_pb2_grpc # type: ignore + +SERVER_PROXY_PORT = 3307 + +logger = logging.getLogger(__name__) + + +def is_resource_exhausted_error(err: Exception) -> bool: + """Checks whether an exception represents a gRPC RESOURCE_EXHAUSTED error.""" + if isinstance(err, (grpc.aio.AioRpcError, grpc.RpcError)): + try: + return err.code() == grpc.StatusCode.RESOURCE_EXHAUSTED + except Exception: # noqa: BLE001, S110 + pass + if hasattr(err, "code") and callable(err.code): + try: + return err.code() == grpc.StatusCode.RESOURCE_EXHAUSTED + except Exception: # noqa: BLE001, S110 + pass + cause = getattr(err, "__cause__", None) or getattr(err, "__context__", None) + if isinstance(cause, Exception) and cause is not err: + return is_resource_exhausted_error(cause) + return False + + +class SqlDataClient: + def __init__( + self, + endpoint: str, + credentials: Credentials, + quota_project: str | None = None, + timeout: float | None = None, + ): + self._endpoint = endpoint + self._credentials = credentials + self._quota_project = quota_project + self._timeout = timeout + self._server: asyncio.Server | None = None + self._tunnel_tasks: set[asyncio.Task] = set() + self._active_grpc_channels: set[grpc.aio.Channel] = set() + self._active_writers: set[asyncio.StreamWriter] = set() + self._on_close_callbacks: list[Callable[[], None]] = [] + + async def connect_tunnel( + self, + instance_connection_name: str, + region: str, + project: str, + get_conn_info: Callable[[], Any], + enable_iam_auth: bool, + on_fallback: Callable[[str], None], + is_fallback_cached: Callable[[str], bool], + on_resource_exhausted: Callable[[Exception], None] | None = None, + on_success: Callable[[], None] | None = None, + connect_timeout: float = 30.0, + ) -> int: + """Starts a local TCP tunnel and returns the local port. + + If the instance does not support SQL Data Service, it falls back + to a direct TLS connection. + """ + # Start local TCP server + server = await asyncio.start_server( + lambda r, w: self._handle_tunnel( + r, + w, + instance_connection_name, + region, + project, + get_conn_info, + enable_iam_auth, + on_fallback, + is_fallback_cached, + on_resource_exhausted, + on_success, + connect_timeout, + ), + "127.0.0.1", + 0, + ) + + port = server.sockets[0].getsockname()[1] + logger.debug(f"SQL Data tunnel listening on 127.0.0.1:{port}") + + # Keep reference to server to close it + self._server = server + return port + + async def close(self) -> None: + """Closes the local tunnel server, active streams, and channels.""" + if self._server: + self._server.close() + try: + await asyncio.wait_for(self._server.wait_closed(), timeout=0.5) + logger.debug("SQL Data tunnel server closed by client close()") + except (asyncio.TimeoutError, Exception) as e: # noqa: BLE001 + logger.debug(f"Tunnel server wait_closed finished or timed out: {e}") + self._server = None + + for task in list(self._tunnel_tasks): + task.cancel() + + for channel in list(self._active_grpc_channels): + try: + await channel.close() + except Exception as e: # noqa: BLE001 + logger.debug(f"Error closing gRPC channel: {e}") + self._active_grpc_channels.clear() + + for writer in list(self._active_writers): + try: + writer.close() + except Exception as e: # noqa: BLE001 + logger.debug(f"Error closing stream writer: {e}") + self._active_writers.clear() + + for cb in self._on_close_callbacks: + try: + cb() + except Exception: # noqa: BLE001, S110 + pass + + async def _open_direct_connection( + self, + target_ip: str, + port: int, + ssl_context: Any, + connect_timeout: float, + ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + return await asyncio.wait_for( + asyncio.open_connection( + target_ip, port, ssl=ssl_context, server_hostname=target_ip + ), + timeout=connect_timeout, + ) + + async def _handle_tunnel( + self, + client_reader: asyncio.StreamReader, + client_writer: asyncio.StreamWriter, + instance_connection_name: str, + region: str, + project: str, + get_conn_info: Callable[[], Any], + enable_iam_auth: bool, + on_fallback: Callable[[str], None], + is_fallback_cached: Callable[[str], bool], + on_resource_exhausted: Callable[[Exception], None] | None = None, + on_success: Callable[[], None] | None = None, + connect_timeout: float = 30.0, + ): + logger.debug("Accepted local connection for SQL Data tunnel") + # Close the server so no more connections are accepted on this port + if self._server: + self._server.close() + self._active_writers.add(client_writer) + + t_client: asyncio.Task | None = None + t_backend: asyncio.Task | None = None + grpc_channel: grpc.aio.Channel | None = None + backend_writer: asyncio.StreamWriter | None = None + backend_reader: asyncio.StreamReader | None = None + grpc_stream: Any | None = None + client_write_buffer = bytearray() + first_read_done = False + fallback_triggered = False + fallback_ready = asyncio.Event() + + # Check if fallback is already cached + use_fallback = is_fallback_cached(instance_connection_name) + + async def connect_grpc() -> tuple[grpc.aio.Channel, Any]: + auth_request = Request() + plugin = AuthMetadataPlugin(self._credentials, auth_request) + call_creds = grpc.metadata_call_credentials(plugin) + channel_creds = grpc.composite_channel_credentials( + grpc.ssl_channel_credentials(), call_creds + ) + + endpoint = self._endpoint.removeprefix("https://").removeprefix("http://") + + logger.debug(f"Creating secure channel to {endpoint}") + channel = grpc.aio.secure_channel(endpoint, channel_creds) + self._active_grpc_channels.add(channel) + stub = sql_data_service_pb2_grpc.SqlDataServiceStub(channel) + + instance_id = f"projects/{project}/instances/{instance_connection_name.split(':')[-1]}" + location_id = f"locations/{region}" + + metadata = [] + quota_project_in_creds = getattr(self._credentials, "quota_project_id", None) + if self._quota_project and self._quota_project != quota_project_in_creds: + metadata.append(("x-goog-user-project", self._quota_project)) + metadata.append( + ( + "x-goog-request-params", + f"instance_id={instance_id}&location_id={location_id}", + ) + ) + + # Start stream + logger.debug(f"Starting StreamSqlData with metadata {metadata}") + stream = stub.StreamSqlData(metadata=metadata, timeout=self._timeout) + + # Send StartSession + start_session = sql_data_service_pb2.StartSession( # type: ignore[attr-defined] + instance_id=instance_id, location_id=location_id + ) + req = sql_data_service_pb2.StreamSqlDataRequest( # type: ignore[attr-defined] + start_session=start_session + ) + logger.debug("Writing StartSession to stream...") + await stream.write(req) + logger.debug("StartSession written successfully") + return channel, stream + + async def connect_direct() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + logger.debug("Fallback triggered, fetching connection info...") + conn_info = await get_conn_info() + # Find a fallback IP address, prioritizing PRIVATE, PSC, PUBLIC + targets: list[str] = [] + for t in [IPTypes.PRIVATE, IPTypes.PSC, IPTypes.PUBLIC]: + try: + targets.extend(conn_info.get_preferred_ips(t)) + except CloudSQLIPTypeError as e: + logger.debug(f"IP type {t} not available: {e}") + continue + if not targets: + raise ValueError("Cannot fallback to direct connection: no IP address available.") + ssl_context = await conn_info.create_ssl_context(enable_iam_auth) + last_ex: Exception | None = None + for target_ip in targets: + logger.debug(f"Connecting directly to {target_ip}:{SERVER_PROXY_PORT}") + try: + r, w = await self._open_direct_connection( + target_ip, + SERVER_PROXY_PORT, + ssl_context, + connect_timeout, + ) + self._active_writers.add(w) + return r, w + except Exception as e: # noqa: BLE001 + logger.debug(f"Direct connection to {target_ip} failed: {e}") + last_ex = e + if last_ex: + raise last_ex + raise ValueError("Cannot fallback to direct connection: no IP address available.") + + # Task to read from client and write to backend + async def client_to_backend(): + nonlocal first_read_done, fallback_triggered, backend_writer, grpc_stream + try: + while True: + data = await client_reader.read(4096) + if not data: + logger.debug("Client socket EOF") + break + + if not first_read_done and not fallback_triggered: + client_write_buffer.extend(data) + + if fallback_triggered: + await fallback_ready.wait() + if backend_writer: + backend_writer.write(data) + await backend_writer.drain() + else: + packet = sql_data_service_pb2.DataPacket(data=data) # type: ignore[attr-defined] + req = sql_data_service_pb2.StreamSqlDataRequest( # type: ignore[attr-defined] + data=packet + ) + if grpc_stream: + try: + await grpc_stream.write(req) + except Exception as e: + if is_resource_exhausted_error(e): + if on_resource_exhausted: + on_resource_exhausted(e) + raise + if fallback_triggered or not first_read_done: + logger.debug( + f"Write to gRPC stream failed while fallback pending or triggered: {e}" + ) + else: + raise + except Exception as e: + logger.error(f"Error in client_to_backend: {e}") + if is_resource_exhausted_error(e) and on_resource_exhausted: + on_resource_exhausted(e) + raise + finally: + if fallback_triggered: + if backend_writer: + backend_writer.write_eof() + else: + if grpc_stream: + try: + await grpc_stream.done_writing() + except Exception as e: # noqa: BLE001 + logger.debug(f"Error calling done_writing: {e}") + logger.debug("Client to backend task finished") + + # Task to read from backend and write to client + async def backend_to_client(): + nonlocal first_read_done, fallback_triggered, backend_reader, backend_writer, grpc_stream, grpc_channel + try: + if fallback_triggered: + # If we started with fallback, just copy + while True: + if not backend_reader: + break + data = await backend_reader.read(4096) + if not data: + break + client_writer.write(data) + await client_writer.drain() + else: + # gRPC read loop + try: + if not grpc_stream: + return + async for resp in grpc_stream: + if not first_read_done: + first_read_done = True + if on_success: + on_success() + msg_type = resp.WhichOneof("message") + if msg_type == "session_metadata": + logger.debug("Received SessionMetadata") + elif msg_type == "data": + data = resp.data.data + logger.debug(f"Received {len(data)} bytes") + client_writer.write(data) + await client_writer.drain() + elif msg_type == "terminate_session": + logger.debug("Received TerminateSession") + break + except grpc.aio.AioRpcError as e: + logger.debug(f"gRPC stream error: {e}") + if is_resource_exhausted_error(e): + if on_resource_exhausted: + on_resource_exhausted(e) + raise + # Check for fallback condition + if ( + not first_read_done + and e.code() == grpc.StatusCode.FAILED_PRECONDITION + ): + logger.info( + f"SQL Data Service not supported for {instance_connection_name}. " + "Falling back to direct connection." + ) + fallback_triggered = True + on_fallback(instance_connection_name) + + # Clean up gRPC + if grpc_channel: + await grpc_channel.close() + + # Connect direct + backend_reader, backend_writer = await connect_direct() + + # Replay buffered client data + if client_write_buffer: + logger.debug(f"Replaying {len(client_write_buffer)} bytes to fallback connection") + backend_writer.write(bytes(client_write_buffer)) + await backend_writer.drain() + + fallback_ready.set() + + # Start copying from direct connection + while True: + data = await backend_reader.read(4096) + if not data: + break + client_writer.write(data) + await client_writer.drain() + else: + # Other gRPC error, re-raise to close connection + raise + except Exception as e: + logger.error(f"Error in backend_to_client: {e}") + raise + finally: + client_writer.close() + try: + await client_writer.wait_closed() + except Exception as e: # noqa: BLE001 + logger.debug(f"Error waiting for client writer to close: {e}") + if fallback_triggered and backend_writer: + backend_writer.close() + try: + await backend_writer.wait_closed() + except Exception as e: # noqa: BLE001 + logger.debug(f"Error waiting for backend writer to close: {e}") + elif grpc_channel: + await grpc_channel.close() + logger.debug("Backend to client task finished") + + try: + # Initialize connection + if use_fallback: + logger.debug("Using cached fallback connection") + backend_reader, backend_writer = await connect_direct() + fallback_triggered = True + fallback_ready.set() + else: + try: + grpc_channel, grpc_stream = await connect_grpc() + except Exception as e: + logger.debug(f"Failed to initialize gRPC stream: {e}") + if is_resource_exhausted_error(e): + if on_resource_exhausted: + on_resource_exhausted(e) + raise + # Try fallback immediately for non-resource-exhausted errors + backend_reader, backend_writer = await connect_direct() + fallback_triggered = True + fallback_ready.set() + on_fallback(instance_connection_name) + + # Run both tasks with explicit lifecycle and cancellation management + t_client = asyncio.create_task(client_to_backend()) + t_backend = asyncio.create_task(backend_to_client()) + self._tunnel_tasks.add(t_client) + self._tunnel_tasks.add(t_backend) + done, pending = await asyncio.wait( + [t_client, t_backend], + return_when=asyncio.FIRST_EXCEPTION, + ) + for task in pending: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + except Exception as e: # noqa: BLE001 + logger.debug(f"Error awaiting cancelled tunnel task: {e}") + for task in done: + if not task.cancelled(): + exc = task.exception() + if exc is not None: + raise exc + finally: + if t_client: + self._tunnel_tasks.discard(t_client) + if t_backend: + self._tunnel_tasks.discard(t_backend) + if grpc_channel: + self._active_grpc_channels.discard(grpc_channel) + try: + await grpc_channel.close() + except Exception as e: # noqa: BLE001 + logger.debug(f"Error closing gRPC channel: {e}") + self._active_writers.discard(client_writer) + if backend_writer: + self._active_writers.discard(backend_writer) + logger.debug("Closing client socket in _handle_tunnel finally") + try: + client_writer.close() + sock = client_writer.get_extra_info("socket") + if sock: + sock.close() + except Exception as e: # noqa: BLE001 + logger.debug(f"Error closing client writer: {e}") + for cb in self._on_close_callbacks: + try: + cb() + except Exception: # noqa: BLE001, S110 + pass + logger.debug("SQL Data tunnel handler finished") + + +class FallbackSocket(socket.socket): + def connect(self, *args: Any, **kwargs: Any) -> None: + # Already connected, do nothing. + # This is needed because some drivers (like pymysql) try to call connect() + # internally even if passed an already connected socket. + pass diff --git a/tests/system/test_asyncpg_connection.py b/tests/system/test_asyncpg_connection.py index 00589942..51fd17b9 100644 --- a/tests/system/test_asyncpg_connection.py +++ b/tests/system/test_asyncpg_connection.py @@ -20,6 +20,7 @@ from typing import Any import asyncpg +import pytest import sqlalchemy import sqlalchemy.ext.asyncio @@ -283,3 +284,109 @@ async def test_lazy_connection_with_asyncpg() -> None: assert res[0][0] == 1 await connector.close_async() + + +async def test_AIDE_sqlalchemy_connection_with_asyncpg() -> None: + """Basic test to get time from database using AIDE instance.""" + if "POSTGRES_AIDE_CONNECTION_NAME" not in os.environ: + pytest.skip("POSTGRES_AIDE_CONNECTION_NAME not set") + inst_conn_name = os.environ["POSTGRES_AIDE_CONNECTION_NAME"] + user = os.environ.get("POSTGRES_AIDE_USER", os.environ.get("POSTGRES_USER", "postgres")) + password = os.environ.get("POSTGRES_AIDE_PASS", os.environ.get("POSTGRES_PASS", "")) + db = os.environ.get("POSTGRES_AIDE_DB", os.environ.get("POSTGRES_DB", "postgres")) + + pool, connector = await create_sqlalchemy_engine( + inst_conn_name, + user, + password, + db, + ip_type="sqldata", + ) + + async with pool.connect() as conn: + res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone() + assert res[0] == 1 + + await connector.close_async() + + +async def test_AIDE_connection_with_asyncpg() -> None: + """Basic test to get time from database using AIDE instance.""" + if "POSTGRES_AIDE_CONNECTION_NAME" not in os.environ: + pytest.skip("POSTGRES_AIDE_CONNECTION_NAME not set") + inst_conn_name = os.environ["POSTGRES_AIDE_CONNECTION_NAME"] + user = os.environ.get("POSTGRES_AIDE_USER", os.environ.get("POSTGRES_USER", "postgres")) + password = os.environ.get("POSTGRES_AIDE_PASS", os.environ.get("POSTGRES_PASS", "")) + db = os.environ.get("POSTGRES_AIDE_DB", os.environ.get("POSTGRES_DB", "postgres")) + + pool, connector = await create_asyncpg_pool( + inst_conn_name, + user, + password, + db, + ip_type="sqldata", + ) + + async with pool.acquire() as conn: + res = await conn.fetch("SELECT 1") + assert res[0][0] == 1 + + await connector.close_async() + + +async def test_sqldata_fallback_sqlalchemy_connection_with_asyncpg() -> None: + """Test connecting to a non-AIDE instance with ip_type='sqldata'. + + The server returns FAILED_PRECONDITION for standard (non-Developer Edition) instances, + and the connector falls back to connecting via public IP. + """ + if "POSTGRES_FALLBACK_CONNECTION_NAME" not in os.environ: + pytest.skip("POSTGRES_FALLBACK_CONNECTION_NAME not set") + inst_conn_name = os.environ["POSTGRES_FALLBACK_CONNECTION_NAME"] + user = os.environ.get("POSTGRES_FALLBACK_USER", os.environ.get("POSTGRES_USER", "postgres")) + password = os.environ.get("POSTGRES_FALLBACK_PASS", os.environ.get("POSTGRES_PASS", "")) + db = os.environ.get("POSTGRES_FALLBACK_DB", os.environ.get("POSTGRES_DB", "postgres")) + + pool, connector = await create_sqlalchemy_engine( + inst_conn_name, + user, + password, + db, + ip_type="sqldata", + ) + + async with pool.connect() as conn: + res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone() + assert res[0] == 1 + + await connector.close_async() + + +async def test_sqldata_fallback_connection_with_asyncpg() -> None: + """Test connecting to a non-AIDE instance with ip_type='sqldata' via raw pool. + + The server returns FAILED_PRECONDITION for standard (non-Developer Edition) instances, + and the connector falls back to connecting via public IP. + """ + if "POSTGRES_FALLBACK_CONNECTION_NAME" not in os.environ: + pytest.skip("POSTGRES_FALLBACK_CONNECTION_NAME not set") + inst_conn_name = os.environ["POSTGRES_FALLBACK_CONNECTION_NAME"] + user = os.environ.get("POSTGRES_FALLBACK_USER", os.environ.get("POSTGRES_USER", "postgres")) + password = os.environ.get("POSTGRES_FALLBACK_PASS", os.environ.get("POSTGRES_PASS", "")) + db = os.environ.get("POSTGRES_FALLBACK_DB", os.environ.get("POSTGRES_DB", "postgres")) + + pool, connector = await create_asyncpg_pool( + inst_conn_name, + user, + password, + db, + ip_type="sqldata", + ) + + async with pool.acquire() as conn: + res = await conn.fetch("SELECT 1") + assert res[0][0] == 1 + + await connector.close_async() + + diff --git a/tests/system/test_pg8000_connection.py b/tests/system/test_pg8000_connection.py index 0fef1b96..0da6ebe2 100644 --- a/tests/system/test_pg8000_connection.py +++ b/tests/system/test_pg8000_connection.py @@ -19,6 +19,7 @@ import os # [START cloud_sql_connector_postgres_pg8000] +import pytest import sqlalchemy from google.cloud.sql.connector import Connector @@ -209,3 +210,58 @@ def test_MCP_pg8000_connection() -> None: curr_time = time[0] assert type(curr_time) is datetime connector.close() + + +def test_AIDE_pg8000_connection() -> None: + """Basic test to get time from database using AIDE instance.""" + if "POSTGRES_AIDE_CONNECTION_NAME" not in os.environ: + pytest.skip("POSTGRES_AIDE_CONNECTION_NAME not set") + inst_conn_name = os.environ["POSTGRES_AIDE_CONNECTION_NAME"] + user = os.environ.get("POSTGRES_AIDE_USER", os.environ.get("POSTGRES_USER", "postgres")) + password = os.environ.get("POSTGRES_AIDE_PASS", os.environ.get("POSTGRES_PASS", "")) + db = os.environ.get("POSTGRES_AIDE_DB", os.environ.get("POSTGRES_DB", "postgres")) + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, + user, + password, + db, + ip_type="sqldata", + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + +def test_sqldata_fallback_pg8000_connection() -> None: + """Test connecting to a non-AIDE instance with ip_type='sqldata'. + + The server returns FAILED_PRECONDITION for standard (non-Developer Edition) instances, + and the connector falls back to connecting via public IP. + """ + if "POSTGRES_FALLBACK_CONNECTION_NAME" not in os.environ: + pytest.skip("POSTGRES_FALLBACK_CONNECTION_NAME not set") + inst_conn_name = os.environ["POSTGRES_FALLBACK_CONNECTION_NAME"] + user = os.environ.get("POSTGRES_FALLBACK_USER", os.environ.get("POSTGRES_USER", "postgres")) + password = os.environ.get("POSTGRES_FALLBACK_PASS", os.environ.get("POSTGRES_PASS", "")) + db = os.environ.get("POSTGRES_FALLBACK_DB", os.environ.get("POSTGRES_DB", "postgres")) + + engine, connector = create_sqlalchemy_engine( + inst_conn_name, + user, + password, + db, + ip_type="sqldata", + ) + with engine.connect() as conn: + time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone() + conn.commit() + curr_time = time[0] + assert type(curr_time) is datetime + connector.close() + + + diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index 8ea8c88c..3a904d90 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -18,6 +18,8 @@ import asyncio import os from threading import Thread +from unittest.mock import AsyncMock +from unittest.mock import MagicMock from unittest.mock import patch from aiohttp import ClientResponseError @@ -35,6 +37,8 @@ from google.cloud.sql.connector.exceptions import IncompatibleDriverError from google.cloud.sql.connector.instance import RefreshAheadCache from google.cloud.sql.connector.resolver import DnsResolver +from google.cloud.sql.connector.sqldata_client import FallbackSocket +from google.cloud.sql.connector.sqldata_client import SqlDataClient @pytest.mark.asyncio @@ -212,6 +216,22 @@ async def test_Connector_Init_async_context_manager( IPTypes.PSC, IPTypes.PSC, ), + ( + "sqldata", + IPTypes.SQL_DATA, + ), + ( + "SQLDATA", + IPTypes.SQL_DATA, + ), + ( + "SQL_DATA", + IPTypes.SQL_DATA, + ), + ( + IPTypes.SQL_DATA, + IPTypes.SQL_DATA, + ), ], ) def test_Connector_init_ip_type( @@ -234,7 +254,7 @@ def test_Connector_Init_bad_ip_type(fake_credentials: Credentials) -> None: assert ( exc_info.value.args[0] == f"Incorrect value for ip_type, got '{bad_ip_type.upper()}'. " - "Want one of: 'PRIMARY', 'PRIVATE', 'PSC', 'PUBLIC'." + "Want one of: 'PRIMARY', 'PRIVATE', 'PSC', 'SQL_DATA', 'PUBLIC'." ) @@ -257,7 +277,7 @@ def test_Connector_connect_bad_ip_type( assert ( exc_info.value.args[0] == f"Incorrect value for ip_type, got '{bad_ip_type.upper()}'. " - "Want one of: 'PRIMARY', 'PRIVATE', 'PSC', 'PUBLIC'." + "Want one of: 'PRIMARY', 'PRIVATE', 'PSC', 'SQL_DATA', 'PUBLIC'." ) @@ -722,3 +742,428 @@ async def test_Connector_connect_async_custom_dns_resolver_no_fallback_psc_to_pr fake_client.instance.psc_enabled = False +def test_Connector_Init_sqldata_options(fake_credentials: Credentials) -> None: + """Test that Connector initializes with custom SQL data endpoint and timeout.""" + with Connector( + credentials=fake_credentials, + sql_data_endpoint="custom.sqladmin.googleapis.com", + sql_data_stream_timeout=3600, + ) as connector: + assert connector._sql_data_endpoint == "custom.sqladmin.googleapis.com" + assert connector._sql_data_stream_timeout == 3600 + + +@pytest.mark.asyncio +async def test_create_async_connector_sqldata_options( + fake_credentials: Credentials, +) -> None: + """Test that create_async_connector properly forwards SQL data options.""" + connector = await create_async_connector( + credentials=fake_credentials, + sql_data_endpoint="custom.sqladmin.googleapis.com", + sql_data_stream_timeout=1800, + ) + assert connector._sql_data_endpoint == "custom.sqladmin.googleapis.com" + assert connector._sql_data_stream_timeout == 1800 + await connector.close_async() + + +@pytest.mark.asyncio +async def test_Connector_connect_async_sqldata_iam_auth( + fake_credentials: Credentials, + fake_client: CloudSQLClient, +) -> None: + """Test that connect_async with SQL_DATA and IAM auth properly maps driver engine without KeyError.""" + connect_string = "test-project:test-region:test-instance" + async with Connector( + credentials=fake_credentials, + loop=asyncio.get_running_loop(), + ip_type=IPTypes.SQL_DATA, + ) as connector: + connector._client = fake_client + + with patch("google.cloud.sql.connector.connector.SqlDataClient") as mock_sqldata_cls: + mock_client_instance = MagicMock() + mock_client_instance.connect_tunnel = AsyncMock(return_value=3307) + mock_client_instance.close = AsyncMock() + mock_sqldata_cls.return_value = mock_client_instance + + with patch("google.cloud.sql.connector.asyncpg.connect") as mock_connect: + mock_connect.return_value = True + + connection = await connector.connect_async( + connect_string, + "asyncpg", + user="test-sa@test-project.iam.gserviceaccount.com", + db="my-db", + enable_iam_auth=True, + ) + assert connection is True + # Verify IAM user was formatted and passed without error + assert mock_connect.called + _, kwargs = mock_connect.call_args + assert kwargs["user"] == "test-sa@test-project.iam" + + +def test_sqldata_client_init(fake_credentials: Credentials) -> None: + """Test that SqlDataClient initializes with expected properties.""" + client = SqlDataClient( + endpoint="custom.sqladmin.googleapis.com", + credentials=fake_credentials, + quota_project="test-quota-project", + timeout=3600, + ) + assert client._endpoint == "custom.sqladmin.googleapis.com" + assert client._credentials == fake_credentials + assert client._quota_project == "test-quota-project" + assert client._timeout == 3600 + assert client._server is None + assert len(client._tunnel_tasks) == 0 + + +@pytest.mark.asyncio +async def test_sqldata_client_close(fake_credentials: Credentials) -> None: + """Test that SqlDataClient.close cleanly cancels tasks and closes resources.""" + client = SqlDataClient( + endpoint="sqladmin.googleapis.com", + credentials=fake_credentials, + ) + mock_server = MagicMock() + mock_server.close = MagicMock() + mock_server.wait_closed = AsyncMock() + client._server = mock_server + + mock_channel = AsyncMock() + client._active_grpc_channels.add(mock_channel) + + mock_writer = MagicMock() + client._active_writers.add(mock_writer) + + callback_called = False + + def on_close() -> None: + nonlocal callback_called + callback_called = True + + client._on_close_callbacks.append(on_close) + + async def dummy_task(): + await asyncio.sleep(100) + + task = asyncio.create_task(dummy_task()) + client._tunnel_tasks.add(task) + + await client.close() + + try: + await task + except asyncio.CancelledError: + pass + + assert client._server is None + assert mock_server.close.called + assert mock_channel.close.called + assert mock_writer.close.called + assert task.cancelled() + assert callback_called + + +def test_fallback_socket() -> None: + """Test that FallbackSocket ignores connect calls.""" + sock = FallbackSocket() + sock.connect("127.0.0.1", 3307) + sock.close() + + +@pytest.mark.asyncio +async def test_sqldata_client_connect_tunnel(fake_credentials: Credentials) -> None: + """Test that connect_tunnel binds to a local port.""" + client = SqlDataClient( + endpoint="sqladmin.googleapis.com", + credentials=fake_credentials, + ) + get_conn_info = AsyncMock() + on_fallback = MagicMock() + is_fallback_cached = MagicMock(return_value=False) + + port = await client.connect_tunnel( + instance_connection_name="proj:reg:inst", + region="reg", + project="proj", + get_conn_info=get_conn_info, + enable_iam_auth=False, + on_fallback=on_fallback, + is_fallback_cached=is_fallback_cached, + ) + assert isinstance(port, int) + assert port > 0 + await client.close() + + +def test_Connector_Init_resource_exhausted_options( + fake_credentials: Credentials, +) -> None: + """Test that Connector initializes with resource_exhausted_cooldown_period.""" + with Connector( + credentials=fake_credentials, + resource_exhausted_cooldown_period=10.0, + ) as connector: + assert connector._resource_exhausted_cooldown_period == 10.0 + + +@pytest.mark.asyncio +async def test_create_async_connector_resource_exhausted_options( + fake_credentials: Credentials, +) -> None: + """Test that create_async_connector forwards resource_exhausted_cooldown_period.""" + connector = await create_async_connector( + credentials=fake_credentials, + resource_exhausted_cooldown_period=8.5, + ) + assert connector._resource_exhausted_cooldown_period == 8.5 + await connector.close_async() + + +def test_cooldown_backoff_calculation() -> None: + """Test exponential backoff with jitter calculation.""" + from google.cloud.sql.connector.connector import _cooldown_backoff + + base = 5.0 + for attempt in range(1, 6): + backoff = _cooldown_backoff(base, attempt) + # 1.618^(attempt-1) <= multiplier <= 1.618^attempt + min_expected = base * (1.618 ** (attempt - 1)) + max_expected = base * (1.618**attempt) + assert min_expected <= backoff <= max_expected + + +def test_is_resource_exhausted_error_helper() -> None: + """Test is_resource_exhausted_error helper with various exception types.""" + import grpc + + from google.cloud.sql.connector.sqldata_client import is_resource_exhausted_error + + class MockRpcError(Exception): + def __init__(self, code): + self._code = code + + def code(self): + return self._code + + assert is_resource_exhausted_error( + MockRpcError(grpc.StatusCode.RESOURCE_EXHAUSTED) + ) + assert not is_resource_exhausted_error( + MockRpcError(grpc.StatusCode.FAILED_PRECONDITION) + ) + assert not is_resource_exhausted_error(Exception("other error")) + + # Test wrapped cause + wrapped = Exception("wrapper error") + wrapped.__cause__ = MockRpcError(grpc.StatusCode.RESOURCE_EXHAUSTED) + assert is_resource_exhausted_error(wrapped) + + +def test_SqlDataConnState_methods() -> None: + """Test SqlDataConnState state transitions and helper methods.""" + import time + + from google.cloud.sql.connector.connector import SqlDataConnState + + state = SqlDataConnState() + assert state.allowed is True + assert state.is_cooldown_active() is False + + err = Exception("resource busy") + backoff = state.record_exhausted(err, base_cooldown=2.0) + assert state.backoff_counter == 1 + assert state.last_err is err + assert state.cooldown_until is not None + assert state.cooldown_until > time.time() + assert state.is_cooldown_active() is True + assert backoff > 0 + + state.record_success() + assert state.backoff_counter == 0 + assert state.cooldown_until is None + assert state.last_err is None + assert state.is_cooldown_active() is False + + state.record_fallback() + assert state.allowed is False + assert state.is_cooldown_active() is False + + +@pytest.mark.asyncio +async def test_ResourceExhausted_cooldown_blocks_connection( + fake_credentials: Credentials, + fake_client: CloudSQLClient, +) -> None: + """Test that active cooldown raises ResourceExhaustedError without connecting.""" + import time + + from google.cloud.sql.connector.connector import SqlDataConnState + from google.cloud.sql.connector.exceptions import ResourceExhaustedError + + connect_string = "proj:reg:inst" + async with Connector( + credentials=fake_credentials, + loop=asyncio.get_running_loop(), + ip_type=IPTypes.SQL_DATA, + resource_exhausted_cooldown_period=2.0, + ) as connector: + connector._client = fake_client + + # Manually set state to cooldown active + state = SqlDataConnState() + state.cooldown_until = time.time() + 10.0 + state.backoff_counter = 1 + state.last_err = Exception("resource busy") + connector._sql_data_conn_state[connect_string] = state + + with ( + patch( + "google.cloud.sql.connector.connector.SqlDataClient" + ) as mock_sqldata_cls, + pytest.raises(ResourceExhaustedError) as exc_info, + ): + await connector.connect_async( + connect_string, + "asyncpg", + user="test-user", + db="test-db", + ) + assert "cooldown active" in str(exc_info.value) + assert not mock_sqldata_cls.called + + +@pytest.mark.asyncio +async def test_ResourceExhausted_callbacks_lifecycle( + fake_credentials: Credentials, + fake_client: CloudSQLClient, +) -> None: + """Test that on_resource_exhausted and on_success callbacks properly update state.""" + import time + + from google.cloud.sql.connector.exceptions import ResourceExhaustedError + + connect_string = "proj:reg:inst" + async with Connector( + credentials=fake_credentials, + loop=asyncio.get_running_loop(), + ip_type=IPTypes.SQL_DATA, + resource_exhausted_cooldown_period=0.5, + ) as connector: + connector._client = fake_client + + captured_on_resource_exhausted = None + captured_on_success = None + + mock_sqldata_instance = MagicMock() + + async def mock_connect_tunnel(**kwargs): + nonlocal captured_on_resource_exhausted, captured_on_success + captured_on_resource_exhausted = kwargs.get("on_resource_exhausted") + captured_on_success = kwargs.get("on_success") + return 3307 + + mock_sqldata_instance.connect_tunnel = AsyncMock( + side_effect=mock_connect_tunnel + ) + mock_sqldata_instance.close = AsyncMock() + + with ( + patch( + "google.cloud.sql.connector.connector.SqlDataClient", + return_value=mock_sqldata_instance, + ), + patch("google.cloud.sql.connector.asyncpg.connect", return_value=True), + ): + # 1. Connect and trigger on_resource_exhausted + await connector.connect_async( + connect_string, + "asyncpg", + user="test-user", + db="test-db", + ) + assert captured_on_resource_exhausted is not None + assert captured_on_success is not None + + state = connector._sql_data_conn_state[connect_string] + assert state.backoff_counter == 0 + assert state.cooldown_until is None + + # Trigger resource exhausted + captured_on_resource_exhausted(Exception("resource exhausted")) + assert state.backoff_counter == 1 + assert state.cooldown_until is not None + assert state.cooldown_until > time.time() + + # Second connect attempt during cooldown fails with ResourceExhaustedError + with pytest.raises(ResourceExhaustedError): + await connector.connect_async( + connect_string, + "asyncpg", + user="test-user", + db="test-db", + ) + + # Reset via on_success + captured_on_success() + assert state.backoff_counter == 0 + assert state.cooldown_until is None + assert state.last_err is None + + +@pytest.mark.asyncio +async def test_sqldata_fallback_ip_order(fake_credentials: Credentials) -> None: + """Test that direct fallback queries IP addresses in PRIVATE, PSC, PUBLIC order.""" + client = SqlDataClient( + endpoint="sqladmin.googleapis.com", + credentials=fake_credentials, + ) + mock_conn_info = MagicMock() + queried_ip_types: list[IPTypes] = [] + + def mock_get_preferred_ips(ip_type: IPTypes): + queried_ip_types.append(ip_type) + if ip_type == IPTypes.PUBLIC: + return ["1.2.3.4"] + from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError + + raise CloudSQLIPTypeError(f"{ip_type} not available") + + mock_conn_info.get_preferred_ips.side_effect = mock_get_preferred_ips + mock_conn_info.create_ssl_context = AsyncMock(return_value=None) + get_conn_info = AsyncMock(return_value=mock_conn_info) + + mock_reader = AsyncMock() + mock_reader.read = AsyncMock(return_value=b"") + mock_writer = MagicMock() + mock_writer.wait_closed = AsyncMock() + client._open_direct_connection = AsyncMock( + return_value=(mock_reader, mock_writer) + ) + + port = await client.connect_tunnel( + instance_connection_name="proj:reg:inst", + region="reg", + project="proj", + get_conn_info=get_conn_info, + enable_iam_auth=False, + on_fallback=MagicMock(), + is_fallback_cached=MagicMock(return_value=True), + ) + + # Trigger client connection to tunnel + _r, w = await asyncio.open_connection("127.0.0.1", port) + await asyncio.sleep(0.1) + w.close() + await w.wait_closed() + + assert queried_ip_types == [IPTypes.PRIVATE, IPTypes.PSC, IPTypes.PUBLIC] + await client.close() + + + + diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index 3dbba59b..95b00bf2 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/test_instance.py @@ -26,6 +26,7 @@ from google.cloud.sql.connector.connection_info import ConnectionInfo from google.cloud.sql.connector.connection_name import ConnectionName from google.cloud.sql.connector.exceptions import AutoIAMAuthNotSupported +from google.cloud.sql.connector.exceptions import CloudSQLConnectionError from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError from google.cloud.sql.connector.instance import RefreshAheadCache from google.cloud.sql.connector.rate_limiter import AsyncRateLimiter @@ -293,3 +294,14 @@ async def test_ConnectionInfo_caches_sslcontext() -> None: # calling create_ssl_context should no-op with an existing 'context' await info.create_ssl_context() assert info.context == "context" + + +@pytest.mark.asyncio +async def test_ConnectionInfo_missing_server_ca_cert() -> None: + """Test that create_ssl_context raises CloudSQLConnectionError when server_ca_cert is None.""" + info = ConnectionInfo( + "", "cert", None, b"key", {}, "POSTGRES", datetime.datetime.now(datetime.timezone.utc) + ) + with pytest.raises(CloudSQLConnectionError) as exc_info: + await info.create_ssl_context() + assert "server CA certificate is missing" in str(exc_info.value)