# Copyright 2023 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 asyncio from datetime import datetime from datetime import timedelta from datetime import timezone import logging import re from typing import TYPE_CHECKING from google.cloud.alloydb.connector.connection_info import ConnectionInfo from google.cloud.alloydb.connector.exceptions import RefreshError from google.cloud.alloydb.connector.rate_limiter import AsyncRateLimiter from google.cloud.alloydb.connector.refresh_utils import _is_valid from google.cloud.alloydb.connector.refresh_utils import _seconds_until_refresh if TYPE_CHECKING: from cryptography.hazmat.primitives.asymmetric import rsa from google.cloud.alloydb.connector.client import AlloyDBClient logger = logging.getLogger(name=__name__) INSTANCE_URI_REGEX = re.compile( "projects/([^:]+(:[^:]+)?)/locations/([^:]+)/clusters/([^:]+)/instances/([^:]+)" ) def _parse_instance_uri(instance_uri: str) -> tuple[str, str, str, str]: # should take form "projects//locations//clusters//instances/" if INSTANCE_URI_REGEX.fullmatch(instance_uri) is None: raise ValueError( "Arg `instance_uri` must have " "format: projects//locations//clusters//instances/, projects/:/locations//clusters//instances/" f"got {instance_uri}." ) instance_uri_split = INSTANCE_URI_REGEX.split(instance_uri) return ( instance_uri_split[1], instance_uri_split[3], instance_uri_split[4], instance_uri_split[5], ) class RefreshAheadCache: """ Manages the information used to connect to the AlloyDB instance. Periodically calls the AlloyDB API, automatically refreshing the required information approximately 4 minutes before the previous certificate expires (every ~56 minutes). Args: instance_uri (str): The instance URI of the AlloyDB instance. ex. projects//locations//clusters//instances/ client (AlloyDBClient): Client used to make requests to AlloyDB APIs. keys (tuple[rsa.RSAPrivateKey, str]): Private and Public key pair. """ def __init__( self, instance_uri: str, client: AlloyDBClient, keys: asyncio.Future[tuple[rsa.RSAPrivateKey, str]], ) -> None: # validate and parse instance_uri self._project, self._region, self._cluster, self._name = _parse_instance_uri( instance_uri ) self._instance_uri = instance_uri self._client = client self._keys = keys self._refresh_rate_limiter = AsyncRateLimiter( max_capacity=2, rate=1 / 30, ) self._refresh_in_progress = asyncio.locks.Event() # For the initial refresh operation, set current = next so that # connection requests block until the first refresh is complete. self._current: asyncio.Task = self._schedule_refresh(0) self._next: asyncio.Task = self._current async def _perform_refresh(self) -> ConnectionInfo: """ Perform a refresh operation on an AlloyDB instance. Retrieves metadata and generates new client certificate required to connect securely to the AlloyDB instance. Returns: ConnectionInfo: Result of the refresh operation. """ self._refresh_in_progress.set() logger.debug( f"['{self._instance_uri}']: Connection info refresh operation started" ) try: await self._refresh_rate_limiter.acquire() connection_info = await self._client.get_connection_info( self._project, self._region, self._cluster, self._name, self._keys, ) logger.debug( f"['{self._instance_uri}']: Connection info refresh operation" " complete" ) logger.debug( f"['{self._instance_uri}']: Current certificate expiration = " f"{connection_info.expiration.isoformat()}" ) except Exception as e: logger.debug( f"['{self._instance_uri}']: Connection info refresh operation" f" failed: {str(e)}" ) raise finally: self._refresh_in_progress.clear() return connection_info def _schedule_refresh(self, delay: int) -> asyncio.Task: """ Schedule a refresh operation. Args: delay (int): Time in seconds to sleep before performing refresh. Returns: asyncio.Task[ConnectionInfo]: A task representing the scheduled refresh operation. """ return asyncio.create_task(self._refresh_operation(delay)) async def _refresh_operation(self, delay: int) -> ConnectionInfo: """ A coroutine that sleeps for the specified amount of time before running _perform_refresh. Args: delay (int): Time in seconds to sleep before performing refresh. Returns: ConnectionInfo: Refresh result for an AlloyDB instance. """ refresh_task: asyncio.Task try: if delay > 0: await asyncio.sleep(delay) refresh_task = asyncio.create_task(self._perform_refresh()) refresh_result = await refresh_task # check that refresh is valid if not await _is_valid(refresh_task): raise RefreshError( f"['{self._instance_uri}']: Invalid refresh operation. Certficate appears to be expired." ) except asyncio.CancelledError: logger.debug( f"['{self._instance_uri}']: Scheduled refresh operation cancelled" ) raise # bad refresh attempt except Exception: logger.info( f"['{self._instance_uri}']: " "An error occurred while performing refresh. " "Scheduling another refresh attempt immediately" ) # check if current refresh result is invalid (expired), # don't want to replace valid result 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 refresh result and schedule next refresh self._current = refresh_task # calculate refresh delay based on certificate expiration delay = _seconds_until_refresh(refresh_result.expiration) logger.debug( f"['{self._instance_uri}']: 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_result async def force_refresh(self) -> None: """ Schedules a new refresh operation 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 connect_info(self) -> ConnectionInfo: """Retrieves ConnectionInfo instance for establishing a secure connection to the AlloyDB instance. """ return await self._current async def close(self) -> None: """ Cancel refresh tasks. """ logger.debug( f"['{self._instance_uri}']: 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)