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__ importannotations

importasyncio
fromdatetimeimportdatetime
fromdatetimeimporttimedelta
fromdatetimeimporttimezone
importlogging

fromgoogle.cloud.sql.connector.clientimportCloudSQLClient
fromgoogle.cloud.sql.connector.connection_infoimportConnectionInfo
fromgoogle.cloud.sql.connector.connection_infoimportConnectionInfoCache
fromgoogle.cloud.sql.connector.connection_nameimportConnectionName
fromgoogle.cloud.sql.connector.exceptionsimportRefreshNotValidError
fromgoogle.cloud.sql.connector.rate_limiterimportAsyncRateLimiter
fromgoogle.cloud.sql.connector.refresh_utilsimport_is_valid
fromgoogle.cloud.sql.connector.refresh_utilsimport_seconds_until_refresh

logger=logging.getLogger(name=__name__)

APPLICATION_NAME="cloud-sql-python-connector"


classRefreshAheadCache(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
defconn_name(self) ->ConnectionName:
returnself._conn_name

@property
defclosed(self) ->bool:
returnself._closed

asyncdefforce_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
ifnotself._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
ifnotawait_is_valid(self._current):
self._current=self._next

asyncdef_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:
awaitself._refresh_rate_limiter.acquire()
connection_info=awaitself._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()}"
 )

exceptExceptionase:
logger.debug(
f"['{self._conn_name}']: Connection info "
f"refresh operation failed: {str(e)}"
 )
raise

finally:
self._refresh_in_progress.clear()
returnconnection_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.
 """

asyncdef_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:
ifdelay>0:
awaitasyncio.sleep(delay)
refresh_task=asyncio.create_task(self._perform_refresh())
refresh_data=awaitrefresh_task
# check that refresh is valid
ifnotawait_is_valid(refresh_task):
raiseRefreshNotValidError(
f"['{self._conn_name}']: Invalid refresh operation. Certficate appears to be expired."
 )
exceptasyncio.CancelledError:
logger.debug(
f"['{self._conn_name}']: Scheduled refresh"" operation cancelled"
 )
raise
# bad refresh attempt
exceptExceptionase:
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
ifnotawait_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)

returnrefresh_data

# schedule refresh task and return it
scheduled_task=asyncio.create_task(_refresh_task(self, delay))
returnscheduled_task

asyncdefconnect_info(self) ->ConnectionInfo:
"""Retrieves ConnectionInfo instance for establishing a secure
 connection to the Cloud SQL instance.
 """
returnawaitself._current

asyncdefclose(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)
awaitasyncio.wait_for(tasks, timeout=2.0)
self._closed=True