Skip to content

Commit

Permalink
feat: validate driver with database engine (#1108)
Browse files Browse the repository at this point in the history
Removing ambiguous behavior if a user configures a database
driver to connect to an incompatible database engine.

i.e. pymysql (MySQL driver) -> Cloud SQL for PostgreSQL database

This behavior would timeout trying to connect and then throw
an ambiguous timeout error.

The Connector.connect already know which database engine
the Cloud SQL database is prior to attempting to connect via the
driver. So we can validate if the driver is compatible and throw an
actionable error if not.
  • Loading branch information
jackwotherspoon authored Jun 3, 2024
1 parent 821cef8 commit a7c7861
Show file tree
Hide file tree
Showing 4 changed files with 81 additions and 1 deletion.
3 changes: 3 additions & 0 deletions google/cloud/sql/connector/connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

import google.cloud.sql.connector.asyncpg as asyncpg
from google.cloud.sql.connector.client import CloudSQLClient
from google.cloud.sql.connector.enums import DriverMapping
from google.cloud.sql.connector.exceptions import ConnectorLoopError
from google.cloud.sql.connector.exceptions import DnsNameResolutionError
from google.cloud.sql.connector.instance import IPTypes
Expand Down Expand Up @@ -332,6 +333,8 @@ async def connect_async(
# attempt to make connection to Cloud SQL instance
try:
conn_info = await cache.connect_info()
# validate driver matches intended database engine
DriverMapping.validate_engine(driver, conn_info.database_version)
ip_address = conn_info.get_preferred_ip(ip_type)
# resolve DNS name into IP address for PSC
if ip_type.value == "PSC":
Expand Down
46 changes: 46 additions & 0 deletions google/cloud/sql/connector/enums.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
# 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 enum import Enum

from google.cloud.sql.connector.exceptions import IncompatibleDriverError


class DriverMapping(Enum):
"""Maps a given database driver to it's corresponding database engine."""

ASYNCPG = "POSTGRES"
PG8000 = "POSTGRES"
PYMYSQL = "MYSQL"
PYTDS = "SQLSERVER"

@staticmethod
def validate_engine(driver: str, engine_version: str) -> None:
"""Validate that the given driver is compatible with the given engine.
Args:
driver (str): Database driver being used. (i.e. "pg8000")
engine_version (str): Database engine version. (i.e. "POSTGRES_16")
Raises:
IncompatibleDriverError: If the given driver is not compatible with
the given engine.
"""
mapping = DriverMapping[driver.upper()]
if not engine_version.startswith(mapping.value):
raise IncompatibleDriverError(
f"Database driver '{driver}' is incompatible with database "
f"version '{engine_version}'. Given driver can "
f"only be used with Cloud SQL {mapping.value} databases."
)
7 changes: 7 additions & 0 deletions google/cloud/sql/connector/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,3 +70,10 @@ class RefreshNotValidError(Exception):
"""

pass


class IncompatibleDriverError(Exception):
"""
Exception to be raised when the database driver given is for the wrong
database engine. (i.e. asyncpg for a MySQL database)
"""
26 changes: 25 additions & 1 deletion tests/unit/test_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
See the License for the specific language governing permissions and
limitations under the License.
"""

import asyncio
from typing import Union

Expand All @@ -25,6 +26,7 @@
from google.cloud.sql.connector import IPTypes
from google.cloud.sql.connector.client import CloudSQLClient
from google.cloud.sql.connector.exceptions import ConnectorLoopError
from google.cloud.sql.connector.exceptions import IncompatibleDriverError
from google.cloud.sql.connector.instance import RefreshAheadCache


Expand All @@ -46,10 +48,32 @@ def test_connect_enable_iam_auth_error(
"If you require both for your use case, please use a new "
"connector.Connector object."
)
# remove cache entrry to avoid destructor warnings
# remove cache entry to avoid destructor warnings
connector._cache = {}


async def test_connect_incompatible_driver_error(
fake_credentials: Credentials,
fake_client: CloudSQLClient,
) -> None:
"""Test that calling connect() with driver that is incompatible with
database version throws error."""
connect_string = "test-project:test-region:test-instance"
async with Connector(
credentials=fake_credentials, loop=asyncio.get_running_loop()
) as connector:
connector._client = fake_client
# try to connect using pymysql driver to a Postgres database
with pytest.raises(IncompatibleDriverError) as exc_info:
await connector.connect_async(connect_string, "pymysql")
assert (
exc_info.value.args[0]
== "Database driver 'pymysql' is incompatible with database version"
" 'POSTGRES_15'. Given driver can only be used with Cloud SQL MYSQL"
" databases."
)


def test_connect_with_unsupported_driver(fake_credentials: Credentials) -> None:
with Connector(credentials=fake_credentials) as connector:
# try to connect using unsupported driver, should raise KeyError
Expand Down

0 comments on commit a7c7861

Please sign in to comment.