google/cloud/sql/connector/instance.py

"""
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_info import ConnectionInfoCache
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(ConnectionInfoCache):
    """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
        self._closed = False

    @property
    def conn_name(self) -> ConnectionName:
        return self._conn_name

    @property
    def closed(self) -> bool:
        return self._closed

    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)
        self._closed = True