cloud-sql-python-connector 1.9.2__py2.py3-none-any.whl → 1.11.0__py2.py3-none-any.whl

google/cloud/sql/connector/connector.py

@@ -19,10 +19,9 @@ from __future__ import annotations
 import asyncio
 from functools import partial
 import logging
-import socket
 from threading import Thread
 from types import TracebackType
-from typing import Any, Dict, Optional, Type
+from typing import Any, Dict, Optional, Type, Union
 
 import google.auth
 from google.auth.credentials import Credentials
@@ -30,10 +29,12 @@ 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.exceptions import ConnectorLoopError
-from google.cloud.sql.connector.exceptions import DnsNameResolutionError
-from google.cloud.sql.connector.instance import IPTypes
 from google.cloud.sql.connector.instance import RefreshAheadCache
+from google.cloud.sql.connector.lazy import LazyRefreshCache
 import google.cloud.sql.connector.pg8000 as pg8000
 import google.cloud.sql.connector.pymysql as pymysql
 import google.cloud.sql.connector.pytds as pytds
@@ -62,6 +63,7 @@ class Connector:
         sqladmin_api_endpoint: Optional[str] = None,
         user_agent: Optional[str] = None,
         universe_domain: Optional[str] = None,
+        refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
     ) -> None:
         """Initializes a Connector instance.
 
@@ -98,22 +100,39 @@ class Connector:
             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
         """
+        # 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
-            self._keys: asyncio.Future = loop.create_task(generate_keys())
+            # 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()
-            self._keys = asyncio.wrap_future(
-                asyncio.run_coroutine_threadsafe(generate_keys(), self._loop),
-                loop=self._loop,
-            )
-        self._cache: Dict[str, RefreshAheadCache] = {}
+            # 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,
+                )
+        self._cache: Dict[str, Union[RefreshAheadCache, LazyRefreshCache]] = {}
         self._client: Optional[CloudSQLClient] = None
 
         # initialize credentials
@@ -242,9 +261,9 @@ class Connector:
                 and then subsequent attempt with IAM database authentication.
             KeyError: Unsupported database driver Must be one of pymysql, asyncpg,
                 pg8000, and pytds.
