From d976ceb882e223e9ecdbe643fee3502a7933f031 Mon Sep 17 00:00:00 2001 From: Jaeyeon Lee Date: Tue, 29 Sep 2026 17:43:31 +0200 Subject: [PATCH 1/7] 197: team synchronization --- aqueduct/management/auth.py | 175 +++++++++--------- .../tests/test_oauth_team_creation.py | 52 +++++- 2 files changed, 133 insertions(+), 94 deletions(-) diff --git a/aqueduct/management/auth.py b/aqueduct/management/auth.py index 96b7b5f9..9d9dc885 100644 --- a/aqueduct/management/auth.py +++ b/aqueduct/management/auth.py @@ -110,6 +110,86 @@ def _update_user_role( user.is_superuser = is_superuser user.save() + @staticmethod + def _add_team_membership(profile: UserProfile, teams): + if not getattr(settings, "ENABLE_OAUTH_GROUP_CREATION", True): + log.info( + "Skipping creation of teams %s (ENABLE_OAUTH_GROUP_CREATION=False)", + sorted(team_name for team_name, _ in teams), + ) + return + + org = profile.org + for team_name, original_group_name in teams: + # Look up by oauth_group_name first, so renaming the mapping + # function renames existing teams instead of creating duplicates. + team = Team.objects.filter(oauth_group_name=original_group_name, org=org).first() + created = False + + if not team: + team, created = Team.objects.get_or_create( + name=team_name, org=org, defaults={"oauth_group_name": original_group_name} + ) + + if team and team.name != team_name: + collision = ( + Team.objects.filter(name=team_name, org=org).exclude(pk=team.pk).exists() + ) + + if collision: + log.warning( + "Cannot rename team '%s' -> '%s': name collision (org: %s, " + "oauth_group: '%s'). Reusing existing team as-is.", + team.name, + team_name, + org.name, + original_group_name, + ) + else: + log.info( + "Renaming team '%s' -> '%s' (org: %s, oauth_group: '%s')", + team.name, + team_name, + org.name, + original_group_name, + ) + team.name = team_name + team.save(update_fields=["name"]) + + if created: + log.info("Created team '%s' for org '%s'", team_name, org.name) + else: + log.info("Reused existing team '%s' for org '%s'", team_name, org.name) + + TeamMembership.objects.get_or_create(user_profile=profile, team=team) + log.info("Added user '%s' to team '%s' (%s)", profile.user.email, team_name, org.name) + + @staticmethod + def _remove_team_membership(profile: UserProfile, teams): + if not getattr(settings, "ENABLE_OAUTH_GROUP_REMOVAL", True): + log.info( + "Skipping removal from teams %s for user '%s' (ENABLE_OAUTH_GROUP_REMOVAL=False)", + sorted(teams), + profile.user.email, + ) + return + + org = profile.org + for team_name in teams: + try: + team = Team.objects.get(name=team_name, org=org) + TeamMembership.objects.filter(user_profile=profile, team=team).delete() + log.info( + "Removed user '%s' from team '%s' (%s)", profile.user.email, team_name, org.name + ) + except Team.DoesNotExist: + log.warning( + "Team '%s' not found for removal (org: %s, user: %s)", + team_name, + org.name, + profile.user.email, + ) + def _sync_team_membership(self, user: User, profile: UserProfile, team_names: list[str]): """ Synchronize team membership based on OAuth claims. @@ -122,6 +202,10 @@ def _sync_team_membership(self, user: User, profile: UserProfile, team_names: li - Respects org boundaries (teams must belong to user's org) """ if not getattr(settings, "ENABLE_OAUTH_GROUP_MANAGEMENT", False): + log.info( + "Skipping synchronization of teams %s (ENABLE_OAUTH_GROUP_MANAGEMENT=False)", + sorted(team_name for team_name, _ in team_names), + ) return if not team_names: @@ -129,12 +213,11 @@ def _sync_team_membership(self, user: User, profile: UserProfile, team_names: li team_mappings = self._get_teams(team_names=team_names) - org = profile.org with transaction.atomic(): existing_memberships = set( - TeamMembership.objects.filter(user_profile=profile).values_list( - "team__name", flat=True - ) + TeamMembership.objects.filter(user_profile=profile) + .exclude(team__oauth_group_name__in=["", None]) + .values_list("team__name", flat=True) ) target_team_names = {team_name for team_name, _ in team_mappings} @@ -144,88 +227,10 @@ def _sync_team_membership(self, user: User, profile: UserProfile, team_names: li (name, team_name_to_original[name]) for name in target_team_names - existing_memberships ] - teams_to_remove = existing_memberships - target_team_names + self._add_team_membership(profile=profile, teams=teams_to_add) - enable_creation = getattr(settings, "ENABLE_OAUTH_GROUP_CREATION", True) - enable_removal = getattr(settings, "ENABLE_OAUTH_GROUP_REMOVAL", True) - - for team_name, original_group_name in teams_to_add: - # Look up by oauth_group_name first, so renaming the mapping - # function renames existing teams instead of creating duplicates. - existing = Team.objects.filter( - oauth_group_name=original_group_name, org=org - ).first() - - if existing is not None: - if existing.name != team_name: - # Check for name collision before renaming - collision = ( - Team.objects.filter(name=team_name, org=org) - .exclude(pk=existing.pk) - .exists() - ) - if collision: - log.warning( - "Cannot rename team '%s' -> '%s': name collision (org: %s, " - "oauth_group: '%s'). Reusing existing team as-is.", - existing.name, - team_name, - org.name, - original_group_name, - ) - else: - log.info( - "Renaming team '%s' -> '%s' (org: %s, oauth_group: '%s')", - existing.name, - team_name, - org.name, - original_group_name, - ) - existing.name = team_name - existing.save(update_fields=["name"]) - team = existing - created = False - elif ( - not enable_creation - and not Team.objects.filter(name=team_name, org=org).exists() - ): - log.info("Skipping team '%s' (ENABLE_OAUTH_GROUP_CREATION=False)", team_name) - continue - else: - team, created = Team.objects.get_or_create( - name=team_name, org=org, defaults={"oauth_group_name": original_group_name} - ) - - if created: - log.info("Created team '%s' for org '%s'", team_name, org.name) - else: - log.info("Reused existing team '%s' for org '%s'", team_name, org.name) - - TeamMembership.objects.get_or_create(user_profile=profile, team=team) - log.info("Added user '%s' to team '%s' (%s)", user.email, team_name, org.name) - - for team_name in teams_to_remove: - try: - team = Team.objects.get(name=team_name, org=org) - is_oauth_managed = bool(team.oauth_group_name) - if is_oauth_managed or enable_removal: - TeamMembership.objects.filter(user_profile=profile, team=team).delete() - log.info( - "Removed user '%s' from team '%s' (%s)", user.email, team_name, org.name - ) - else: - log.info( - "Skipping removal from non-OAuth team '%s' for user '%s'", - team_name, - user.email, - ) - except Team.DoesNotExist: - log.warning( - "Team '%s' not found for removal (org: %s, user: %s)", - team_name, - org.name, - user.email, - ) + teams_to_remove = existing_memberships - target_team_names + self._remove_team_membership(profile=profile, teams=teams_to_remove) def create_user(self, claims: dict[str, Any]) -> User | None: org = self._org(claims) diff --git a/aqueduct/management/tests/test_oauth_team_creation.py b/aqueduct/management/tests/test_oauth_team_creation.py index a654a520..d5461236 100644 --- a/aqueduct/management/tests/test_oauth_team_creation.py +++ b/aqueduct/management/tests/test_oauth_team_creation.py @@ -244,8 +244,9 @@ def test_user_removed_from_teams_when_group_removed(self): self.assertEqual(memberships.count(), 1) self.assertEqual(memberships.first().team.name, "E123") - def test_user_removed_from_oauth_managed_teams_even_when_removal_disabled(self): - """Test that user is always removed from OAuth-managed teams.""" + @override_settings(ENABLE_OAUTH_GROUP_REMOVAL=False) + def test_user_not_removed_from_oauth_managed_teams_when_removal_disabled(self): + """Test that user stays in OAuth-managed teams when ENABLE_OAUTH_GROUP_REMOVAL=False.""" user = User.objects.create_user(username="testuser", email="test@example.com") user.groups.add(self.user_group) profile = UserProfile.objects.create(user=user, org=self.org) @@ -259,14 +260,30 @@ def test_user_removed_from_oauth_managed_teams_even_when_removal_disabled(self): self.assertTrue(team1.managed_by_oauth) self.assertTrue(team2.managed_by_oauth) - with override_settings(ENABLE_OAUTH_GROUP_REMOVAL=False): - backend = OIDCBackend() - updated_groups = {"email": "test@example.com", "groups": ["E123-Students"]} - sync_teams(backend, user, profile, updated_groups) + backend = OIDCBackend() + updated_groups = {"email": "test@example.com", "groups": ["E123-Students"]} + sync_teams(backend, user, profile, updated_groups) + + memberships = TeamMembership.objects.filter(user_profile=profile) + self.assertEqual(memberships.count(), 2) + + def test_user_removed_from_oauth_managed_teams(self): + """Test that user is removed from OAuth-managed teams when removal is enabled.""" + user = User.objects.create_user(username="testuser", email="test@example.com") + user.groups.add(self.user_group) + profile = UserProfile.objects.create(user=user, org=self.org) + + initial_groups = {"email": "test@example.com", "groups": ["E123-Students", "E456-Staff"]} + sync_teams(self.backend, user, profile, initial_groups) + + self.assertEqual(TeamMembership.objects.filter(user_profile=profile).count(), 2) + + updated_groups = {"email": "test@example.com", "groups": ["E123-Students"]} + sync_teams(self.backend, user, profile, updated_groups) - memberships = TeamMembership.objects.filter(user_profile=profile) - self.assertEqual(memberships.count(), 1) - self.assertEqual(memberships.first().team.name, "E123") + memberships = TeamMembership.objects.filter(user_profile=profile) + self.assertEqual(memberships.count(), 1) + self.assertEqual(memberships.first().team.name, "E123") @override_settings(ENABLE_OAUTH_GROUP_REMOVAL=False) def test_user_not_removed_from_non_oauth_teams_when_removal_disabled(self): @@ -294,6 +311,23 @@ def test_user_not_removed_from_non_oauth_teams_when_removal_disabled(self): self.assertIn("E123", team_names) self.assertIn("ManualTeam", team_names) + def test_user_not_removed_from_non_oauth_teams_when_removal_enabled(self): + """Regression: non-OAuth teams are never removed, even with removal enabled.""" + user = User.objects.create_user(username="testuser", email="test@example.com") + user.groups.add(self.user_group) + profile = UserProfile.objects.create(user=user, org=self.org) + + manual_team = Team.objects.create(name="ManualTeam", org=self.org, oauth_group_name="") + TeamMembership.objects.create(user_profile=profile, team=manual_team) + + # Default ENABLE_OAUTH_GROUP_REMOVAL=True; group no longer maps to the manual team. + updated_groups = {"email": "test@example.com", "groups": ["OtherGroup"]} + sync_teams(self.backend, user, profile, updated_groups) + + memberships = TeamMembership.objects.filter(user_profile=profile) + self.assertEqual(memberships.count(), 1) + self.assertEqual(memberships.first().team.name, "ManualTeam") + def test_membership_sync_on_update(self): """Test membership sync on update_user().""" claims_initial = {"email": "test@example.com", "groups": ["E123-Students"]} From d1e27fc1fd32946d6a82f41668970890858cdb19 Mon Sep 17 00:00:00 2001 From: Jaeyeon Lee Date: Tue, 29 Sep 2026 17:50:38 +0200 Subject: [PATCH 2/7] Add testcase --- .../tests/test_oauth_team_creation.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/aqueduct/management/tests/test_oauth_team_creation.py b/aqueduct/management/tests/test_oauth_team_creation.py index d5461236..75f22e5a 100644 --- a/aqueduct/management/tests/test_oauth_team_creation.py +++ b/aqueduct/management/tests/test_oauth_team_creation.py @@ -328,6 +328,27 @@ def test_user_not_removed_from_non_oauth_teams_when_removal_enabled(self): self.assertEqual(memberships.count(), 1) self.assertEqual(memberships.first().team.name, "ManualTeam") + def test_only_oauth_teams_managed_and_non_oauth_preserved(self): + """Only OAuth-managed teams are added/removed; non-OAuth memberships are preserved.""" + user = User.objects.create_user(username="testuser", email="test@example.com") + user.groups.add(self.user_group) + profile = UserProfile.objects.create(user=user, org=self.org) + + initial_groups = {"email": "test@example.com", "groups": ["E123-Students", "E456-Staff"]} + sync_teams(self.backend, user, profile, initial_groups) + self.assertEqual(TeamMembership.objects.filter(user_profile=profile).count(), 2) + + # Manually attach to a non-OAuth (manual) team; sync must never touch it. + manual_team = Team.objects.create(name="ManualTeam", org=self.org, oauth_group_name="") + TeamMembership.objects.create(user_profile=profile, team=manual_team) + + # E456-Staff is dropped, E123-Students kept; OtherGroup maps to no OAuth team. + updated_groups = {"email": "test@example.com", "groups": ["E123-Students", "OtherGroup"]} + sync_teams(self.backend, user, profile, updated_groups) + + team_names = {m.team.name for m in TeamMembership.objects.filter(user_profile=profile)} + self.assertEqual(team_names, {"E123", "ManualTeam"}) + def test_membership_sync_on_update(self): """Test membership sync on update_user().""" claims_initial = {"email": "test@example.com", "groups": ["E123-Students"]} From 8d4ed721cabf555e1aa9fbadc141ae0fbee4e386 Mon Sep 17 00:00:00 2001 From: Jaeyeon Lee Date: Mon, 5 Oct 2026 16:47:27 +0200 Subject: [PATCH 3/7] Edit creation logic --- aqueduct/management/auth.py | 11 ++++------- .../tests/test_oauth_team_creation.py | 19 +++++++++++++++++++ 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/aqueduct/management/auth.py b/aqueduct/management/auth.py index 9d9dc885..0ce394cc 100644 --- a/aqueduct/management/auth.py +++ b/aqueduct/management/auth.py @@ -112,13 +112,7 @@ def _update_user_role( @staticmethod def _add_team_membership(profile: UserProfile, teams): - if not getattr(settings, "ENABLE_OAUTH_GROUP_CREATION", True): - log.info( - "Skipping creation of teams %s (ENABLE_OAUTH_GROUP_CREATION=False)", - sorted(team_name for team_name, _ in teams), - ) - return - + enable_creation = getattr(settings, "ENABLE_OAUTH_GROUP_CREATION", True) org = profile.org for team_name, original_group_name in teams: # Look up by oauth_group_name first, so renaming the mapping @@ -127,6 +121,9 @@ def _add_team_membership(profile: UserProfile, teams): created = False if not team: + if not enable_creation: + log.info("Skipping team '%s' (ENABLE_OAUTH_GROUP_CREATION=False)", team_name) + continue team, created = Team.objects.get_or_create( name=team_name, org=org, defaults={"oauth_group_name": original_group_name} ) diff --git a/aqueduct/management/tests/test_oauth_team_creation.py b/aqueduct/management/tests/test_oauth_team_creation.py index 75f22e5a..e5eb16f5 100644 --- a/aqueduct/management/tests/test_oauth_team_creation.py +++ b/aqueduct/management/tests/test_oauth_team_creation.py @@ -539,6 +539,25 @@ def test_creation_disabled_flag(self): self.assertEqual(Team.objects.filter(org=self.org).count(), 0) self.assertEqual(TeamMembership.objects.filter(user_profile=profile).count(), 0) + def test_creation_disabled_adds_membership_to_existing_team(self): + """ENABLE_OAUTH_GROUP_CREATION=False still adds memberships to existing teams.""" + seed_active_config(SNIPPET_TEAM_NAMES_AND_MAP) + Team.objects.create(name="E123", org=self.org, oauth_group_name="E123-Students") + + with override_settings(ENABLE_OAUTH_GROUP_CREATION=False): + claims = {"email": "test@example.com", "groups": ["E123-Students", "E456-Staff"]} + + user = User.objects.create_user(username="testuser", email="test@example.com") + user.groups.add(self.user_group) + profile = UserProfile.objects.create(user=user, org=self.org) + + sync_teams(self.backend, user, profile, claims) + + # E456 is not created; E123 already exists and the user is added to it. + self.assertEqual(Team.objects.filter(org=self.org).count(), 1) + self.assertEqual(TeamMembership.objects.filter(user_profile=profile).count(), 1) + self.assertEqual(TeamMembership.objects.get(user_profile=profile).team.name, "E123") + def test_manual_team_not_affected_by_oauth_sync(self): """Test that manually created teams without oauth_group_name are not affected.""" manual_team = Team.objects.create(name="ManualTeam", org=self.org, oauth_group_name="") From 06046527bbebfab43d7ebaf1d66c521cb26cb4a8 Mon Sep 17 00:00:00 2001 From: Jaeyeon Lee Date: Fri, 9 Oct 2026 16:46:38 +0200 Subject: [PATCH 4/7] Refactor --- aqueduct/management/auth.py | 90 ++++++++++--------- .../tests/test_oauth_team_creation.py | 38 ++++++-- 2 files changed, 78 insertions(+), 50 deletions(-) diff --git a/aqueduct/management/auth.py b/aqueduct/management/auth.py index 0ce394cc..58847f5a 100644 --- a/aqueduct/management/auth.py +++ b/aqueduct/management/auth.py @@ -120,39 +120,39 @@ def _add_team_membership(profile: UserProfile, teams): team = Team.objects.filter(oauth_group_name=original_group_name, org=org).first() created = False - if not team: - if not enable_creation: - log.info("Skipping team '%s' (ENABLE_OAUTH_GROUP_CREATION=False)", team_name) - continue + if team is not None: + if team.name != team_name: + collision = ( + Team.objects.filter(name=team_name, org=org).exclude(pk=team.pk).exists() + ) + + if collision: + log.warning( + "Cannot rename team '%s' -> '%s': name collision (org: %s, " + "oauth_group: '%s'). Reusing existing team as-is.", + team.name, + team_name, + org.name, + original_group_name, + ) + else: + log.info( + "Renaming team '%s' -> '%s' (org: %s, oauth_group: '%s')", + team.name, + team_name, + org.name, + original_group_name, + ) + team.name = team_name + team.save(update_fields=["name"]) + elif not enable_creation and not Team.objects.filter(name=team_name, org=org).exists(): + log.info("Skipping team '%s' (ENABLE_OAUTH_GROUP_CREATION=False)", team_name) + continue + else: team, created = Team.objects.get_or_create( name=team_name, org=org, defaults={"oauth_group_name": original_group_name} ) - if team and team.name != team_name: - collision = ( - Team.objects.filter(name=team_name, org=org).exclude(pk=team.pk).exists() - ) - - if collision: - log.warning( - "Cannot rename team '%s' -> '%s': name collision (org: %s, " - "oauth_group: '%s'). Reusing existing team as-is.", - team.name, - team_name, - org.name, - original_group_name, - ) - else: - log.info( - "Renaming team '%s' -> '%s' (org: %s, oauth_group: '%s')", - team.name, - team_name, - org.name, - original_group_name, - ) - team.name = team_name - team.save(update_fields=["name"]) - if created: log.info("Created team '%s' for org '%s'", team_name, org.name) else: @@ -163,22 +163,28 @@ def _add_team_membership(profile: UserProfile, teams): @staticmethod def _remove_team_membership(profile: UserProfile, teams): - if not getattr(settings, "ENABLE_OAUTH_GROUP_REMOVAL", True): - log.info( - "Skipping removal from teams %s for user '%s' (ENABLE_OAUTH_GROUP_REMOVAL=False)", - sorted(teams), - profile.user.email, - ) - return - + enable_removal = getattr(settings, "ENABLE_OAUTH_GROUP_REMOVAL", True) org = profile.org + for team_name in teams: try: team = Team.objects.get(name=team_name, org=org) - TeamMembership.objects.filter(user_profile=profile, team=team).delete() - log.info( - "Removed user '%s' from team '%s' (%s)", profile.user.email, team_name, org.name - ) + is_oauth_managed = bool(team.oauth_group_name) + + if is_oauth_managed or enable_removal: + TeamMembership.objects.filter(user_profile=profile, team=team).delete() + log.info( + "Removed user '%s' from team '%s' (%s)", + profile.user.email, + team_name, + org.name, + ) + else: + log.info( + "Skipping removal from non-OAuth team '%s' for user '%s'", + team_name, + profile.user.email, + ) except Team.DoesNotExist: log.warning( "Team '%s' not found for removal (org: %s, user: %s)", @@ -201,7 +207,7 @@ def _sync_team_membership(self, user: User, profile: UserProfile, team_names: li if not getattr(settings, "ENABLE_OAUTH_GROUP_MANAGEMENT", False): log.info( "Skipping synchronization of teams %s (ENABLE_OAUTH_GROUP_MANAGEMENT=False)", - sorted(team_name for team_name, _ in team_names), + sorted(team_names), ) return diff --git a/aqueduct/management/tests/test_oauth_team_creation.py b/aqueduct/management/tests/test_oauth_team_creation.py index e5eb16f5..b2965b54 100644 --- a/aqueduct/management/tests/test_oauth_team_creation.py +++ b/aqueduct/management/tests/test_oauth_team_creation.py @@ -244,9 +244,8 @@ def test_user_removed_from_teams_when_group_removed(self): self.assertEqual(memberships.count(), 1) self.assertEqual(memberships.first().team.name, "E123") - @override_settings(ENABLE_OAUTH_GROUP_REMOVAL=False) - def test_user_not_removed_from_oauth_managed_teams_when_removal_disabled(self): - """Test that user stays in OAuth-managed teams when ENABLE_OAUTH_GROUP_REMOVAL=False.""" + def test_user_removed_from_oauth_managed_teams_even_when_removal_disabled(self): + """Test that user is always removed from OAuth-managed teams.""" user = User.objects.create_user(username="testuser", email="test@example.com") user.groups.add(self.user_group) profile = UserProfile.objects.create(user=user, org=self.org) @@ -260,12 +259,14 @@ def test_user_not_removed_from_oauth_managed_teams_when_removal_disabled(self): self.assertTrue(team1.managed_by_oauth) self.assertTrue(team2.managed_by_oauth) - backend = OIDCBackend() - updated_groups = {"email": "test@example.com", "groups": ["E123-Students"]} - sync_teams(backend, user, profile, updated_groups) + with override_settings(ENABLE_OAUTH_GROUP_REMOVAL=False): + backend = OIDCBackend() + updated_groups = {"email": "test@example.com", "groups": ["E123-Students"]} + sync_teams(backend, user, profile, updated_groups) - memberships = TeamMembership.objects.filter(user_profile=profile) - self.assertEqual(memberships.count(), 2) + memberships = TeamMembership.objects.filter(user_profile=profile) + self.assertEqual(memberships.count(), 1) + self.assertEqual(memberships.first().team.name, "E123") def test_user_removed_from_oauth_managed_teams(self): """Test that user is removed from OAuth-managed teams when removal is enabled.""" @@ -558,6 +559,27 @@ def test_creation_disabled_adds_membership_to_existing_team(self): self.assertEqual(TeamMembership.objects.filter(user_profile=profile).count(), 1) self.assertEqual(TeamMembership.objects.get(user_profile=profile).team.name, "E123") + def test_creation_disabled_reuses_manual_team_by_name(self): + """Creation disabled reuses an existing manual team matching the display name.""" + seed_active_config(SNIPPET_TEAM_NAMES_AND_MAP) + manual_team = Team.objects.create(name="E123", org=self.org, oauth_group_name="") + + with override_settings(ENABLE_OAUTH_GROUP_CREATION=False): + claims = {"email": "test@example.com", "groups": ["E123-Students"]} + + user = User.objects.create_user(username="testuser", email="test@example.com") + user.groups.add(self.user_group) + profile = UserProfile.objects.create(user=user, org=self.org) + + sync_teams(self.backend, user, profile, claims) + + # No new team created; the manual team is reused and the user is added to it. + self.assertEqual(Team.objects.filter(org=self.org).count(), 1) + memberships = TeamMembership.objects.filter(user_profile=profile) + self.assertEqual(memberships.count(), 1) + self.assertEqual(memberships.first().team, manual_team) + self.assertEqual(manual_team.oauth_group_name, "") + def test_manual_team_not_affected_by_oauth_sync(self): """Test that manually created teams without oauth_group_name are not affected.""" manual_team = Team.objects.create(name="ManualTeam", org=self.org, oauth_group_name="") From b8c4fec996887530f1f9415f439f550aeeae9b82 Mon Sep 17 00:00:00 2001 From: Jaeyeon Lee Date: Mon, 12 Oct 2026 16:11:16 +0200 Subject: [PATCH 5/7] Solve N+1 problem --- aqueduct/management/auth.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/aqueduct/management/auth.py b/aqueduct/management/auth.py index 58847f5a..19cd5317 100644 --- a/aqueduct/management/auth.py +++ b/aqueduct/management/auth.py @@ -162,13 +162,14 @@ def _add_team_membership(profile: UserProfile, teams): log.info("Added user '%s' to team '%s' (%s)", profile.user.email, team_name, org.name) @staticmethod - def _remove_team_membership(profile: UserProfile, teams): + def _remove_team_membership(profile: UserProfile, team_names): enable_removal = getattr(settings, "ENABLE_OAUTH_GROUP_REMOVAL", True) org = profile.org - for team_name in teams: + teams = Team.objects.filter(name__in=team_names, org=org) + + for team in teams: try: - team = Team.objects.get(name=team_name, org=org) is_oauth_managed = bool(team.oauth_group_name) if is_oauth_managed or enable_removal: @@ -176,19 +177,19 @@ def _remove_team_membership(profile: UserProfile, teams): log.info( "Removed user '%s' from team '%s' (%s)", profile.user.email, - team_name, + team.name, org.name, ) else: log.info( "Skipping removal from non-OAuth team '%s' for user '%s'", - team_name, + team.name, profile.user.email, ) except Team.DoesNotExist: log.warning( "Team '%s' not found for removal (org: %s, user: %s)", - team_name, + team.name, org.name, profile.user.email, ) @@ -233,7 +234,7 @@ def _sync_team_membership(self, user: User, profile: UserProfile, team_names: li self._add_team_membership(profile=profile, teams=teams_to_add) teams_to_remove = existing_memberships - target_team_names - self._remove_team_membership(profile=profile, teams=teams_to_remove) + self._remove_team_membership(profile=profile, team_names=teams_to_remove) def create_user(self, claims: dict[str, Any]) -> User | None: org = self._org(claims) From 945c1a7e29bf47e2ab2b3ab8fdbf939def5d315f Mon Sep 17 00:00:00 2001 From: Jaeyeon Lee Date: Mon, 12 Oct 2026 17:22:57 +0200 Subject: [PATCH 6/7] Refactoring --- aqueduct/db_test.sqlite3-journal | Bin 0 -> 189296 bytes aqueduct/management/auth.py | 127 ++++++++++++++++++++----------- 2 files changed, 82 insertions(+), 45 deletions(-) create mode 100644 aqueduct/db_test.sqlite3-journal diff --git a/aqueduct/db_test.sqlite3-journal b/aqueduct/db_test.sqlite3-journal new file mode 100644 index 0000000000000000000000000000000000000000..9888638e178e39f650f86438a92f0405d38d6afa GIT binary patch literal 189296 zcmeI*%WfNW6bE1@gVdyIJS8Dih)Pjgu|Or1QYf@s%Kf4Oq7MMUNruD_U$C7;pMh=I z@C2~y7V#1Y@c>nKzzeVgJDjN2U+j^NK5ydVvFthD$?B&w<5uh8iezwe*H)NSLRmrMt}eT0t5&UAV7cs0RjZ_3M^OW7CO7VZogWo zHWmh>z0Kol@l>_Gd^m1Kk^3yoasaVU-v0xg5g&+Vi7x%&J1je56UZU5c= zrTu67r}mF&5i~}C009C72oNAZfB*pk|Es{-Bh}@F&f#Q#YuJo>-SN0P=nvbMs!Pke zUv>I>gROCQ(yU*st}gXD{mx#~Yx}0- zCQYaJ=#$mc^RG=8=g+TIx8`4(y{*^uwwuv-zdLL%R~P<0x4qHea9CYDTisrY)7ohc zn$g{@?WQxCoq2b#+ufOT#$P|RTHX1t_w0-Y{YhtgytCiz9v(FH53PF6{mkiVW9e97 z*2DC`+T6j}({ojIZQP)7p zjDxES4#v45FkK+009C72oNAZpb&^0;1&A- z1&opa0RjXF5FkK+009C72qYAU9N+~xK*Hf_i2wlt1PBlyK!5-N0t5&Y0+9n;lLHhm zN&*B35FkK+009C72oNBUP#|)E=iMJ5;c&G?fB*pk1PBlyK!5-N0t5gkpm@5i z2oNAZfB*pk1PBmFC=fZo=YB6B;c&G?fB*pk1PBlyK!5-N0t5$&$&g&d%OQ4%0PfB*pk1PBly nK!5;&gaVNR+;&|+!r^L(009C72oNAZfB*pk1PBxYvmD?b4>}>E literal 0 HcmV?d00001 diff --git a/aqueduct/management/auth.py b/aqueduct/management/auth.py index 19cd5317..676884f6 100644 --- a/aqueduct/management/auth.py +++ b/aqueduct/management/auth.py @@ -111,54 +111,72 @@ def _update_user_role( user.save() @staticmethod - def _add_team_membership(profile: UserProfile, teams): - enable_creation = getattr(settings, "ENABLE_OAUTH_GROUP_CREATION", True) + def _update_team_membership(profile: UserProfile, teams, teams_by_name): org = profile.org - for team_name, original_group_name in teams: - # Look up by oauth_group_name first, so renaming the mapping - # function renames existing teams instead of creating duplicates. - team = Team.objects.filter(oauth_group_name=original_group_name, org=org).first() - created = False - - if team is not None: - if team.name != team_name: - collision = ( - Team.objects.filter(name=team_name, org=org).exclude(pk=team.pk).exists() + added = [] + + for team, team_name, original_group_name in teams: + if team.name != team_name: + existing = teams_by_name.get(team_name) + collision = existing is not None and existing.pk != team.pk + + if collision: + log.warning( + "Cannot rename team '%s' -> '%s': name collision (org: %s, " + "oauth_group: '%s'). Reusing existing team as-is.", + team.name, + team_name, + org.name, + original_group_name, + ) + else: + log.info( + "Renaming team '%s' -> '%s' (org: %s, oauth_group: '%s')", + team.name, + team_name, + org.name, + original_group_name, ) - if collision: - log.warning( - "Cannot rename team '%s' -> '%s': name collision (org: %s, " - "oauth_group: '%s'). Reusing existing team as-is.", - team.name, - team_name, - org.name, - original_group_name, - ) - else: - log.info( - "Renaming team '%s' -> '%s' (org: %s, oauth_group: '%s')", - team.name, - team_name, - org.name, - original_group_name, - ) - team.name = team_name - team.save(update_fields=["name"]) - elif not enable_creation and not Team.objects.filter(name=team_name, org=org).exists(): - log.info("Skipping team '%s' (ENABLE_OAUTH_GROUP_CREATION=False)", team_name) - continue - else: - team, created = Team.objects.get_or_create( - name=team_name, org=org, defaults={"oauth_group_name": original_group_name} - ) + old_name = team.name + team.name = team_name + team.save(update_fields=["name"]) + teams_by_name.pop(old_name, None) + teams_by_name[team_name] = team + + added.append(team) + + TeamMembership.objects.bulk_create( + [TeamMembership(user_profile=profile, team=team) for team in added], + ignore_conflicts=True, + ) + + for team in added: + log.info("Added user '%s' to team '%s' (%s)", profile.user.email, team.name, org.name) + + @staticmethod + def _add_team_membership(profile: UserProfile, teams): + """Create OAuth teams (or reuse an existing team by name) and ensure membership.""" + org = profile.org + added = [] + + for team_name, original_group_name in teams: + team, created = Team.objects.get_or_create( + name=team_name, org=org, defaults={"oauth_group_name": original_group_name} + ) if created: log.info("Created team '%s' for org '%s'", team_name, org.name) else: log.info("Reused existing team '%s' for org '%s'", team_name, org.name) + added.append((team_name, team)) - TeamMembership.objects.get_or_create(user_profile=profile, team=team) + TeamMembership.objects.bulk_create( + [TeamMembership(user_profile=profile, team=team) for _, team in added], + ignore_conflicts=True, + ) + + for team_name, _ in added: log.info("Added user '%s' to team '%s' (%s)", profile.user.email, team_name, org.name) @staticmethod @@ -216,6 +234,7 @@ def _sync_team_membership(self, user: User, profile: UserProfile, team_names: li return team_mappings = self._get_teams(team_names=team_names) + enable_creation = getattr(settings, "ENABLE_OAUTH_GROUP_CREATION", True) with transaction.atomic(): existing_memberships = set( @@ -225,13 +244,31 @@ def _sync_team_membership(self, user: User, profile: UserProfile, team_names: li ) target_team_names = {team_name for team_name, _ in team_mappings} - team_name_to_original = dict(team_mappings) - teams_to_add = [ - (name, team_name_to_original[name]) - for name in target_team_names - existing_memberships - ] - self._add_team_membership(profile=profile, teams=teams_to_add) + org_teams = list(Team.objects.filter(org=profile.org)) + teams_by_oauth = { + org_team.oauth_group_name: org_team + for org_team in org_teams + if org_team.oauth_group_name + } + teams_by_name = {org_team.name: org_team for org_team in org_teams} + + to_update, to_add = [], [] + for team_name, original_group_name in team_mappings: + if team_name in existing_memberships: + continue + team = teams_by_oauth.get(original_group_name) + if team is not None: + to_update.append((team, team_name, original_group_name)) + elif enable_creation or team_name in teams_by_name: + to_add.append((team_name, original_group_name)) + else: + log.info("Skipping team '%s' (ENABLE_OAUTH_GROUP_CREATION=False)", team_name) + + self._update_team_membership( + profile=profile, teams=to_update, teams_by_name=teams_by_name + ) + self._add_team_membership(profile=profile, teams=to_add) teams_to_remove = existing_memberships - target_team_names self._remove_team_membership(profile=profile, team_names=teams_to_remove) From 8a1b6e33e2d6581dd53d0f11cfa20eed7b62e179 Mon Sep 17 00:00:00 2001 From: Jaeyeon Lee Date: Mon, 12 Oct 2026 17:55:18 +0200 Subject: [PATCH 7/7] Edit --- aqueduct/db_test.sqlite3-journal | Bin 189296 -> 0 bytes aqueduct/management/auth.py | 50 ++++++++++++++----------------- 2 files changed, 22 insertions(+), 28 deletions(-) delete mode 100644 aqueduct/db_test.sqlite3-journal diff --git a/aqueduct/db_test.sqlite3-journal b/aqueduct/db_test.sqlite3-journal deleted file mode 100644 index 9888638e178e39f650f86438a92f0405d38d6afa..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 189296 zcmeI*%WfNW6bE1@gVdyIJS8Dih)Pjgu|Or1QYf@s%Kf4Oq7MMUNruD_U$C7;pMh=I z@C2~y7V#1Y@c>nKzzeVgJDjN2U+j^NK5ydVvFthD$?B&w<5uh8iezwe*H)NSLRmrMt}eT0t5&UAV7cs0RjZ_3M^OW7CO7VZogWo zHWmh>z0Kol@l>_Gd^m1Kk^3yoasaVU-v0xg5g&+Vi7x%&J1je56UZU5c= zrTu67r}mF&5i~}C009C72oNAZfB*pk|Es{-Bh}@F&f#Q#YuJo>-SN0P=nvbMs!Pke zUv>I>gROCQ(yU*st}gXD{mx#~Yx}0- zCQYaJ=#$mc^RG=8=g+TIx8`4(y{*^uwwuv-zdLL%R~P<0x4qHea9CYDTisrY)7ohc zn$g{@?WQxCoq2b#+ufOT#$P|RTHX1t_w0-Y{YhtgytCiz9v(FH53PF6{mkiVW9e97 z*2DC`+T6j}({ojIZQP)7p zjDxES4#v45FkK+009C72oNAZpb&^0;1&A- z1&opa0RjXF5FkK+009C72qYAU9N+~xK*Hf_i2wlt1PBlyK!5-N0t5&Y0+9n;lLHhm zN&*B35FkK+009C72oNBUP#|)E=iMJ5;c&G?fB*pk1PBlyK!5-N0t5gkpm@5i z2oNAZfB*pk1PBmFC=fZo=YB6B;c&G?fB*pk1PBlyK!5-N0t5$&$&g&d%OQ4%0PfB*pk1PBly nK!5;&gaVNR+;&|+!r^L(009C72oNAZfB*pk1PBxYvmD?b4>}>E diff --git a/aqueduct/management/auth.py b/aqueduct/management/auth.py index 676884f6..def46b6a 100644 --- a/aqueduct/management/auth.py +++ b/aqueduct/management/auth.py @@ -111,11 +111,16 @@ def _update_user_role( user.save() @staticmethod - def _update_team_membership(profile: UserProfile, teams, teams_by_name): + def _update_team_membership( + profile: UserProfile, + teams_to_update: list[tuple[Team, str, str]], + teams_by_name: dict[str, Team], + ): + """Rename OAuth teams to match current group names and ensure membership.""" org = profile.org added = [] - for team, team_name, original_group_name in teams: + for team, team_name, original_group_name in teams_to_update: if team.name != team_name: existing = teams_by_name.get(team_name) collision = existing is not None and existing.pk != team.pk @@ -156,12 +161,12 @@ def _update_team_membership(profile: UserProfile, teams, teams_by_name): log.info("Added user '%s' to team '%s' (%s)", profile.user.email, team.name, org.name) @staticmethod - def _add_team_membership(profile: UserProfile, teams): + def _add_team_membership(profile: UserProfile, teams_to_add: list[tuple[str, str]]): """Create OAuth teams (or reuse an existing team by name) and ensure membership.""" org = profile.org added = [] - for team_name, original_group_name in teams: + for team_name, original_group_name in teams_to_add: team, created = Team.objects.get_or_create( name=team_name, org=org, defaults={"oauth_group_name": original_group_name} ) @@ -180,35 +185,24 @@ def _add_team_membership(profile: UserProfile, teams): log.info("Added user '%s' to team '%s' (%s)", profile.user.email, team_name, org.name) @staticmethod - def _remove_team_membership(profile: UserProfile, team_names): + def _remove_team_membership(profile: UserProfile, team_names: set[str]): + """Remove the user from teams no longer present in their OAuth groups.""" enable_removal = getattr(settings, "ENABLE_OAUTH_GROUP_REMOVAL", True) org = profile.org teams = Team.objects.filter(name__in=team_names, org=org) - for team in teams: - try: - is_oauth_managed = bool(team.oauth_group_name) + is_oauth_managed = bool(team.oauth_group_name) - if is_oauth_managed or enable_removal: - TeamMembership.objects.filter(user_profile=profile, team=team).delete() - log.info( - "Removed user '%s' from team '%s' (%s)", - profile.user.email, - team.name, - org.name, - ) - else: - log.info( - "Skipping removal from non-OAuth team '%s' for user '%s'", - team.name, - profile.user.email, - ) - except Team.DoesNotExist: - log.warning( - "Team '%s' not found for removal (org: %s, user: %s)", + if is_oauth_managed or enable_removal: + TeamMembership.objects.filter(user_profile=profile, team=team).delete() + log.info( + "Removed user '%s' from team '%s' (%s)", profile.user.email, team.name, org.name + ) + else: + log.info( + "Skipping removal from non-OAuth team '%s' for user '%s'", team.name, - org.name, profile.user.email, ) @@ -266,9 +260,9 @@ def _sync_team_membership(self, user: User, profile: UserProfile, team_names: li log.info("Skipping team '%s' (ENABLE_OAUTH_GROUP_CREATION=False)", team_name) self._update_team_membership( - profile=profile, teams=to_update, teams_by_name=teams_by_name + profile=profile, teams_to_update=to_update, teams_by_name=teams_by_name ) - self._add_team_membership(profile=profile, teams=to_add) + self._add_team_membership(profile=profile, teams_to_add=to_add) teams_to_remove = existing_memberships - target_team_names self._remove_team_membership(profile=profile, team_names=teams_to_remove)