feat: support syncing service accounts with GroupSync (#73)

Author: jackwotherspoon

iam_groups_authn/postgres.py

@@ -22,6 +22,23 @@
 from google.auth.transport.requests import Request
 
 
+def postgres_username(iam_email):
+    """Get Postgres username from user or service account email.
+
+    Given an IAM user or IAM service account email, format their Postgres
+    database username accordingly. Do nothing for user emails, remove
+    '.gserviceaccount.com' suffix from service account emails.
+
+    Args:
+        iam_email: An IAM user or service account email.
+
+    Returns:
+        username: The IAM user or service account Postgres DB username.
+    """
+    username = iam_email.removesuffix(".gserviceaccount.com")
+    return username
+
+
 class PostgresRoleService(RoleService):
     """Class for managing a Postgres DB user's role grants."""
 

iam_groups_authn/sql_admin.py

@@ -16,6 +16,7 @@
 
 from typing import NamedTuple
 from iam_groups_authn.mysql import mysql_username
+from iam_groups_authn.postgres import postgres_username
 
 
 class InstanceConnectionName(NamedTuple):
@@ -76,10 +77,14 @@ async def add_missing_db_users(
             [user for user in iam_users if mysql_username(user) not in db_users]
         )
     else:
-        missing_db_users = set([user for user in iam_users if user not in db_users])
+        missing_db_users = set(
+            [user for user in iam_users if postgres_username(user) not in db_users]
+        )
     # add missing users to database instance
     for user in missing_db_users:
         await user_service.insert_db_user(
-            user, InstanceConnectionName(*instance_connection_name.split(":"))
+            user,
+            InstanceConnectionName(*instance_connection_name.split(":")),
+            database_type,
         )
     return missing_db_users

iam_groups_authn/sync.py

@@ -37,6 +37,7 @@
 from iam_groups_authn.postgres import (
     init_postgres_connection_engine,
     PostgresRoleService,
+    postgres_username,
 )
 
 
@@ -298,7 +299,7 @@ async def get_db_users(self, instance_connection_name):
                 f"Error: Failed to get the database users for instance `{instance_connection_name}`. Verify instance connection name and instance details."
             ) from e
 
-    async def insert_db_user(self, user_email, instance_connection_name):
+    async def insert_db_user(self, user_email, instance_connection_name, database_type):
         """Create DB user from IAM user.
 
         Given an IAM user's email, insert the IAM user as a DB user for Cloud SQL instance.
@@ -308,12 +309,27 @@ async def insert_db_user(self, user_email, instance_connection_name):
             instance_connection_name: InstanceConnectionName namedTuple.
                 (e.g. InstanceConnectionName(project='my-project', region='my-region',
                 instance='my-instance'))
+            database_type: Cloud SQL database version.
         """
         # build request to SQL Admin API
         project = instance_connection_name.project
         instance = instance_connection_name.instance
         url = f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{project}/instances/{instance}/users"
-        user = {"name": user_email, "type": "CLOUD_IAM_USER"}
+        # if service account, add service account IAM database user
+        if user_email.endswith(".gserviceaccount.com"):
+            # the Cloud SQL Admin API doesn't format Postgres usernames, but does format MySQL usernames
+            if database_type.is_mysql():
+                user = {
+                    "name": user_email,
+                    "type": "CLOUD_IAM_SERVICE_ACCOUNT",
+                }
+            else:
+                user = {
+                    "name": user_email.removesuffix(".gserviceaccount.com"),
+                    "type": "CLOUD_IAM_SERVICE_ACCOUNT",
+                }
+        else:
+            user = {"name": user_email, "type": "CLOUD_IAM_USER"}
 
         try:
             # call the SQL Admin API