-            DnsNameResolutionError: Could not resolve PSC IP address from DNS
-                host name.
         """
+        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(
@@ -265,11 +284,30 @@ class Connector:
                     "connector.Connector object."
                 )
         else:
-            cache = RefreshAheadCache(
-                instance_connection_string,
-                self._client,
-                self._keys,
-                enable_iam_auth,
+            if self._refresh_strategy == RefreshStrategy.LAZY:
+                logger.debug(
+                    f"['{instance_connection_string}']: Refresh strategy is set"
+                    " to lazy refresh"
+                )
+                cache = LazyRefreshCache(
+                    instance_connection_string,
+                    self._client,
+                    self._keys,
+                    enable_iam_auth,
+                )
+            else:
+                logger.debug(
+                    f"['{instance_connection_string}']: Refresh strategy is set"
+                    " to backgound refresh"
+                )
+                cache = RefreshAheadCache(
+                    instance_connection_string,
+                    self._client,
+                    self._keys,
+                    enable_iam_auth,
+                )
+            logger.debug(
+                f"['{instance_connection_string}']: Connection info added to cache"
             )
             self._cache[instance_connection_string] = cache
 
@@ -298,42 +336,44 @@ class Connector:
         kwargs.pop("ssl", None)
         kwargs.pop("port", None)
 
-        # attempt to make connection to Cloud SQL instance
+        # attempt to get connection info for Cloud SQL instance
         try:
-            instance_data, ip_address = await cache.connect_info(ip_type)
-            # resolve DNS name into IP address for PSC
-            if ip_type.value == "PSC":
-                addr_info = await self._loop.getaddrinfo(
-                    ip_address, None, family=socket.AF_INET, type=socket.SOCK_STREAM
-                )
-                # getaddrinfo returns a list of 5-tuples that contain socket
-                # connection info in the form
-                # (family, type, proto, canonname, sockaddr), where sockaddr is a
-                # 2-tuple in the form (ip_address, port)
-                try:
-                    ip_address = addr_info[0][4][0]
-                except IndexError as e:
-                    raise DnsNameResolutionError(
-                        f"['{instance_connection_string}']: DNS name could not be resolved into IP address"
-                    ) from e
-
-            # format `user` param for automatic IAM database authn
-            if enable_iam_auth:
-                formatted_user = format_database_user(
-                    instance_data.database_version, kwargs["user"]
+            conn_info = await 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(instance_connection_string)
+            raise
+        logger.debug(
+            f"['{instance_connection_string}']: 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}"
                 )
-                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
-
+                kwargs["user"] = formatted_user
+        try:
             # async drivers are unblocking and can be awaited directly
             if driver in ASYNC_DRIVERS:
-                return await connector(ip_address, instance_data.context, **kwargs)
+                return await connector(
+                    ip_address,
+                    conn_info.create_ssl_context(enable_iam_auth),
+                    **kwargs,
+                )
             # synchronous drivers are blocking and run using executor
             connect_partial = partial(
-                connector, ip_address, instance_data.context, **kwargs
+                connector,
+                ip_address,
+                conn_info.create_ssl_context(enable_iam_auth),
+                **kwargs,
             )
             return await self._loop.run_in_executor(None, connect_partial)
 
@@ -342,6 +382,17 @@ class Connector:
             await cache.force_refresh()
             raise
 
+    async def _remove_cached(self, instance_connection_string: str) -> 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)
+        await cache.close()
+
     def __enter__(self) -> Any:
         """Enter context manager by returning Connector object"""
         return self

google/cloud/sql/connector/enums.py

@@ -0,0 +1,85 @@
+# Copyright 2024 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
+
+from enum import Enum
+
+from google.cloud.sql.connector.exceptions import IncompatibleDriverError
+
+
+class RefreshStrategy(Enum):
+    LAZY: str = "LAZY"
+    BACKGROUND: str = "BACKGROUND"
+
+    @classmethod
+    def _missing_(cls, value: object) -> None:
+        raise ValueError(
+            f"Incorrect value for refresh_strategy, got '{value}'. Want one of: "
+            f"{', '.join([repr(m.value) for m in cls])}."
+        )
+
+    @classmethod
+    def _from_str(cls, refresh_strategy: str) -> RefreshStrategy:
+        """Convert refresh strategy from a str into RefreshStrategy."""
+        return cls(refresh_strategy.upper())
+
+
+class IPTypes(Enum):
+    PUBLIC: str = "PRIMARY"
+    PRIVATE: str = "PRIVATE"
+    PSC: str = "PSC"
+
+    @classmethod
+    def _missing_(cls, value: object) -> None:
+        raise ValueError(
+            f"Incorrect value for ip_type, got '{value}'. Want one of: "
+            f"{', '.join([repr(m.value) for m in cls])}, 'PUBLIC'."
+        )
+
+    @classmethod
+    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"
+        return cls(ip_type_str.upper())
+
+
+class DriverMapping(Enum):
+    """Maps a given database driver to it's corresponding database engine."""
+
+    ASYNCPG = "POSTGRES"
+    PG8000 = "POSTGRES"
+    PYMYSQL = "MYSQL"
+    PYTDS = "SQLSERVER"
+
+    @staticmethod
+    def validate_engine(driver: str, engine_version: str) -> None:
+        """Validate that the given driver is compatible with the given engine.
+
+        Args:
+            driver (str): Database driver being used. (i.e. "pg8000")
+            engine_version (str): Database engine version. (i.e. "POSTGRES_16")
+
+        Raises:
+            IncompatibleDriverError: If the given driver is not compatible with
+                the given engine.
+        """
+        mapping = DriverMapping[driver.upper()]
+        if not engine_version.startswith(mapping.value):
+            raise IncompatibleDriverError(
+                f"Database driver '{driver}' is incompatible with database "
+                f"version '{engine_version}'. Given driver can "
+                f"only be used with Cloud SQL {mapping.value} databases."
+            )

google/cloud/sql/connector/exceptions.py

@@ -57,8 +57,16 @@ class AutoIAMAuthNotSupported(Exception):
     pass
 
 
