refactor: use one ClientSession per Connector (#1007)

Author: jackwotherspoon

google/cloud/sql/connector/client.py

@@ -0,0 +1,221 @@
+# 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
+
+import datetime
+import logging
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
+
+import aiohttp
+from cryptography.hazmat.backends import default_backend
+from cryptography.x509 import load_pem_x509_certificate
+import google.auth.transport.requests
+
+from google.cloud.sql.connector.refresh_utils import _downscope_credentials
+from google.cloud.sql.connector.version import __version__ as version
+
+if TYPE_CHECKING:
+    from google.auth.credentials import Credentials
+
+USER_AGENT: str = f"cloud-sql-python-connector/{version}"
+API_VERSION: str = "v1beta4"
+
+logger = logging.getLogger(name=__name__)
+
+
+def _format_user_agent(driver: Optional[str], custom: Optional[str]) -> str:
+    agent = f"{USER_AGENT}+{driver}" if driver else USER_AGENT
+    if custom:
+        agent = f"{agent} {custom}"
+    return agent
+
+
+class CloudSQLClient:
+    def __init__(
+        self,
+        sqladmin_api_endpoint: str,
+        quota_project: Optional[str],
+        credentials: Credentials,
+        client: Optional[aiohttp.ClientSession] = None,
+        driver: Optional[str] = None,
+        user_agent: Optional[str] = None,
+    ) -> None:
+        """
+        Establish the client to be used for Cloud SQL Admin API requests.
+
+        Args:
+            sqladmin_api_endpoint (str): Base URL to use when calling
+                the Cloud SQL Admin API endpoints.
+            quota_project (str): The Project ID for an existing Google Cloud
+                project. The project specified is used for quota and
+                billing purposes.
+            credentials (google.auth.credentials.Credentials):
+                A credentials object created from the google-auth Python library.
+                Must have the Cloud SQL Admin scopes. For more info check out
+                https://google-auth.readthedocs.io/en/latest/.
+            client (aiohttp.ClientSession): Async client used to make requests to
+                Cloud SQL Admin APIs.
+                Optional, defaults to None and creates new client.
+            driver (str): Database driver to be used by the client.
+        """
+        user_agent = _format_user_agent(driver, user_agent)
+        headers = {
+            "x-goog-api-client": user_agent,
+            "User-Agent": user_agent,
+            "Content-Type": "application/json",
+        }
+        if quota_project:
+            headers["x-goog-user-project"] = quota_project
+
+        self._client = client if client else aiohttp.ClientSession(headers=headers)
+        self._credentials = credentials
+        self._sqladmin_api_endpoint = sqladmin_api_endpoint
+        self._user_agent = user_agent
+
+    async def _get_metadata(
+        self,
+        project: str,
+        region: str,
+        instance: str,
+    ) -> Dict[str, Any]:
+        """Requests metadata from the Cloud SQL Instance
+        and returns a dictionary containing the IP addresses and certificate
+        authority of the Cloud SQL Instance.
+
+        :type project: str
+        :param project:
+            A string representing the name of the project.
+
+        :type region: str
+        :param region : A string representing the name of the region.
+
+        :type instance: str
+        :param instance: A string representing the name of the instance.
+
+        :rtype: Dict[str: Union[Dict, str]]
+        :returns: Returns a dictionary containing a dictionary of all IP
+            addresses and their type and a string representing the
+            certificate authority.
+        """
+        if not self._credentials.valid:
+            request = google.auth.transport.requests.Request()
+            self._credentials.refresh(request)
+
+        headers = {
+            "Authorization": f"Bearer {self._credentials.token}",
+        }
+
+        url = f"{self._sqladmin_api_endpoint}/sql/{API_VERSION}/projects/{project}/instances/{instance}/connectSettings"
+
+        logger.debug(f"['{instance}']: Requesting metadata")
+
+        resp = await self._client.get(url, headers=headers, raise_for_status=True)
+        ret_dict = await resp.json()
+
+        if ret_dict["region"] != region:
+            raise ValueError(
+                f'[{project}:{region}:{instance}]: Provided region was mismatched - got region {region}, expected {ret_dict["region"]}.'
+            )
+
+        ip_addresses = (
+            {ip["type"]: ip["ipAddress"] for ip in ret_dict["ipAddresses"]}
+            if "ipAddresses" in ret_dict
+            else {}
+        )
+        if "dnsName" in ret_dict:
+            ip_addresses["PSC"] = ret_dict["dnsName"]
+
+        return {
+            "ip_addresses": ip_addresses,
+            "server_ca_cert": ret_dict["serverCaCert"]["cert"],
+            "database_version": ret_dict["databaseVersion"],
+        }
+
+    async def _get_ephemeral(
+        self,
+        project: str,
+        instance: str,
+        pub_key: str,
+        enable_iam_auth: bool = False,
+    ) -> Tuple[str, datetime.datetime]:
+        """Asynchronously requests an ephemeral certificate from the Cloud SQL Instance.
+
+        :type project: str
+        :param project : A string representing the name of the project.
+
+        :type instance: str
+        :param instance: A string representing the name of the instance.
+
+        :type pub_key:
+        :param str: A string representing PEM-encoded RSA public key.
+
+        :type enable_iam_auth: bool
+        :param enable_iam_auth
+            Enables automatic IAM database authentication for Postgres or MySQL
+            instances.
+
+        :rtype: str
+        :returns: An ephemeral certificate from the Cloud SQL instance that allows
+            authorized connections to the instance.
+        """
+
+        logger.debug(f"['{instance}']: Requesting ephemeral certificate")
+
+        if not self._credentials.valid:
+            request = google.auth.transport.requests.Request()
+            self._credentials.refresh(request)
+
+        headers = {
+            "Authorization": f"Bearer {self._credentials.token}",
+        }
+
+        url = f"{self._sqladmin_api_endpoint}/sql/{API_VERSION}/projects/{project}/instances/{instance}:generateEphemeralCert"
+
+        data = {"public_key": pub_key}
+
+        if enable_iam_auth:
+            # down-scope credentials with only IAM login scope (refreshes them too)
+            login_creds = _downscope_credentials(self._credentials)
+            data["access_token"] = login_creds.token
+
+        resp = await self._client.post(
+            url, headers=headers, json=data, raise_for_status=True
+        )
+
+        ret_dict = await resp.json()
+
+        ephemeral_cert: str = ret_dict["ephemeralCert"]["cert"]
+
+        # decode cert to read expiration
+        x509 = load_pem_x509_certificate(
+            ephemeral_cert.encode("UTF-8"), default_backend()
+        )
+        expiration = x509.not_valid_after_utc
+        # for IAM authentication OAuth2 token is embedded in cert so it
+        # must still be valid for successful connection
+        if enable_iam_auth:
+            token_expiration: datetime.datetime = login_creds.expiry
+            # google.auth library strips timezone info for backwards compatibality
+            # reasons with Python 2. Add it back to allow timezone aware datetimes.
+            # Ref: https://github.com/googleapis/google-auth-library-python/blob/49a5ff7411a2ae4d32a7d11700f9f961c55406a9/google/auth/_helpers.py#L93-L99
+            token_expiration = token_expiration.replace(tzinfo=datetime.timezone.utc)
+
+            if expiration > token_expiration:
+                expiration = token_expiration
+        return ephemeral_cert, expiration
+
+    async def close(self) -> None:
+        """Close CloudSQLClient gracefully."""
+        await self._client.close()

google/cloud/sql/connector/connector.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from __future__ import annotations
 
 import asyncio
@@ -28,6 +29,7 @@
 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.exceptions import ConnectorLoopError
 from google.cloud.sql.connector.exceptions import DnsNameResolutionError
 from google.cloud.sql.connector.instance import Instance
@@ -109,11 +111,12 @@ def __init__(
                 loop=self._loop,
             )
         self._instances: Dict[str, Instance] = {}
+        self._client: Optional[CloudSQLClient] = None
 
         # initialize credentials
         scopes = ["https://www.googleapis.com/auth/sqlservice.admin"]
         if credentials:
-            # verfiy custom credentials are proper type
+            # verify custom credentials are proper type
             # and atleast base class of google.auth.credentials
             if not isinstance(credentials, Credentials):
                 raise TypeError(
@@ -207,27 +210,32 @@ async def connect_async(
         # Use the Instance to establish an SSL Connection.
         #
         # Return a DBAPI connection
+        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)
         if instance_connection_string in self._instances:
             instance = self._instances[instance_connection_string]
             if enable_iam_auth != instance._enable_iam_auth:
                 raise ValueError(
-                    f"connect() called with `enable_iam_auth={enable_iam_auth}`, "
-                    f"but previously used enable_iam_auth={instance._enable_iam_auth}`. "
+                    f"connect() called with 'enable_iam_auth={enable_iam_auth}', "
+                    f"but previously used 'enable_iam_auth={instance._enable_iam_auth}'. "
                     "If you require both for your use case, please use a new "
                     "connector.Connector object."
                 )
         else:
             instance = Instance(
                 instance_connection_string,
-                driver,
+                self._client,
                 self._keys,
                 self._loop,
-                self._credentials,
                 enable_iam_auth,
-                self._quota_project,
-                self._sqladmin_api_endpoint,
-                user_agent=self._user_agent,
             )
             self._instances[instance_connection_string] = instance
 
@@ -329,8 +337,8 @@ def close(self) -> None:
             close_future = asyncio.run_coroutine_threadsafe(
                 self.close_async(), loop=self._loop
             )
-            # Will attempt to safely shut down tasks for 5s
-            close_future.result(timeout=5)
+            # 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():
@@ -345,6 +353,8 @@ async def close_async(self) -> None:
         await asyncio.gather(
             *[instance.close() for instance in self._instances.values()]
         )
+        if self._client:
+            await self._client.close()
 
     def __del__(self) -> None:
         """Close Connector as part of garbage collection"""