feat: improve concurrency of independent tasks (#37)

* Update get_instance_users for single instance

* Update get_iam_users for single group

* Improve concurrency of tasks

* Update two tests

Author: jackwotherspoon

app.py

@@ -18,12 +18,14 @@
 from google.cloud.sql.connector.instance_connection_manager import IPTypes
 from iam_groups_authn.sync import (
     get_credentials,
-    get_users_to_add,
-    manage_instance_users,
+    get_users_with_roles,
+    revoke_iam_group_role,
+    grant_iam_group_role,
     UserService,
 )
-from iam_groups_authn.sql_admin import get_instance_users, InstanceConnectionName
+from iam_groups_authn.sql_admin import get_instance_users, add_missing_db_users
 from iam_groups_authn.iam_admin import get_iam_users
+from iam_groups_authn.mysql import init_connection_engine, RoleService, mysql_username
 
 # define scopes
 SCOPES = [
@@ -65,6 +67,9 @@ async def run_groups_authn():
             400,
         )
 
+    # set ip_type to proper type for connector
+    ip_type = IPTypes.PRIVATE if private_ip else IPTypes.PUBLIC
+
     # grab default creds from cloud run service account
     creds, project = default()
     # update default credentials with IAM and SQL admin scopes
@@ -73,36 +78,76 @@ async def run_groups_authn():
     # create UserService object for API calls
     user_service = UserService(creds)
 
-    iam_users, instance_users = await asyncio.gather(
-        get_iam_users(user_service, iam_groups),
-        get_instance_users(user_service, sql_instances),
-    )
-
-    # get IAM users of each IAM group
-    for group_name, user_list in iam_users.items():
-        print(f"IAM Users in Group {group_name}: {user_list}")
-
-    # get all instance DB users
-    for instance_name, db_users in instance_users.items():
-        print(f"DB Users in instance `{instance_name}`: {db_users}")
-
-    # find IAM users who are missing as DB users
-    users_to_add = get_users_to_add(iam_users, instance_users)
-    for instance, users in users_to_add.items():
-        print(f"Missing IAM DB users for instance `{instance}`: {users}")
-        for user in users:
-            user_service.insert_db_user(
-                user, InstanceConnectionName(*instance.split(":"))
+    # keep track of IAM group and database instance tasks
+    group_tasks = {}
+    instance_tasks = {}
+
+    # loop iam_groups and sql_instances creating async tasks
+    for group in iam_groups:
+        group_task = asyncio.create_task(get_iam_users(user_service, group))
+        group_tasks[group] = group_task
+
+    for instance in sql_instances:
+        instance_task = asyncio.create_task(get_instance_users(user_service, instance))
+        instance_tasks[instance] = instance_task
+
+    # create pairings of iam groups and instances
+    for group in iam_groups:
+        for instance in sql_instances:
+            # add missing IAM group members to database
+            add_users_task = asyncio.create_task(
+                add_missing_db_users(
+                    user_service, group_tasks[group], instance_tasks[instance], instance
+                )
             )
 
-    # set ip_type to proper type for connector
-    ip_type = IPTypes.PRIVATE if private_ip else IPTypes.PUBLIC
+            # initialize database engine
+            db = init_connection_engine(instance, updated_creds, ip_type)
+            role_service = RoleService(db)
+
+            # verify role for IAM group exists on database, create if does not exist
+            role = mysql_username(group)
+            verify_role_task = asyncio.create_task(role_service.create_group_role(role))
 
-    # for each instance manage users and group role permissions
-    instance_coroutines = [
-        manage_instance_users(instance, iam_users, updated_creds, ip_type)
-        for instance in sql_instances
-    ]
-    await asyncio.gather(*instance_coroutines)
+            # get database users who have group role
+            users_with_roles_task = asyncio.create_task(
+                get_users_with_roles(role_service, role)
+            )
+
+            # await dependent tasks
+            results = await asyncio.gather(
+                add_users_task, verify_role_task, return_exceptions=True
+            )
+            # raise exception if found
+            for result in results:
+                if issubclass(type(result), Exception):
+                    raise result
+
+            # revoke group role from users no longer in IAM group
+            revoke_role_task = asyncio.create_task(
+                revoke_iam_group_role(
+                    role_service,
+                    role,
+                    users_with_roles_task,
+                    group_tasks[group],
+                )
+            )
+
+            # grant group role to IAM users who are missing it on database
+            grant_role_task = asyncio.create_task(
+                grant_iam_group_role(
+                    role_service,
+                    role,
+                    users_with_roles_task,
+                    group_tasks[group],
+                )
+            )
+            results = await asyncio.gather(
+                revoke_role_task, grant_role_task, return_exceptions=True
+            )
+            # raise exception if found
+            for result in results:
+                if issubclass(type(result), Exception):
+                    raise result
 
     return "Sync successful.", 200

iam_groups_authn/iam_admin.py

@@ -15,52 +15,39 @@
 # iam_admin.py contains functions for interacting with the Admin Directory API
 # to access IAM groups and their users
 
-from quart.utils import run_sync
-from collections import defaultdict
-from functools import partial
 
+async def get_iam_users(user_service, group):
+    """Get list of all IAM users within an IAM group.
 
-async def get_iam_users(user_service, groups):
-    """Get list of all IAM users within IAM groups.
-
-    Given a list of IAM groups, get all IAM users that are members within one or
-    more of the groups or a nested child group.
+    Given the email of an IAM group, get all IAM users that are members within
+    the group or a nested child group.
 
     Args:
         user_service: Instance of a UserService object.
-        groups: List of IAM groups. (e.g., ["group@example.com", "abc@example.com"])
+        group: Email of an IAM group. (e.g., "group@example.com")
 
     Returns:
-        iam_users: Set containing all IAM users found within IAM groups.
+        iam_users: Set containing all IAM users found within IAM group.
     """
-    # keep track of iam users using set for no duplicates
-    iam_users = defaultdict(list)
-    # loop through groups and get their IAM users
-    for group in groups:
-        group_queue = [group]
-        # set initial groups searched to input groups
-        searched_groups = group_queue.copy()
-        group_users = set()
-        while group_queue:
-            current_group = group_queue.pop(0)
-            # get all members of current IAM group
-            members_partial = partial(user_service.get_group_members, current_group)
-            members = await run_sync(members_partial)()
-            # check if member is a group, otherwise they are a user
-            for member in members:
-                if member["type"] == "GROUP":
-                    if member["email"] not in searched_groups:
-                        # add current group to searched groups
-                        searched_groups.append(member["email"])
-                        # add group to queue
-                        group_queue.append(member["email"])
-                elif member["type"] == "USER":
-                    # add user to list of group users
-                    group_users.add(member["email"])
-                else:
-                    continue
-        # only add to dict if group has members, allows skipping of not valid groups
-        if group_users:
-            iam_users[group] = group_users
-
-    return iam_users
+    group_queue = [group]
+    # set initial groups searched to input group
+    searched_groups = set(group)
+    group_users = set()
+    while group_queue:
+        current_group = group_queue.pop(0)
+        # get all members of current IAM group
+        members = await user_service.get_group_members(current_group)
+        # check if member is a group, otherwise they are a user
+        for member in members:
+            if member["type"] == "GROUP":
+                if member["email"] not in searched_groups:
+                    # add current group to searched groups
+                    searched_groups.add(member["email"])
+                    # add group to queue
+                    group_queue.append(member["email"])
+            elif member["type"] == "USER":
+                # add user to list of group users
+                group_users.add(member["email"])
+            else:
+                continue
+    return group_users

iam_groups_authn/mysql.py

@@ -15,11 +15,10 @@
 # mysql.py contains all database specific functions for connecting
 # and querying a MySQL database
 
-import asyncio
 import sqlalchemy
 from google.cloud.sql.connector import connector
 from google.cloud.sql.connector.instance_connection_manager import IPTypes
-from functools import partial, wraps
+from iam_groups_authn.utils import async_wrap
 
 
 def mysql_username(iam_email):
@@ -38,23 +37,6 @@ def mysql_username(iam_email):
     return username
 
 
-def async_wrap(func):
-    """Wrapper function to turn synchronous functions into async functions.
-
-    Args:
-        func: Synchronous function to wrap.
-    """
-
-    @wraps(func)
-    async def run(*args, loop=None, executor=None, **kwargs):
-        if loop is None:
-            loop = asyncio.get_event_loop()
-        pfunc = partial(func, *args, **kwargs)
-        return await loop.run_in_executor(executor, pfunc)
-
-    return run
-
-
 class RoleService:
     """Class for managing a DB user's role grants."""
 
@@ -76,11 +58,13 @@ def fetch_role_grants(self, group_name):
         Returns:
             results: List of results for given query.
         """
-        # query role_edges table
-        stmt = sqlalchemy.text(
-            "SELECT FROM_USER, TO_USER FROM mysql.role_edges WHERE FROM_USER= :group_name"
-        )
-        results = self.db.execute(stmt, {"group_name": group_name}).fetchall()
+        # create connection to db instance
+        with self.db.connect() as db_connection:
+            # query role_edges table
+            stmt = sqlalchemy.text(
+                "SELECT FROM_USER, TO_USER FROM mysql.role_edges WHERE FROM_USER= :group_name"
+            )
+            results = db_connection.execute(stmt, {"group_name": group_name}).fetchall()
         return results
 
     @async_wrap
@@ -94,8 +78,10 @@ def create_group_role(self, group):
             db: Database connection pool instance.
             group: Name of group to be verified as role or created as new role.
         """
-        stmt = sqlalchemy.text("CREATE ROLE IF NOT EXISTS :role")
-        self.db.execute(stmt, {"role": group})
+        # create connection to db instance
+        with self.db.connect() as db_connection:
+            stmt = sqlalchemy.text("CREATE ROLE IF NOT EXISTS :role")
+            db_connection.execute(stmt, {"role": group})
 
     @async_wrap
     def grant_group_role(self, role, users):
@@ -108,9 +94,11 @@ def grant_group_role(self, role, users):
             role: Name of DB role to grant to users.
             users: List of DB users' usernames.
         """
-        stmt = sqlalchemy.text("GRANT :role TO :user")
-        for user in users:
-            self.db.execute(stmt, {"role": role, "user": user})
+        # create connection to db instance
+        with self.db.connect() as db_connection:
+            stmt = sqlalchemy.text("GRANT :role TO :user")
+            for user in users:
+                db_connection.execute(stmt, {"role": role, "user": user})
 
     @async_wrap
     def revoke_group_role(self, role, users):
@@ -123,9 +111,11 @@ def revoke_group_role(self, role, users):
             role: Name of DB role to revoke from users.
             users: List of DB users' usernames.
         """
-        stmt = sqlalchemy.text("REVOKE :role FROM :user")
-        for user in users:
-            self.db.execute(stmt, {"role": role, "user": user})
+        # create connection to db instance
+        with self.db.connect() as db_connection:
+            stmt = sqlalchemy.text("REVOKE :role FROM :user")
+            for user in users:
+                db_connection.execute(stmt, {"role": role, "user": user})
 
 
 def init_connection_engine(instance_connection_name, creds, ip_type=IPTypes.PUBLIC):

iam_groups_authn/sql_admin.py

@@ -14,10 +14,8 @@
 
 # sql_admin.py contains functions for interacting with the SQL Admin API
 
-from functools import partial
-from collections import defaultdict
-from quart.utils import run_sync
 from typing import NamedTuple
+from iam_groups_authn.sync import get_users_to_add
 
 
 class InstanceConnectionName(NamedTuple):
@@ -34,28 +32,47 @@ class InstanceConnectionName(NamedTuple):
     instance: str
 
 
-async def get_instance_users(user_service, instance_connection_names):
-    """Get users that belong to each Cloud SQL instance.
+async def get_instance_users(user_service, instance_connection_name):
+    """Get users that belong to a Cloud SQL instance.
 
-    Given a list of Cloud SQL instance names and a Google Cloud project, get a list
-    of database users that belong to each instance.
+    Given a Cloud SQL instance name and a Google Cloud project, get a list
+    of database users that belong to that instance.
 
     Args:
         user_service: A UserService object for calling SQL admin APIs.
-        instance_connection_names: List of Cloud SQL instance connection names.
-            (e.g., ["my-project:my-region:my-instance", "my-project:my-region:my-other-instance"])
+        instance_connection_name: Cloud SQL instance connection name.
+            (e.g., "my-project:my-region:my-instance")
 
     Returns:
-        db_users: A dict with the instance names mapping to their list of database users.
+        db_users: A list with the names of database users for the given instance.
     """
-    # create dict to hold database users of each instance
-    db_users = defaultdict(list)
-    for connection_name in instance_connection_names:
-        get_users = partial(
-            user_service.get_db_users,
-            InstanceConnectionName(*connection_name.split(":")),
-        )
-        users = await run_sync(get_users)()
-        for user in users:
-            db_users[connection_name].append(user["name"])
+    db_users = []
+    # get database users for instance
+    users = await user_service.get_db_users(
+        InstanceConnectionName(*instance_connection_name.split(":"))
+    )
+    for user in users:
+        db_users.append(user["name"])
     return db_users
+
+
+async def add_missing_db_users(
+    user_service, iam_future, db_future, instance_connection_name
+):
+    """Add missing IAM users as database users on instance.
+
+    Args:
+        user_service: A UserService object for calling SQL admin APIs.
+        iam_future: Future for list of IAM users who are members of IAM group.
+        db_future: Future for list of DB users on Cloud SQL database instance.
+        instance_connection_name: Cloud SQL instance connection name.
+            (e.g., "my-project:my-region:my-instance")
+    """
+    iam_users, db_users = await iam_future, await db_future
+    # find IAM users who are missing as DB users
+    missing_db_users = get_users_to_add(iam_users, db_users)
+    # add missing users to database instance
+    for user in missing_db_users:
+        user_service.insert_db_user(
+            user, InstanceConnectionName(*instance_connection_name.split(":"))
+        )