-class DnsNameResolutionError(Exception):
+class RefreshNotValidError(Exception):
     """
-    Exception to be raised when the DnsName of a PSC connection to a
-    Cloud SQL instance can not be resolved to a proper IP address.
+    Exception to be raised when the task returned from refresh is not valid.
+    """
+
+    pass
+
+
+class IncompatibleDriverError(Exception):
+    """
+    Exception to be raised when the database driver given is for the wrong
+    database engine. (i.e. asyncpg for a MySQL database)
     """

google/cloud/sql/connector/instance.py

@@ -17,26 +17,21 @@ limitations under the License.
 from __future__ import annotations
 
 import asyncio
-from enum import Enum
+from datetime import datetime
+from datetime import timedelta
+from datetime import timezone
 import logging
 import re
-import ssl
-from tempfile import TemporaryDirectory
-from typing import Any, Dict, Tuple, TYPE_CHECKING
+from typing import Tuple
 
 import aiohttp
 
 from google.cloud.sql.connector.client import CloudSQLClient
-from google.cloud.sql.connector.exceptions import AutoIAMAuthNotSupported
-from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
-from google.cloud.sql.connector.exceptions import TLSVersionError
+from google.cloud.sql.connector.connection_info import ConnectionInfo
+from google.cloud.sql.connector.exceptions import RefreshNotValidError
 from google.cloud.sql.connector.rate_limiter import AsyncRateLimiter
 from google.cloud.sql.connector.refresh_utils import _is_valid
 from google.cloud.sql.connector.refresh_utils import _seconds_until_refresh
-from google.cloud.sql.connector.utils import write_to_file
-
-if TYPE_CHECKING:
-    import datetime
 
 logger = logging.getLogger(name=__name__)
 