@@ -483,6 +499,9 @@ async def grant_iam_group_role(
     if database_type.is_mysql():
         # truncate mysql_usernames
         iam_users = [mysql_username(user) for user in iam_users]
+    elif database_type.is_postgres():
+        # truncate postgres service accounts
+        iam_users = [postgres_username(user) for user in iam_users]
 
     # find DB users who are part of IAM group that need role granted to them
     users_to_grant = [user for user in iam_users if user not in users_with_roles]

tests/integration/test_mysql_sync.py

@@ -29,6 +29,7 @@
 sql_instance = os.environ["MYSQL_INSTANCE"]
 iam_groups = [os.environ["IAM_GROUPS"]]
 test_user = os.environ["TEST_USER"]
+sa_user = os.environ["SA_USER"]
 
 scopes = [
     "https://www.googleapis.com/auth/admin.directory.group.member",
@@ -73,8 +74,12 @@ def setup_and_teardown():
     try:
         # cleanup user from database
         delete_database_user(sql_instance, mysql_username(test_user), credentials)
+        # cleanup service account from database
+        delete_database_user(sql_instance, mysql_username(sa_user), credentials)
         # re-add member to IAM group
         add_iam_member(iam_groups[0], test_user, credentials)
+        # re-add service account to IAM group
+        add_iam_member(iam_groups[0], sa_user, credentials)
         # wait 30 seconds, adding IAM member is slow
         time.sleep(30)
     except Exception:
@@ -95,11 +100,12 @@ async def test_service_mysql(credentials):
         - Verifies test user no longer has group role
     """
 
-    # remove database user if they exist
+    # remove database users if they exist
     try:
         delete_database_user(sql_instance, mysql_username(test_user), credentials)
+        delete_database_user(sql_instance, mysql_username(sa_user), credentials)
     except Exception:
-        print("Database user must already have been deleted!")
+        print("Database users must already have been deleted!")
 
     # create aiohttp client session for async API calls
     client_session = ClientSession(headers={"Content-Type": "application/json"})
@@ -108,21 +114,24 @@ async def test_service_mysql(credentials):
     user_service = UserService(client_session, credentials)
     db_users = await get_instance_users(user_service, sql_instance)
     assert mysql_username(test_user) not in db_users
+    assert mysql_username(sa_user) not in db_users
 
-    # make sure test_user is member of IAM group
+    # make sure users are members of IAM group
     try:
         add_iam_member(iam_groups[0], test_user, credentials)
+        add_iam_member(iam_groups[0], sa_user, credentials)
         # wait 30 seconds, adding IAM member is slow
         time.sleep(30)
     except Exception:
-        print("Member must already belong to IAM Group.")
+        print("Members must already belong to IAM Group.")
 
     # run groups sync
     await groups_sync(iam_groups, [sql_instance], credentials, dict(), False)
 
     # check that test_user has been created as database user
     db_users = await get_instance_users(user_service, sql_instance)
     assert mysql_username(test_user) in db_users
+    assert mysql_username(sa_user) in db_users
 
     # create database connection to instance
     pool = init_mysql_connection_engine(sql_instance, credentials)
@@ -134,8 +143,9 @@ async def test_service_mysql(credentials):
         for member in iam_members:
             assert mysql_username(member) in users_with_role
 
-    # remove test_user from IAM group
+    # remove users from IAM group
     delete_iam_member(iam_groups[0], test_user, credentials)
+    delete_iam_member(iam_groups[0], sa_user, credentials)
 
     # wait 30 seconds, deleting IAM member is slow
     time.sleep(30)
@@ -146,6 +156,7 @@ async def test_service_mysql(credentials):
     # verify test_user no longer has group role
     users_with_role = check_role_mysql(pool, mysql_username(iam_groups[0]))
     assert mysql_username(test_user) not in users_with_role
+    assert mysql_username(sa_user) not in users_with_role
 
     # close aiohttp client session for graceful exit
     if not client_session.closed:

tests/integration/test_postgres_sync.py

@@ -21,7 +21,7 @@
 from helpers import delete_database_user, delete_iam_member, add_iam_member
 from iam_groups_authn.iam_admin import get_iam_users
 from iam_groups_authn.mysql import mysql_username
-from iam_groups_authn.postgres import init_postgres_connection_engine
+from iam_groups_authn.postgres import init_postgres_connection_engine, postgres_username
 from iam_groups_authn.sql_admin import get_instance_users
 from iam_groups_authn.sync import groups_sync, UserService
 import time
@@ -30,6 +30,7 @@
 sql_instance = os.environ["POSTGRES_INSTANCE"]
 iam_groups = [os.environ["IAM_GROUPS"]]
 test_user = os.environ["TEST_USER"]
+sa_user = os.environ["SA_USER"]
 
 scopes = [
     "https://www.googleapis.com/auth/admin.directory.group.member",
@@ -74,8 +75,12 @@ def setup_and_teardown():
     try:
         # cleanup user from database
         delete_database_user(sql_instance, test_user, credentials)
+        # cleanup service account from database
+        delete_database_user(sql_instance, postgres_username(sa_user), credentials)
         # re-add member to IAM group
         add_iam_member(iam_groups[0], test_user, credentials)
+        # re-add service account to IAM group
+        add_iam_member(iam_groups[0], sa_user, credentials)
         # wait 30 seconds, adding IAM member is slow
         time.sleep(30)
     except Exception:
@@ -96,23 +101,26 @@ async def test_service_postgres(credentials):
         - Verifies test user no longer has group role
     """
 
-    # remove database user if they exist
+    # remove database users if they exist
     try:
         delete_database_user(sql_instance, test_user, credentials)
+        delete_database_user(sql_instance, postgres_username(sa_user), credentials)
     except Exception:
-        print("Database user must already have been deleted!")
+        print("Database users must already have been deleted!")
 
     # create aiohttp client session for async API calls
     client_session = ClientSession(headers={"Content-Type": "application/json"})
 
-    # check that test_user is not a database user
+    # check that users are not a database user
     user_service = UserService(client_session, credentials)
     db_users = await get_instance_users(user_service, sql_instance)
     assert test_user not in db_users
+    assert postgres_username(sa_user) not in db_users
 
-    # make sure test_user is member of IAM group
+    # make sure users are members of IAM group
     try:
         add_iam_member(iam_groups[0], test_user, credentials)
+        add_iam_member(iam_groups[0], sa_user, credentials)
         # wait 30 seconds, adding IAM member is slow
         time.sleep(30)
     except Exception:
@@ -121,9 +129,10 @@ async def test_service_postgres(credentials):
     # run groups sync
     await groups_sync(iam_groups, [sql_instance], credentials, dict(), False)
 
-    # check that test_user has been created as database user
+    # check that users has been created as database users
     db_users = await get_instance_users(user_service, sql_instance)
     assert test_user in db_users
+    assert postgres_username(sa_user) in db_users
 
     # create database connection to instance
     pool = init_postgres_connection_engine(sql_instance, credentials)
@@ -133,10 +142,11 @@ async def test_service_postgres(credentials):
         users_with_role = check_role_postgres(pool, mysql_username(iam_group))
         iam_members = await get_iam_users(user_service, iam_group)
         for member in iam_members:
-            assert member in users_with_role
+            assert postgres_username(member) in users_with_role
 
-    # remove test_user from IAM group
+    # remove users from IAM group
     delete_iam_member(iam_groups[0], test_user, credentials)
+    delete_iam_member(iam_groups[0], sa_user, credentials)
 
     # wait 30 seconds, deleting IAM member is slow
     time.sleep(30)
@@ -147,6 +157,7 @@ async def test_service_postgres(credentials):
     # verify test_user no longer has group role
     users_with_role = check_role_postgres(pool, mysql_username(iam_groups[0]))
     assert test_user not in users_with_role
+    assert postgres_username(sa_user) not in users_with_role
 
     # close aiohttp client session for graceful exit
     if not client_session.closed:

tests/unit/test_add_missing_db_users.py

@@ -24,7 +24,7 @@ class FakeUserService:
     def __init__(self):
         pass
 
-    async def insert_db_user(self, user, instance_connection_name):
+    async def insert_db_user(self, user, instance_connection_name, database_type):
         pass
 
 
@@ -35,11 +35,13 @@ async def test_no_missing_users():
     """
     user_service = FakeUserService()
     iam_future = asyncio.Future()
-    iam_future.set_result(["user1@test.com", "user2@test.com"])
+    iam_future.set_result(
+        ["user1@test.com", "user2@test.com", "sa@test.iam.gserviceaccount.com"]
+    )
     mysql_users = asyncio.Future()
-    mysql_users.set_result(["user1", "user2"])
+    mysql_users.set_result(["user1", "user2", "sa"])
     postgres_users = asyncio.Future()
-    postgres_users.set_result(["user1@test.com", "user2@test.com"])
+    postgres_users.set_result(["user1@test.com", "user2@test.com", "sa@test.iam"])
 
     missing_mysql_users = await add_missing_db_users(
         user_service,
@@ -60,13 +62,20 @@ async def test_no_missing_users():
 
 
 @pytest.mark.asyncio
-async def test_missing__users():
+async def test_missing_users():
     """Test where there are IAM users missing corresponding database user.
     Should return set of the emails of IAM users missing database user.
     """
     user_service = FakeUserService()
     iam_future = asyncio.Future()
-    iam_future.set_result(["user1@test.com", "user2@test.com", "user3@test.com"])
+    iam_future.set_result(
+        [
+            "user1@test.com",
+            "user2@test.com",
+            "user3@test.com",
+            "sa@test.iam.gserviceaccount.com",
+        ]
+    )
     postgres_users = asyncio.Future()
     postgres_users.set_result(["user1@test.com"])
     mysql_users = asyncio.Future()
@@ -79,7 +88,9 @@ async def test_missing__users():
         "group:region:instance",
         DatabaseVersion.POSTGRES_13,
     )
-    assert missing_iam_users == set(["user2@test.com", "user3@test.com"])
+    assert missing_iam_users == set(
+        ["user2@test.com", "user3@test.com", "sa@test.iam.gserviceaccount.com"]
+    )
 
     missing_iam_users = await add_missing_db_users(
         user_service,
@@ -88,7 +99,9 @@ async def test_missing__users():
         "group:region:instance",
         DatabaseVersion.MYSQL_8_0,
     )
-    assert missing_iam_users == set(["user2@test.com", "user3@test.com"])
+    assert missing_iam_users == set(
+        ["user2@test.com", "user3@test.com", "sa@test.iam.gserviceaccount.com"]
+    )
 
 
 @pytest.mark.asyncio