feat: format user argument for automatic IAM authn (GoogleCloudPlatform#449)

Author: jackwotherspoon

google/cloud/sql/connector/connector.py

@@ -24,7 +24,7 @@
 import google.cloud.sql.connector.pg8000 as pg8000
 import google.cloud.sql.connector.pytds as pytds
 import google.cloud.sql.connector.asyncpg as asyncpg
-from google.cloud.sql.connector.utils import generate_keys
+from google.cloud.sql.connector.utils import generate_keys, format_database_user
 from google.cloud.sql.connector.exceptions import ConnectorLoopError
 from google.auth.credentials import Credentials
 from threading import Thread
@@ -235,6 +235,18 @@ async def connect_async(
         # helper function to wrap in timeout
         async def get_connection() -> Any:
             instance_data, ip_address = await instance.connect_info(ip_type)
+
+            # format `user` param for automatic IAM database authn
+            if enable_iam_auth:
+                formatted_user = format_database_user(
+                    instance_data.database_version, kwargs["user"]
+                )
+                if formatted_user != kwargs["user"]:
+                    logger.debug(
+                        f"['{instance_connection_string}']: Truncated IAM database username from {kwargs['user']} to {formatted_user}"
+                    )
+                    kwargs["user"] = formatted_user
+
             # async drivers are unblocking and can be awaited directly
             if driver in ASYNC_DRIVERS:
                 return await connector(ip_address, instance_data.context, **kwargs)

google/cloud/sql/connector/instance.py

@@ -64,18 +64,21 @@ class IPTypes(Enum):
 class InstanceMetadata:
     ip_addrs: Dict[str, Any]
     context: ssl.SSLContext
+    database_version: str
     expiration: datetime.datetime
 
     def __init__(
         self,
         ephemeral_cert: str,
+        database_version: str,
         ip_addrs: Dict[str, Any],
         private_key: bytes,
         server_ca_cert: str,
         expiration: datetime.datetime,
         enable_iam_auth: bool,
     ) -> None:
         self.ip_addrs = ip_addrs
+        self.database_version = database_version
         self.context = ssl.SSLContext(ssl.PROTOCOL_TLS)
 
         # verify OpenSSL version supports TLSv1.3
@@ -370,6 +373,7 @@ async def _perform_refresh(self) -> InstanceMetadata:
 
         return InstanceMetadata(
             ephemeral_cert,
+            metadata["database_version"],
             metadata["ip_addresses"],
             priv_key,
             metadata["server_ca_cert"],

google/cloud/sql/connector/utils.py

@@ -77,3 +77,28 @@ def write_to_file(
         priv_out.write(priv_key)
 
     return (ca_filename, cert_filename, key_filename)
+
+
+def format_database_user(database_version: str, user: str) -> str:
+    """
+    Format database `user` param for Cloud SQL automatic IAM authentication.
+
+    :type database_version: str
+    :param database_version
+        Cloud SQL database version. (i.e. POSTGRES_14, MYSQL8_0, etc.)
+
+    :type user: str
+    :param user
+        Database username to connect to Cloud SQL database with.
+    """
+    # remove suffix for Postgres service accounts
+    if database_version.startswith("POSTGRES"):
+        suffix = ".gserviceaccount.com"
+        user = user[: -len(suffix)] if user.endswith(suffix) else user
+        return user
+
+    # remove everything after and including the @ for MySQL
+    if database_version.startswith("MYSQL") and "@" in user:
+        return user.split("@")[0]
+
+    return user

tests/system/test_connector_object.py

@@ -21,7 +21,9 @@
 import logging
 import google.auth
 from google.cloud.sql.connector import Connector
-from google.cloud.sql.connector.instance import AutoIAMAuthNotSupported
+from google.cloud.sql.connector.exceptions import (
+    AutoIAMAuthNotSupported,
+)
 import datetime
 import concurrent.futures
 from threading import Thread

tests/unit/test_utils.py

@@ -39,3 +39,39 @@ async def test_generate_keys_returns_bytes_and_str() -> None:
 
     res1, res2 = await utils.generate_keys()
     assert isinstance(res1, bytes) and (isinstance(res2, str))
+
+
+def test_format_database_user_postgres() -> None:
+    """
+    Test that format_database_user properly formats Postgres IAM database users.
+    """
+    service_account = utils.format_database_user(
+        "POSTGRES_14", "[email protected]"
+    )
+    service_account2 = utils.format_database_user(
+        "POSTGRES_14", "[email protected]"
+    )
+    assert service_account == "[email protected]"
+    assert service_account2 == "[email protected]"
+    user = utils.format_database_user("POSTGRES_14", "[email protected]")
+    assert user == "[email protected]"
+
+
+def test_format_database_user_mysql() -> None:
+    """
+    Test that format_database _user properly formats MySQL IAM database users.
+    """
+    service_account = utils.format_database_user(
+        "MYSQL_8_0", "[email protected]"
+    )
+    service_account2 = utils.format_database_user(
+        "MYSQL_8_0", "[email protected]"
+    )
+    service_account3 = utils.format_database_user("MYSQL_8_0", "service-account")
+    assert service_account == "service-account"
+    assert service_account2 == "service-account"
+    assert service_account3 == "service-account"
+    user = utils.format_database_user("MYSQL_8_0", "[email protected]")
+    user2 = utils.format_database_user("MYSQL_8_0", "test")
+    assert user == "test"
+    assert user2 == "test"