@@ -59,91 +54,6 @@ def _parse_instance_connection_name(connection_name: str) -> Tuple[str, str, str
     return connection_name_split[1], connection_name_split[3], connection_name_split[4]
 
 
-class IPTypes(Enum):
-    PUBLIC: str = "PRIMARY"
-    PRIVATE: str = "PRIVATE"
-    PSC: str = "PSC"
-
-    @classmethod
-    def _missing_(cls, value: object) -> None:
-        raise ValueError(
-            f"Incorrect value for ip_type, got '{value}'. Want one of: "
-            f"{', '.join([repr(m.value) for m in cls])}, 'PUBLIC'."
-        )
-
-    @classmethod
-    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"
-        return cls(ip_type_str.upper())
-
-
-class ConnectionInfo:
-    ip_addrs: Dict[str, Any]
-    context: ssl.SSLContext
-    database_version: str
-    expiration: datetime.datetime
-
-    def __init__(
-        self,
-        ephemeral_cert: str,
-        database_version: str,
-        ip_addrs: Dict[str, Any],
-        private_key: bytes,
-        server_ca_cert: str,
-        expiration: datetime.datetime,
-        enable_iam_auth: bool,
-    ) -> None:
-        self.ip_addrs = ip_addrs
-        self.database_version = database_version
-        self.context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
-
-        # update ssl.PROTOCOL_TLS_CLIENT default
-        self.context.check_hostname = False
-
-        # verify OpenSSL version supports TLSv1.3
-        if ssl.HAS_TLSv1_3:
-            # force TLSv1.3 if supported by client
-            self.context.minimum_version = ssl.TLSVersion.TLSv1_3
-        # fallback to TLSv1.2 for older versions of OpenSSL
-        else:
-            if enable_iam_auth:
-                raise TLSVersionError(
-                    f"Your current version of OpenSSL ({ssl.OPENSSL_VERSION}) does not "
-                    "support TLSv1.3, which is required to use IAM Authentication.\n"
-                    "Upgrade your OpenSSL version to 1.1.1 for TLSv1.3 support."
-                )
-            logger.warning(
-                "TLSv1.3 is not supported with your version of OpenSSL "
-                f"({ssl.OPENSSL_VERSION}), falling back to TLSv1.2\n"
-                "Upgrade your OpenSSL version to 1.1.1 for TLSv1.3 support."
-            )
-            self.context.minimum_version = ssl.TLSVersion.TLSv1_2
-        self.expiration = expiration
-
-        # tmpdir and its contents are automatically deleted after the CA cert
-        # and ephemeral cert are loaded into the SSLcontext. The values
-        # need to be written to files in order to be loaded by the SSLContext
-        with TemporaryDirectory() as tmpdir:
-            ca_filename, cert_filename, key_filename = write_to_file(
-                tmpdir, server_ca_cert, ephemeral_cert, private_key
-            )
-            self.context.load_cert_chain(cert_filename, keyfile=key_filename)
-            self.context.load_verify_locations(cafile=ca_filename)
-
-    def get_preferred_ip(self, ip_type: IPTypes) -> str:
-        """Returns the first IP address for the instance, according to the preference
-        supplied by ip_type. If no IP addressess with the given preference are found,
-        an error is raised."""
-        if ip_type.value in self.ip_addrs:
-            return self.ip_addrs[ip_type.value]
-        raise CloudSQLIPTypeError(
-            "Cloud SQL instance does not have any IP addresses matching "
-            f"preference: {ip_type.value})"
-        )
-
-
 class RefreshAheadCache:
     """Cache that refreshes connection info in the background prior to expiration.
 
@@ -212,73 +122,47 @@ class RefreshAheadCache:
         """
         self._refresh_in_progress.set()
         logger.debug(
-            f"['{self._instance_connection_string}']: Entered _perform_refresh"
+            f"['{self._instance_connection_string}']: Connection info refresh "
+            "operation started"
         )
 
         try:
             await self._refresh_rate_limiter.acquire()
-            priv_key, pub_key = await self._keys
-
-            logger.debug(f"['{self._instance_connection_string}']: Creating context")
-
-            metadata_task = asyncio.create_task(
-                self._client._get_metadata(
-                    self._project,
-                    self._region,
-                    self._instance,
-                )
+            connection_info = await self._client.get_connection_info(
+                self._project,
+                self._region,
+                self._instance,
+                self._keys,
+                self._enable_iam_auth,
             )
-
-            ephemeral_task = asyncio.create_task(
-                self._client._get_ephemeral(
-                    self._project,
-                    self._instance,
-                    pub_key,
-                    self._enable_iam_auth,
-                )
+            logger.debug(
+                f"['{self._instance_connection_string}']: Connection info "
+                "refresh operation complete"
+            )
+            logger.debug(
+                f"['{self._instance_connection_string}']: Current certificate "
+                f"expiration = {connection_info.expiration.isoformat()}"
             )
-            try:
-                metadata = await metadata_task
-                # check if automatic IAM database authn is supported for database engine
-                if self._enable_iam_auth and not metadata[
-                    "database_version"
-                ].startswith(("POSTGRES", "MYSQL")):
-                    raise AutoIAMAuthNotSupported(
-                        f"'{metadata['database_version']}' does not support automatic IAM authentication. It is only supported with Cloud SQL Postgres or MySQL instances."
-                    )
-            except Exception:
-                # cancel ephemeral cert task if exception occurs before it is awaited
-                ephemeral_task.cancel()
-                raise
-
-            ephemeral_cert, expiration = await ephemeral_task
 
         except aiohttp.ClientResponseError as e:
             logger.debug(
-                f"['{self._instance_connection_string}']: Error occurred during _perform_refresh."
+                f"['{self._instance_connection_string}']: Connection info "
+                f"refresh operation failed: {str(e)}"
             )
             if e.status == 403:
-                e.message = "Forbidden: Authenticated IAM principal does not seeem authorized to make API request. Verify 'Cloud SQL Admin API' is enabled within your GCP project and 'Cloud SQL Client' role has been granted to IAM principal."
+                e.message = "Forbidden: Authenticated IAM principal does not seem authorized to make API request. Verify 'Cloud SQL Admin API' is enabled within your GCP project and 'Cloud SQL Client' role has been granted to IAM principal."
             raise
 
-        except Exception:
+        except Exception as e:
             logger.debug(
-                f"['{self._instance_connection_string}']: Error occurred during _perform_refresh."
+                f"['{self._instance_connection_string}']: Connection info "
+                f"refresh operation failed: {str(e)}"
             )
             raise
 
         finally:
             self._refresh_in_progress.clear()
-
-        return ConnectionInfo(
-            ephemeral_cert,
-            metadata["database_version"],
-            metadata["ip_addresses"],
-            priv_key,
-            metadata["server_ca_cert"],
-            expiration,
-            self._enable_iam_auth,
-        )
+        return connection_info
 
     def _schedule_refresh(self, delay: int) -> asyncio.Task:
         """
@@ -298,14 +182,19 @@ class RefreshAheadCache:
             """
             refresh_task: asyncio.Task
             try:
-                logger.debug(f"['{self._instance_connection_string}']: Entering sleep")
                 if delay > 0:
                     await asyncio.sleep(delay)
                 refresh_task = asyncio.create_task(self._perform_refresh())
                 refresh_data = await refresh_task
+                # check that refresh is valid
+                if not await _is_valid(refresh_task):
+                    raise RefreshNotValidError(
+                        f"['{self._instance_connection_string}']: Invalid refresh operation. Certficate appears to be expired."
+                    )
             except asyncio.CancelledError:
                 logger.debug(
-                    f"['{self._instance_connection_string}']: Schedule refresh task cancelled."
+                    f"['{self._instance_connection_string}']: Scheduled refresh"
+                    " operation cancelled"
                 )
                 raise
             # bad refresh attempt
@@ -327,6 +216,12 @@ class RefreshAheadCache:
             self._current = refresh_task
             # calculate refresh delay based on certificate expiration
             delay = _seconds_until_refresh(refresh_data.expiration)
+            logger.debug(
+                f"['{self._instance_connection_string}']: Connection info refresh"
+                " operation scheduled for "
+                f"{(datetime.now(timezone.utc) + timedelta(seconds=delay)).isoformat(timespec='seconds')} "
+                f"(now + {timedelta(seconds=delay)})"
+            )
             self._next = self._schedule_refresh(delay)
 
             return refresh_data
@@ -335,47 +230,22 @@ class RefreshAheadCache:
         scheduled_task = asyncio.create_task(_refresh_task(self, delay))
         return scheduled_task
 
-    async def connect_info(
-        self,
-        ip_type: IPTypes,
-    ) -> Tuple[ConnectionInfo, str]:
-        """Retrieve instance metadata and ip address required
-        for making connection to Cloud SQL instance.
-
-        Args:
-            ip_type (IPTypes): Enum specifying type of IP address to lookup and
-                use for connection.
-
-        Returns:
-            A tuple with the first item being the ConnectionInfo instance for
-            establishing the connection, and the second item being the IP
-            address of the Cloud SQL instance matching the specified IP type.
+    async def connect_info(self) -> ConnectionInfo:
+        """Retrieves ConnectionInfo instance for establishing a secure
+        connection to the Cloud SQL instance.
         """
-        logger.debug(
-            f"['{self._instance_connection_string}']: Entered connect_info method"
-        )
-
-        instance_data: ConnectionInfo
-
-        instance_data = await self._current
-        ip_address: str = instance_data.get_preferred_ip(ip_type)
-        return instance_data, ip_address
+        return await self._current
 
     async def close(self) -> None:
-        """Cleanup function to make sure ClientSession is closed and tasks have
-        finished to have a graceful exit.
+        """Cleanup function to make sure tasks have finished to have a
+        graceful exit.
         """
         logger.debug(
-            f"['{self._instance_connection_string}']: Waiting for _current to be cancelled"
+            f"['{self._instance_connection_string}']: Canceling connection info "
+            "refresh operation tasks"
         )
         self._current.cancel()
-        logger.debug(
-            f"['{self._instance_connection_string}']: Waiting for _next to be cancelled"
-        )
         self._next.cancel()
-        logger.debug(
-            f"['{self._instance_connection_string}']: Waiting for _client_session to close"
-        )
         # gracefully wait for tasks to cancel
         tasks = asyncio.gather(self._current, self._next, return_exceptions=True)
         await asyncio.wait_for(tasks, timeout=2.0)