Repository: GoogleCloudPlatform/cloud-sql-python-connector (python)

Issue:
Validate email address format when using automatic IAM authentication
When adding IAM database users to either MySQL or Postgres the user or service account username is formatted in a certain way that differs across database engine types.

This is not necessarily trivial to users when using the connectors, we should try validating and formatting the IAM database usernames within the connectors to remove this burden on users.

**Postgres**: Removes `.gserviceaccount.com` domain suffix if it exists.
**MySQL**: Removes everything after and including the `@` sign. (`test-user@gmail.com` -> `test-user`)

Should already have `databaseVersion` from the API request to https://cloud.google.com/sql/docs/mysql/admin-api/rest/v1beta4/connect/get 

Pull request:
feat: format `user` argument for automatic IAM authn
Adding a formatting step to format the database `user` argument when using automatic IAM database authentication. This will abstract away the different formatting semantics between MySQL and Postgres so that users can leverage the connector without needing to worry about learning naming rules.

**Postgres**: Removes `.gserviceaccount.com` domain suffix if it exists.
**MySQL**: Removes everything after and including the `@` sign. (`test-user@gmail.com` -> `test-user`)

Closes #429 

Patch:
diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py
index 94887b8..c836e3c 100755
--- a/google/cloud/sql/connector/connector.py
+++ b/google/cloud/sql/connector/connector.py
@@ -24,7 +24,7 @@ import google.cloud.sql.connector.pymysql as pymysql
 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 @@ class Connector:
         # 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)
diff --git a/google/cloud/sql/connector/instance.py b/google/cloud/sql/connector/instance.py
index 53511e2..3c2cdf0 100644
--- a/google/cloud/sql/connector/instance.py
+++ b/google/cloud/sql/connector/instance.py
@@ -64,11 +64,13 @@ 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,
@@ -76,6 +78,7 @@ class InstanceMetadata:
         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 @@ class Instance:
 
         return InstanceMetadata(
             ephemeral_cert,
+            metadata["database_version"],
             metadata["ip_addresses"],
             priv_key,
             metadata["server_ca_cert"],
diff --git a/google/cloud/sql/connector/utils.py b/google/cloud/sql/connector/utils.py
index fb05235..21bebd6 100755
--- a/google/cloud/sql/connector/utils.py
+++ b/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


Test patch:
diff --git a/tests/system/test_connector_object.py b/tests/system/test_connector_object.py
index 532c160..70b509b 100644
--- a/tests/system/test_connector_object.py
+++ b/tests/system/test_connector_object.py
@@ -21,7 +21,9 @@ import sqlalchemy
 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
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index ffaf189..0b9b72a 100644
--- a/tests/unit/test_utils.py
+++ b/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", "service-account@test.iam"
+    )
+    service_account2 = utils.format_database_user(
+        "POSTGRES_14", "service-account@test.iam.gserviceaccount.com"
+    )
+    assert service_account == "service-account@test.iam"
+    assert service_account2 == "service-account@test.iam"
+    user = utils.format_database_user("POSTGRES_14", "test@test.com")
+    assert user == "test@test.com"
+
+
+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", "service-account@test.iam"
+    )
+    service_account2 = utils.format_database_user(
+        "MYSQL_8_0", "service-account@test.iam.gserviceaccount.com"
+    )
+    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", "test@test.com")
+    user2 = utils.format_database_user("MYSQL_8_0", "test")
+    assert user == "test"
+    assert user2 == "test"