cloud-sql-python-connector 1.17.0.dev0__py3-none-any.whl

google/cloud/sql/connector/instance.py

@@ -0,0 +1,214 @@
+"""
+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 datetime import datetime
+from datetime import timedelta
+from datetime import timezone
+import logging
+
+from google.cloud.sql.connector.client import CloudSQLClient
+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 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
+
+logger = logging.getLogger(name=__name__)
+
+APPLICATION_NAME = "cloud-sql-python-connector"
+
+
+class RefreshAheadCache:
+    """Cache that refreshes connection info in the background prior to expiration.
+
+    Background tasks are used to schedule refresh attempts to get a new
+    ephemeral certificate and Cloud SQL metadata (IP addresses, etc.) ahead of
+    expiration.
+    """
+
+    def __init__(
+        self,
+        conn_name: ConnectionName,
+        client: CloudSQLClient,
+        keys: asyncio.Future,
+        enable_iam_auth: bool = False,
+    ) -> None:
+        """Initializes a RefreshAheadCache instance.
+
+        Args:
+            conn_name (ConnectionName): The Cloud SQL instance's
+                connection name.
+            client (CloudSQLClient): The Cloud SQL Client instance.
+            keys (asyncio.Future): A future to the client's public-private key
+                pair.
+            enable_iam_auth (bool): Enables automatic IAM database authentication
+                (Postgres and MySQL) as the default authentication method for all
+                connections.
+        """
+        self._conn_name = conn_name
+
+        self._enable_iam_auth = enable_iam_auth
+        self._keys = keys
+        self._client = client
+        self._refresh_rate_limiter = AsyncRateLimiter(
+            max_capacity=2,
+            rate=1 / 30,
+        )
+        self._refresh_in_progress = asyncio.locks.Event()
+        self._current: asyncio.Task = self._schedule_refresh(0)
+        self._next: asyncio.Task = self._current
+
+    async def force_refresh(self) -> None:
+        """
+        Forces a new refresh attempt immediately to be used for future connection attempts.
+        """
+        # if next refresh is not already in progress, cancel it and schedule new one immediately
+        if not self._refresh_in_progress.is_set():
+            self._next.cancel()
+            self._next = self._schedule_refresh(0)
+        # block all sequential connection attempts on the next refresh result if current is invalid
+        if not await _is_valid(self._current):
+            self._current = self._next
+
+    async def _perform_refresh(self) -> ConnectionInfo:
+        """Retrieves instance metadata and ephemeral certificate from the
+        Cloud SQL Instance.
+
+        Returns:
+            A ConnectionInfo instance containing a string representing the
+            ephemeral certificate, a dict containing the instances IP adresses,
+            a string representing a PEM-encoded private key and a string
+            representing a PEM-encoded certificate authority.
+        """
+        self._refresh_in_progress.set()
+        logger.debug(
+            f"['{self._conn_name}']: Connection info refresh operation started"
+        )
+
+        try:
+            await self._refresh_rate_limiter.acquire()
+            connection_info = await self._client.get_connection_info(
+                self._conn_name,
+                self._keys,
+                self._enable_iam_auth,
+            )
+            logger.debug(
+                f"['{self._conn_name}']: Connection info refresh operation complete"
+            )
+            logger.debug(
+                f"['{self._conn_name}']: Current certificate "
+                f"expiration = {connection_info.expiration.isoformat()}"
+            )
+
+        except Exception as e:
+            logger.debug(
+                f"['{self._conn_name}']: Connection info "
+                f"refresh operation failed: {str(e)}"
+            )
+            raise
+
+        finally:
+            self._refresh_in_progress.clear()
+        return connection_info
+
+    def _schedule_refresh(self, delay: int) -> asyncio.Task:
+        """
+        Schedule task to sleep and then perform refresh to get ConnectionInfo.
+
+        Args:
+            delay (int): Time in seconds to sleep before performing a refresh.
+
+        Returns:
+            An asyncio.Task representing the scheduled refresh.
+        """
+
+        async def _refresh_task(self: RefreshAheadCache, delay: int) -> ConnectionInfo:
+            """
+            A coroutine that sleeps for the specified amount of time before
+            running _perform_refresh.
+            """
+            refresh_task: asyncio.Task
+            try:
+                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._conn_name}']: Invalid refresh operation. Certficate appears to be expired."
+                    )
+            except asyncio.CancelledError:
+                logger.debug(
+                    f"['{self._conn_name}']: Scheduled refresh" " operation cancelled"
+                )
+                raise
+            # bad refresh attempt
+            except Exception as e:
+                logger.exception(
+                    f"['{self._conn_name}']: "
+                    "An error occurred while performing refresh. "
+                    "Scheduling another refresh attempt immediately",
+                    exc_info=e,
+                )
+                # check if current metadata is invalid (expired),
+                # don't want to replace valid metadata with invalid refresh
+                if not await _is_valid(self._current):
+                    self._current = refresh_task
+                # schedule new refresh attempt immediately
+                self._next = self._schedule_refresh(0)
+                raise
+            # if valid refresh, replace current with valid metadata and schedule next refresh
+            self._current = refresh_task
+            # calculate refresh delay based on certificate expiration
+            delay = _seconds_until_refresh(refresh_data.expiration)
+            logger.debug(
+                f"['{self._conn_name}']: 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
+
+        # schedule refresh task and return it
+        scheduled_task = asyncio.create_task(_refresh_task(self, delay))
+        return scheduled_task
+
+    async def connect_info(self) -> ConnectionInfo:
+        """Retrieves ConnectionInfo instance for establishing a secure
+        connection to the Cloud SQL instance.
+        """
+        return await self._current
+
+    async def close(self) -> None:
+        """Cleanup function to make sure tasks have finished to have a
+        graceful exit.
+        """
+        logger.debug(
+            f"['{self._conn_name}']: Canceling connection info "
+            "refresh operation tasks"
+        )
+        self._current.cancel()
+        self._next.cancel()
+        # gracefully wait for tasks to cancel
+        tasks = asyncio.gather(self._current, self._next, return_exceptions=True)
+        await asyncio.wait_for(tasks, timeout=2.0)

google/cloud/sql/connector/lazy.py

@@ -0,0 +1,124 @@
+# 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.
+
+import asyncio
+from datetime import datetime
+from datetime import timedelta
+from datetime import timezone
+import logging
+from typing import Optional
+
+from google.cloud.sql.connector.client import CloudSQLClient
+from google.cloud.sql.connector.connection_info import ConnectionInfo
+from google.cloud.sql.connector.connection_name import ConnectionName
+from google.cloud.sql.connector.refresh_utils import _refresh_buffer
+
+logger = logging.getLogger(name=__name__)
+
+
+class LazyRefreshCache:
+    """Cache that refreshes connection info when a caller requests a connection.
+
+    Only refreshes the cache when a new connection is requested and the current
+    certificate is close to or already expired.
+
+    This is the recommended option for serverless environments.
+    """
+
+    def __init__(
+        self,
+        conn_name: ConnectionName,
+        client: CloudSQLClient,
+        keys: asyncio.Future,
+        enable_iam_auth: bool = False,
+    ) -> None:
+        """Initializes a LazyRefreshCache instance.
+
+        Args:
+            conn_name (ConnectionName): The Cloud SQL instance's
+                connection name.
+            client (CloudSQLClient): The Cloud SQL Client instance.
+            keys (asyncio.Future): A future to the client's public-private key
+                pair.
+            enable_iam_auth (bool): Enables automatic IAM database authentication
+                (Postgres and MySQL) as the default authentication method for all
+                connections.
+        """
+        self._conn_name = conn_name
+        self._enable_iam_auth = enable_iam_auth
+        self._keys = keys
+        self._client = client
+        self._lock = asyncio.Lock()
+        self._cached: Optional[ConnectionInfo] = None
+        self._needs_refresh = False
+
+    async def force_refresh(self) -> None:
+        """
+        Invalidates the cache and configures the next call to
+        connect_info() to retrieve a fresh ConnectionInfo instance.
+        """
+        async with self._lock:
+            self._needs_refresh = True
+
+    async def connect_info(self) -> ConnectionInfo:
+        """Retrieves ConnectionInfo instance for establishing a secure
+        connection to the Cloud SQL instance.
+        """
+        async with self._lock:
+            # If connection info is cached, check expiration.
+            # Pad expiration with a buffer to give the client plenty of time to
+            # establish a connection to the server with the certificate.
+            if (
+                self._cached
+                and not self._needs_refresh
+                and datetime.now(timezone.utc)
+                < (self._cached.expiration - timedelta(seconds=_refresh_buffer))
+            ):
+                logger.debug(
+                    f"['{self._conn_name}']: Connection info "
+                    "is still valid, using cached info"
+                )
+                return self._cached
+            logger.debug(
+                f"['{self._conn_name}']: Connection info " "refresh operation started"
+            )
+            try:
+                conn_info = await self._client.get_connection_info(
+                    self._conn_name,
+                    self._keys,
+                    self._enable_iam_auth,
+                )
+            except Exception as e:
+                logger.debug(
+                    f"['{self._conn_name}']: Connection info "
+                    f"refresh operation failed: {str(e)}"
+                )
+                raise
+            logger.debug(
+                f"['{self._conn_name}']: Connection info "
+                "refresh operation completed successfully"
+            )
+            logger.debug(
+                f"['{self._conn_name}']: Current certificate "
+                f"expiration = {str(conn_info.expiration)}"
+            )
+            self._cached = conn_info
+            self._needs_refresh = False
+            return conn_info
+
+    async def close(self) -> None:
+        """Close is a no-op and provided purely for a consistent interface with
+        other cache types.
+        """
+        pass

google/cloud/sql/connector/pg8000.py

@@ -0,0 +1,66 @@
+"""
+Copyright 2022 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.
+"""
+
+import socket
+import ssl
+from typing import Any, TYPE_CHECKING
+
+SERVER_PROXY_PORT = 3307
+
+if TYPE_CHECKING:
+    import pg8000
+
+
+def connect(
+    ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
+) -> "pg8000.dbapi.Connection":
+    """Helper function to create a pg8000 DB-API connection object.
+
+    :type ip_address: str
+    :param ip_address: A string containing an IP address for the Cloud SQL
+        instance.
+
+    :type ctx: ssl.SSLContext
+    :param ctx: An SSLContext object created from the Cloud SQL server CA
+        cert and ephemeral cert.
+
+
+    :rtype: pg8000.dbapi.Connection
+    :returns: A pg8000 Connection object for the Cloud SQL instance.
+    """
+    try:
+        import pg8000
+    except ImportError:
+        raise ImportError(
+            'Unable to import module "pg8000." Please install and try again.'
+        )
+
+    # Create socket and wrap with context.
+    sock = ctx.wrap_socket(
+        socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+        server_hostname=ip_address,
+    )
+
+    user = kwargs.pop("user")
+    db = kwargs.pop("db")
+    passwd = kwargs.pop("password", None)
+    return pg8000.dbapi.connect(
+        user,
+        database=db,
+        password=passwd,
+        sock=sock,
+        **kwargs,
+    )

google/cloud/sql/connector/pymysql.py

@@ -0,0 +1,64 @@
+"""
+Copyright 2022 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.
+"""
+
+import socket
+import ssl
+from typing import Any, TYPE_CHECKING
+
+SERVER_PROXY_PORT = 3307
+
+if TYPE_CHECKING:
+    import pymysql
+
+
+def connect(
+    ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
+) -> "pymysql.connections.Connection":
+    """Helper function to create a pymysql DB-API connection object.
+
+    :type ip_address: str
+    :param ip_address: A string containing an IP address for the Cloud SQL
+        instance.
+
+    :type ctx: ssl.SSLContext
+    :param ctx: An SSLContext object created from the Cloud SQL server CA
+        cert and ephemeral cert.
+
+    :rtype: pymysql.Connection
+    :returns: A PyMySQL Connection object for the Cloud SQL instance.
+    """
+    try:
+        import pymysql
+    except ImportError:
+        raise ImportError(
+            'Unable to import module "pymysql." Please install and try again.'
+        )
+
+    # allow automatic IAM database authentication to not require password
+    kwargs["password"] = kwargs["password"] if "password" in kwargs else None
+
+    # Create socket and wrap with context.
+    sock = ctx.wrap_socket(
+        socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+        server_hostname=ip_address,
+    )
+    # pop timeout as timeout arg is called 'connect_timeout' for pymysql
+    timeout = kwargs.pop("timeout")
+    kwargs["connect_timeout"] = kwargs.get("connect_timeout", timeout)
+    # Create pymysql connection object and hand in pre-made connection
+    conn = pymysql.Connection(host=ip_address, defer_connect=True, **kwargs)
+    conn.connect(sock)
+    return conn

google/cloud/sql/connector/pytds.py

@@ -0,0 +1,78 @@
+"""
+Copyright 2022 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.
+"""
+
+import platform
+import socket
+import ssl
+from typing import Any, TYPE_CHECKING
+
+from google.cloud.sql.connector.exceptions import PlatformNotSupportedError
+
+SERVER_PROXY_PORT = 3307
+
+if TYPE_CHECKING:
+    import pytds
+
+
+def connect(ip_address: str, ctx: ssl.SSLContext, **kwargs: Any) -> "pytds.Connection":
+    """Helper function to create a pytds DB-API connection object.
+
+    :type ip_address: str
+    :param ip_address: A string containing an IP address for the Cloud SQL
+        instance.
+
+    :type ctx: ssl.SSLContext
+    :param ctx: An SSLContext object created from the Cloud SQL server CA
+        cert and ephemeral cert.
+
+
+    :rtype: pytds.Connection
+    :returns: A pytds Connection object for the Cloud SQL instance.
+    """
+    try:
+        import pytds
+    except ImportError:
+        raise ImportError(
+            'Unable to import module "pytds." Please install and try again.'
+        )
+
+    db = kwargs.pop("db", None)
+
+    # Create socket and wrap with context.
+    sock = ctx.wrap_socket(
+        socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+        server_hostname=ip_address,
+    )
+    if kwargs.pop("active_directory_auth", False):
+        if platform.system() == "Windows":
+            # Ignore username and password if using active directory auth
+            server_name = kwargs.pop("server_name")
+            return pytds.connect(
+                database=db,
+                auth=pytds.login.SspiAuth(port=1433, server_name=server_name),
+                sock=sock,
+                **kwargs,
+            )
+        else:
+            raise PlatformNotSupportedError(
+                "Active Directory authentication is currently only supported on Windows."
+            )
+
+    user = kwargs.pop("user")
+    passwd = kwargs.pop("password")
+    return pytds.connect(
+        ip_address, database=db, user=user, password=passwd, sock=sock, **kwargs
+    )

google/cloud/sql/connector/rate_limiter.py

@@ -0,0 +1,87 @@
+"""
+Copyright 2021 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.
+"""
+
+import asyncio
+from typing import Optional
+
+
+class AsyncRateLimiter(object):
+    """
+    An asyncio-compatible rate limiter which uses the Token Bucket algorithm
+    (https://en.wikipedia.org/wiki/Token_bucket) to limit the number of function calls over a time interval using an event queue.
+
+    :type max_capacity: int
+    :param: max_capacity:
+        The maximum capacity of tokens the bucket will store at any one time.
+        Default: 1
+
+    :type rate: float
+    :param: rate:
+        The number of tokens that should be added per second.
+
+    :type loop: asyncio.AbstractEventLoop
+    :param: loop:
+        The event loop to use. If not provided, the default event loop will be used.
+
+
+    """
+
+    def __init__(
+        self,
+        max_capacity: int = 1,
+        rate: float = 1 / 60,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+    ) -> None:
+        self.rate = rate
+        self.max_capacity = max_capacity
+        self._loop = loop or asyncio.get_event_loop()
+        self._tokens: float = max_capacity
+        self._last_token_update = self._loop.time()
+        self._lock = asyncio.Lock()
+
+    def _update_token_count(self) -> None:
+        """
+        Calculates how much time has passed since the last leak and removes the
+        appropriate amount of events from the queue.
+        Leaking is done lazily, meaning that if there is a large time gap between
+        leaks, the next set of calls might be a burst if burst_size > 1
+        """
+        now = self._loop.time()
+        time_elapsed = now - self._last_token_update
+        new_tokens = time_elapsed * self.rate
+        self._tokens = min(new_tokens + self._tokens, self.max_capacity)
+        self._last_token_update = now
+
+    async def _wait_for_next_token(self) -> None:
+        """
+        Wait until enough time has elapsed to add another token.
+        """
+        token_deficit = 1 - self._tokens
+        if token_deficit > 0:
+            wait_time = token_deficit / self.rate
+            await asyncio.sleep(wait_time)
+
+    async def acquire(self) -> None:
+        """
+        Waits for a token to become available, if necessary, then subtracts token and allows
+        request to go through.
+        """
+        async with self._lock:
+            self._update_token_count()
+            if self._tokens < 1:
+                await self._wait_for_next_token()
+                self._update_token_count()
+            self._tokens -= 1