Skip to content

Commit

Permalink
fix: throw error if Auto IAM AuthN is unsupported (#476)
Browse files Browse the repository at this point in the history
  • Loading branch information
jackwotherspoon authored Sep 22, 2022
1 parent c4788c2 commit fef0cd7
Show file tree
Hide file tree
Showing 7 changed files with 95 additions and 12 deletions.
29 changes: 24 additions & 5 deletions google/cloud/sql/connector/instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,15 @@ def __init__(self, *args: Any) -> None:
super(ExpiredInstanceMetadata, self).__init__(self, *args)


class AutoIAMAuthNotSupported(Exception):
"""
Exception to be raised when Automatic IAM Authentication is not
supported with database engine version.
"""

pass


class InstanceMetadata:
ip_addrs: Dict[str, Any]
context: ssl.SSLContext
Expand All @@ -115,7 +124,6 @@ def __init__(
enable_iam_auth: bool,
) -> None:
self.ip_addrs = ip_addrs

self.context = ssl.SSLContext(ssl.PROTOCOL_TLS)

# verify OpenSSL version supports TLSv1.3
Expand Down Expand Up @@ -366,10 +374,21 @@ async def _perform_refresh(self) -> InstanceMetadata:
self._enable_iam_auth,
)
)

metadata, ephemeral_cert = await asyncio.gather(
metadata_task, ephemeral_task
)
try:
metadata = await metadata_task
# check if automatic IAM database authn is supported for database engine
if self._enable_iam_auth and not metadata[
"database_version"
].startswith("POSTGRES"):
raise AutoIAMAuthNotSupported(
f"'{metadata['database_version']}' does not support automatic IAM authentication. It is only supported with Cloud SQL Postgres instances."
)
except Exception:
# cancel ephemeral cert task if exception occurs before it is awaited
ephemeral_task.cancel()
raise

ephemeral_cert = await ephemeral_task

x509 = load_pem_x509_certificate(
ephemeral_cert.encode("UTF-8"), default_backend()
Expand Down
1 change: 1 addition & 0 deletions google/cloud/sql/connector/refresh_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ async def _get_metadata(
metadata = {
"ip_addresses": {ip["type"]: ip["ipAddress"] for ip in ret_dict["ipAddresses"]},
"server_ca_cert": ret_dict["serverCaCert"]["cert"],
"database_version": ret_dict["databaseVersion"],
}

return metadata
Expand Down
9 changes: 6 additions & 3 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,20 +162,23 @@ async def instance(
# mock Cloud SQL Admin API calls
with aioresponses() as mocked:
mocked.get(
"https://sqladmin.googleapis.com/sql/v1beta4/projects/my-project/instances/my-instance/connectSettings",
f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{mock_instance.project}/instances/{mock_instance.name}/connectSettings",
status=200,
body=mock_instance.connect_settings(),
repeat=True,
)
mocked.post(
"https://sqladmin.googleapis.com/sql/v1beta4/projects/my-project/instances/my-instance:generateEphemeralCert",
f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{mock_instance.project}/instances/{mock_instance.name}:generateEphemeralCert",
status=200,
body=mock_instance.generate_ephemeral(client_key),
repeat=True,
)

instance = Instance(
"my-project:my-region:my-instance", "pg8000", keys, event_loop
f"{mock_instance.project}:{mock_instance.region}:{mock_instance.name}",
"pg8000",
keys,
event_loop,
)

yield instance
Expand Down
33 changes: 33 additions & 0 deletions tests/system/test_connector_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,13 @@
"""
import asyncio
import os
import pytest
import pymysql
import sqlalchemy
import logging
import google.auth
from google.cloud.sql.connector import Connector
from google.cloud.sql.connector.instance import AutoIAMAuthNotSupported
import datetime
import concurrent.futures
from threading import Thread
Expand Down Expand Up @@ -141,3 +143,34 @@ def test_connector_with_custom_loop() -> None:
assert result[0] == 1
# assert that Connector does not start its own thread
assert connector._thread is None


def test_connector_mysql_iam_auth_error() -> None:
"""
Test that connecting with enable_iam_auth set to True
for MySQL raises exception.
"""
with pytest.raises(AutoIAMAuthNotSupported):
with Connector(enable_iam_auth=True) as connector:
connector.connect(
os.environ["MYSQL_CONNECTION_NAME"],
"pymysql",
user="my-user",
db="my-db",
)


def test_connector_sqlserver_iam_auth_error() -> None:
"""
Test that connecting with enable_iam_auth set to True
for SQL Server raises exception.
"""
with pytest.raises(AutoIAMAuthNotSupported):
with Connector(enable_iam_auth=True) as connector:
connector.connect(
os.environ["SQLSERVER_CONNECTION_NAME"],
"pytds",
user="my-user",
password="my-pass",
db="my-db",
)
10 changes: 8 additions & 2 deletions tests/unit/mocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,11 +178,17 @@ async def create_ssl_context() -> ssl.SSLContext:


class FakeCSQLInstance:
def __init__(self, project: str, region: str, name: str) -> None:
def __init__(
self,
project: str = "my-project",
region: str = "my-region",
name: str = "my-instance",
db_version: str = "POSTGRES_14",
) -> None:
self.project = project
self.region = region
self.name = name
self.db_version = "POSTGRES_14" # arbitrary value
self.db_version = db_version
self.ip_addrs = {"PRIMARY": "0.0.0.0", "PRIVATE": "1.1.1.1"}
self.backend_type = "SECOND_GEN"

Expand Down
19 changes: 19 additions & 0 deletions tests/unit/test_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from google.cloud.sql.connector.rate_limiter import AsyncRateLimiter
from google.auth.credentials import Credentials
from google.cloud.sql.connector.instance import (
AutoIAMAuthNotSupported,
IPTypes,
Instance,
CredentialsTypeError,
Expand Down Expand Up @@ -402,3 +403,21 @@ async def test_ClientResponseError(
)
finally:
await instance.close()


@pytest.mark.asyncio
@pytest.mark.parametrize(
"mock_instance",
[
mocks.FakeCSQLInstance(db_version="SQLSERVER_2019_STANDARD"),
mocks.FakeCSQLInstance(db_version="MYSQL_8_0"),
],
)
async def test_AutoIAMAuthNotSupportedError(instance: Instance) -> None:
"""
Test that AutoIAMAuthNotSupported exception is raised
for SQL Server and MySQL instances.
"""
instance._enable_iam_auth = True
with pytest.raises(AutoIAMAuthNotSupported):
await instance._current
6 changes: 4 additions & 2 deletions tests/unit/test_refresh_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,8 +164,10 @@ async def test_get_metadata(
instance,
)

assert result["ip_addresses"] is not None and isinstance(
result["server_ca_cert"], str
assert (
result["ip_addresses"] is not None
and result["database_version"] == "POSTGRES_14"
and isinstance(result["server_ca_cert"], str)
)


Expand Down

0 comments on commit fef0cd7

Please sign in to comment.