""" Copyright 2019 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. """ from __future__ import annotations import asyncio from functools import partial import logging import os import socket from threading import Thread from types import TracebackType from typing import Any, Callable, Optional, Union import google.auth from google.auth.credentials import Credentials from google.auth.credentials import with_scopes_if_required import google.cloud.sql.connector.asyncpg as asyncpg from google.cloud.sql.connector.client import CloudSQLClient 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.instance import RefreshAheadCache from google.cloud.sql.connector.lazy import LazyRefreshCache from google.cloud.sql.connector.monitored_cache import MonitoredCache import google.cloud.sql.connector.pg8000 as pg8000 import google.cloud.sql.connector.pymysql as pymysql import google.cloud.sql.connector.pytds as pytds from google.cloud.sql.connector.resolver import DefaultResolver from google.cloud.sql.connector.resolver import DnsResolver from google.cloud.sql.connector.utils import format_database_user from google.cloud.sql.connector.utils import generate_keys logger = logging.getLogger(name=__name__) ASYNC_DRIVERS = ["asyncpg"] SERVER_PROXY_PORT = 3307 _DEFAULT_SCHEME = "https://" _DEFAULT_UNIVERSE_DOMAIN = "googleapis.com" _SQLADMIN_HOST_TEMPLATE = "sqladmin.{universe_domain}" class Connector: """Configure and create secure connections to Cloud SQL.""" def __init__( self, ip_type: str | IPTypes = IPTypes.PUBLIC, enable_iam_auth: bool = False, timeout: int = 30, credentials: Optional[Credentials] = None, loop: Optional[asyncio.AbstractEventLoop] = None, quota_project: Optional[str] = None, sqladmin_api_endpoint: Optional[str] = None, user_agent: Optional[str] = None, universe_domain: Optional[str] = None, refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, resolver: type[DefaultResolver] | type[DnsResolver] = DefaultResolver, failover_period: int = 30, ) -> 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 enable_iam_auth (bool): Enables automatic IAM database authentication (Postgres and MySQL) as the default authentication method for all connections. timeout (int): The default time limit in seconds for a connection before raising a TimeoutError. credentials (google.auth.credentials.Credentials): A credentials object created from the google-auth Python library to be used. If not specified, Application Default Credentials (ADC) are used. quota_project (str): The Project ID for an existing Google Cloud project. The project specified is used for quota and billing purposes. If not specified, defaults to project sourced from environment. loop (asyncio.AbstractEventLoop): Event loop to run asyncio tasks, if not specified, defaults to creating new event loop on background thread. sqladmin_api_endpoint (str): Base URL to use when calling the Cloud SQL Admin API endpoint. Defaults to "https://sqladmin.googleapis.com", this argument should only be used in development. universe_domain (str): The universe domain for Cloud SQL API calls. Default: "googleapis.com". refresh_strategy (str | RefreshStrategy): The default refresh strategy used to refresh SSL/TLS cert and instance metadata. Can be one of the following: RefreshStrategy.LAZY ("LAZY") or RefreshStrategy.BACKGROUND ("BACKGROUND"). Default: RefreshStrategy.BACKGROUND resolver (DefaultResolver | DnsResolver): The class name of the resolver to use for resolving the Cloud SQL instance connection name. To resolve a DNS record to an instance connection name, use DnsResolver. Default: DefaultResolver failover_period (int): The time interval in seconds between each attempt to check if a failover has occured for a given instance. Must be used with `resolver=DnsResolver` to have any effect. Default: 30 """ # if refresh_strategy is str, convert to RefreshStrategy enum if isinstance(refresh_strategy, str): refresh_strategy = RefreshStrategy._from_str(refresh_strategy) self._refresh_strategy = refresh_strategy # if event loop is given, use for background tasks if loop: self._loop: asyncio.AbstractEventLoop = loop self._thread: Optional[Thread] = None # if lazy refresh is specified we should lazy init keys if self._refresh_strategy == RefreshStrategy.LAZY: self._keys: Optional[asyncio.Future] = None else: self._keys = loop.create_task(generate_keys()) # if no event loop is given, spin up new loop in background thread else: self._loop = asyncio.new_event_loop() self._thread = Thread(target=self._loop.run_forever, daemon=True) self._thread.start() # if lazy refresh is specified we should lazy init keys if self._refresh_strategy == RefreshStrategy.LAZY: self._keys = None else: self._keys = asyncio.wrap_future( asyncio.run_coroutine_threadsafe(generate_keys(), self._loop), loop=self._loop, ) # initialize dict to store caches, key is a tuple consisting of instance # connection name string and enable_iam_auth boolean flag self._cache: dict[tuple[str, bool], MonitoredCache] = {} self._client: Optional[CloudSQLClient] = None # initialize credentials scopes = ["https://www.googleapis.com/auth/sqlservice.admin"] if credentials: # verify custom credentials are proper type # and atleast base class of google.auth.credentials if not isinstance(credentials, Credentials): raise TypeError( "credentials must be of type google.auth.credentials.Credentials," f" got {type(credentials)}" ) self._credentials = with_scopes_if_required(credentials, scopes=scopes) # otherwise use application default credentials else: self._credentials, _ = google.auth.default(scopes=scopes) # set default params for connections self._timeout = timeout self._enable_iam_auth = enable_iam_auth self._user_agent = user_agent self._resolver = resolver() self._failover_period = failover_period # if ip_type is str, convert to IPTypes enum if isinstance(ip_type, str): ip_type = IPTypes._from_str(ip_type) self._ip_type = ip_type # check for quota project arg and then env var if quota_project: self._quota_project = quota_project else: self._quota_project = os.environ.get("GOOGLE_CLOUD_QUOTA_PROJECT") # type: ignore # check for universe domain arg and then env var if universe_domain: self._universe_domain = universe_domain else: self._universe_domain = os.environ.get("GOOGLE_CLOUD_UNIVERSE_DOMAIN") # type: ignore # construct service endpoint for Cloud SQL Admin API calls if not sqladmin_api_endpoint: self._sqladmin_api_endpoint = ( _DEFAULT_SCHEME + _SQLADMIN_HOST_TEMPLATE.format(universe_domain=self.universe_domain) ) # otherwise if endpoint override is passed in use it else: self._sqladmin_api_endpoint = sqladmin_api_endpoint # validate that the universe domain of the credentials matches the # universe domain of the service endpoint if self._credentials.universe_domain != self.universe_domain: raise ValueError( f"The configured universe domain ({self.universe_domain}) does " "not match the universe domain found in the credentials " f"({self._credentials.universe_domain}). If you haven't " "configured the universe domain explicitly, `googleapis.com` " "is the default." ) @property def universe_domain(self) -> str: return self._universe_domain or _DEFAULT_UNIVERSE_DOMAIN def connect( self, instance_connection_string: str, driver: str, **kwargs: Any ) -> Any: """Connect to a Cloud SQL instance. Prepares and returns a database connection object connected to a Cloud SQL instance using SSL/TLS. Starts a background refresh to periodically retrieve up-to-date ephemeral certificate and instance metadata. Args: instance_connection_string (str): The instance connection name of the Cloud SQL instance to connect to. Takes the form of "project-id:region:instance-name" Example: "my-project:us-central1:my-instance" driver (str): A string representing the database driver to connect with. Supported drivers are pymysql, pg8000, and pytds. **kwargs: Any driver-specific arguments to pass to the underlying driver .connect call. Returns: A DB-API connection to the specified Cloud SQL instance. """ # connect runs sync database connections on background thread. # Async database connections should call 'connect_async' directly to # avoid hanging indefinitely. connect_future = asyncio.run_coroutine_threadsafe( self.connect_async(instance_connection_string, driver, **kwargs), self._loop, ) return connect_future.result() async def connect_async( self, instance_connection_string: str, driver: str, **kwargs: Any ) -> Any: """Connect asynchronously to a Cloud SQL instance. Prepares and returns a database connection object connected to a Cloud SQL instance using SSL/TLS. Schedules a refresh to periodically retrieve up-to-date ephemeral certificate and instance metadata. Async version of Connector.connect. Args: instance_connection_string (str): The instance connection name of the Cloud SQL instance to connect to. Takes the form of "project-id:region:instance-name" Example: "my-project:us-central1:my-instance" driver (str): A string representing the database driver to connect with. Supported drivers are pymysql, asyncpg, pg8000, and pytds. **kwargs: Any driver-specific arguments to pass to the underlying driver .connect call. Returns: A DB-API connection to the specified Cloud SQL instance. Raises: ValueError: Connection attempt with built-in database authentication and then subsequent attempt with IAM database authentication. KeyError: Unsupported database driver Must be one of pymysql, asyncpg, pg8000, and pytds. """ if self._keys is None: self._keys = asyncio.create_task(generate_keys()) if self._client is None: # lazy init client as it has to be initialized in async context self._client = CloudSQLClient( self._sqladmin_api_endpoint, self._quota_project, self._credentials, user_agent=self._user_agent, driver=driver, ) enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth) 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: Union[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 connect_func = { "pymysql": pymysql.connect, "pg8000": pg8000.connect, "asyncpg": asyncpg.connect, "pytds": pytds.connect, } # only accept supported database drivers try: 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 # want the user to specify them. kwargs.pop("host", None) kwargs.pop("ssl", None) kwargs.pop("port", None) # attempt to get connection info for Cloud SQL instance try: conn_info = await monitored_cache.connect_info() # validate driver matches intended database engine DriverMapping.validate_engine(driver, conn_info.database_version) ip_address = conn_info.get_preferred_ip(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 logger.debug(f"['{conn_info.conn_name}']: Connecting to {ip_address}:3307") # 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}" ) kwargs["user"] = formatted_user try: # async drivers are unblocking and can be awaited directly if driver in ASYNC_DRIVERS: return await connector( ip_address, await conn_info.create_ssl_context(enable_iam_auth), **kwargs, ) # Create socket with SSLContext for sync drivers ctx = await conn_info.create_ssl_context(enable_iam_auth) sock = ctx.wrap_socket( socket.create_connection((ip_address, SERVER_PROXY_PORT)), server_hostname=ip_address, ) # 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, ip_address, sock, **kwargs, ) return await self._loop.run_in_executor(None, connect_partial) except Exception: # with any exception, we attempt a force refresh, then throw the error await monitored_cache.force_refresh() raise async def _remove_cached( self, instance_connection_string: str, enable_iam_auth: bool ) -> None: """Stops all background refreshes and deletes the connection info cache from the map of caches. """ logger.debug( f"['{instance_connection_string}']: Removing connection info from cache" ) # remove cache from stored caches and close it cache = self._cache.pop((instance_connection_string, enable_iam_auth)) await cache.close() def __enter__(self) -> Any: """Enter context manager by returning Connector object""" return self def __exit__( self, exc_type: Optional[type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: """Exit context manager by closing Connector""" self.close() async def __aenter__(self) -> Any: """Enter async context manager by returning Connector object""" return self async def __aexit__( self, exc_type: Optional[type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: """Exit async context manager by closing Connector""" await self.close_async() def close(self) -> None: """Close Connector by stopping tasks and releasing resources.""" if self._loop.is_running(): 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) # if background thread exists for Connector, clean it up if self._thread: if self._loop.is_running(): # stop event loop running in background thread self._loop.call_soon_threadsafe(self._loop.stop) # wait for thread to finish closing (i.e. loop to stop) self._thread.join() async def close_async(self) -> None: """Helper function to cancel the cache's tasks and close aiohttp.ClientSession.""" await asyncio.gather(*[cache.close() for cache in self._cache.values()]) if self._client: await self._client.close() async def create_async_connector( ip_type: str | IPTypes = IPTypes.PUBLIC, enable_iam_auth: bool = False, timeout: int = 30, credentials: Optional[Credentials] = None, loop: Optional[asyncio.AbstractEventLoop] = None, quota_project: Optional[str] = None, sqladmin_api_endpoint: Optional[str] = None, user_agent: Optional[str] = None, universe_domain: Optional[str] = None, refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND, resolver: type[DefaultResolver] | type[DnsResolver] = DefaultResolver, failover_period: int = 30, ) -> Connector: """Helper function to create Connector object for asyncio connections. Force use of Connector in an asyncio context. Auto-detect and use current thread's running event loop. 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 enable_iam_auth (bool): Enables automatic IAM database authentication (Postgres and MySQL) as the default authentication method for all connections. timeout (int): The default time limit in seconds for a connection before raising a TimeoutError. credentials (google.auth.credentials.Credentials): A credentials object created from the google-auth Python library to be used. If not specified, Application Default Credentials (ADC) are used. quota_project (str): The Project ID for an existing Google Cloud project. The project specified is used for quota and billing purposes. If not specified, defaults to project sourced from environment. loop (asyncio.AbstractEventLoop): Event loop to run asyncio tasks, if not specified, defaults to creating new event loop on background thread. sqladmin_api_endpoint (str): Base URL to use when calling the Cloud SQL Admin API endpoint. Defaults to "https://sqladmin.googleapis.com", this argument should only be used in development. universe_domain (str): The universe domain for Cloud SQL API calls. Default: "googleapis.com". refresh_strategy (str | RefreshStrategy): The default refresh strategy used to refresh SSL/TLS cert and instance metadata. Can be one of the following: RefreshStrategy.LAZY ("LAZY") or RefreshStrategy.BACKGROUND ("BACKGROUND"). Default: RefreshStrategy.BACKGROUND resolver (DefaultResolver | DnsResolver): The class name of the resolver to use for resolving the Cloud SQL instance connection name. To resolve a DNS record to an instance connection name, use DnsResolver. Default: DefaultResolver failover_period (int): The time interval in seconds between each attempt to check if a failover has occured for a given instance. Must be used with `resolver=DnsResolver` to have any effect. Default: 30 Returns: A Connector instance configured with running event loop. """ # if no loop given, automatically detect running event loop if loop is None: loop = asyncio.get_running_loop() return Connector( ip_type=ip_type, enable_iam_auth=enable_iam_auth, timeout=timeout, credentials=credentials, loop=loop, quota_project=quota_project, sqladmin_api_endpoint=sqladmin_api_endpoint, user_agent=user_agent, universe_domain=universe_domain, refresh_strategy=refresh_strategy, resolver=resolver, failover_period=failover_period, )