Fixed edge case for RegexSubStrategy (#1561)

Author: nkvuong

src/databricks/labs/ucx/workspace_access/groups.py

@@ -205,7 +205,9 @@ def generate_migrated_groups(self):
             temporary_name = f"{self.renamed_groups_prefix}{group.display_name}"
             account_group = self.account_groups_in_account.get(group.display_name)
             if not account_group:
-                logger.info(f"Couldn't find a matching account group for {group.display_name} group")
+                logger.info(
+                    f"Couldn't find a matching account group for {group.display_name} group using name matching"
+                )
                 continue
             yield MigratedGroup(
                 id_in_workspace=group.id,
@@ -241,21 +243,19 @@ def generate_migrated_groups(self):
         for group in workspace_groups.values():
             temporary_name = f"{self.renamed_groups_prefix}{group.display_name}"
             account_group = account_groups_by_id.get(group.external_id)
-            if account_group:
-                yield MigratedGroup(
-                    id_in_workspace=group.id,
-                    name_in_workspace=group.display_name,
-                    name_in_account=account_group.display_name,
-                    temporary_name=temporary_name,
-                    external_id=account_group.external_id,
-                    members=json.dumps([gg.as_dict() for gg in group.members]) if group.members else None,
-                    roles=json.dumps([gg.as_dict() for gg in group.roles]) if group.roles else None,
-                    entitlements=(
-                        json.dumps([gg.as_dict() for gg in group.entitlements]) if group.entitlements else None
-                    ),
-                )
-            else:
+            if not account_group:
                 logger.info(f"Couldn't find a matching account group for {group.display_name} group with external_id")
+                continue
+            yield MigratedGroup(
+                id_in_workspace=group.id,
+                name_in_workspace=group.display_name,
+                name_in_account=account_group.display_name,
+                temporary_name=temporary_name,
+                external_id=account_group.external_id,
+                members=json.dumps([gg.as_dict() for gg in group.members]) if group.members else None,
+                roles=json.dumps([gg.as_dict() for gg in group.roles]) if group.roles else None,
+                entitlements=(json.dumps([gg.as_dict() for gg in group.entitlements]) if group.entitlements else None),
+            )
 
 
 class RegexSubStrategy(GroupMigrationStrategy):
@@ -285,12 +285,18 @@ def generate_migrated_groups(self):
             name_in_account = self._safe_sub(
                 group.display_name, self.workspace_group_regex, self.workspace_group_replace
             )
+            account_group = self.account_groups_in_account.get(name_in_account)
+            if not account_group:
+                logger.info(
+                    f"Couldn't find a matching account group for {group.display_name} group with regex substitution"
+                )
+                continue
             yield MigratedGroup(
                 id_in_workspace=group.id,
                 name_in_workspace=group.display_name,
                 name_in_account=name_in_account,
                 temporary_name=temporary_name,
-                external_id=self.account_groups_in_account[name_in_account].external_id,
+                external_id=account_group.external_id,
                 members=json.dumps([gg.as_dict() for gg in group.members]) if group.members else None,
                 roles=json.dumps([gg.as_dict() for gg in group.roles]) if group.roles else None,
                 entitlements=json.dumps([gg.as_dict() for gg in group.entitlements]) if group.entitlements else None,
@@ -329,21 +335,23 @@ def generate_migrated_groups(self):
         for group_match, ws_group in workspace_groups_by_match.items():
             temporary_name = f"{self.renamed_groups_prefix}{ws_group.display_name}"
             account_group = account_groups_by_match.get(group_match)
-            if account_group:
-                yield MigratedGroup(
-                    id_in_workspace=ws_group.id,
-                    name_in_workspace=ws_group.display_name,
-                    name_in_account=account_group.display_name,
-                    temporary_name=temporary_name,
-                    external_id=account_group.external_id,
-                    members=json.dumps([gg.as_dict() for gg in ws_group.members]) if ws_group.members else None,
-                    roles=json.dumps([gg.as_dict() for gg in ws_group.roles]) if ws_group.roles else None,
-                    entitlements=(
-                        json.dumps([gg.as_dict() for gg in ws_group.entitlements]) if ws_group.entitlements else None
-                    ),
+            if not account_group:
+                logger.info(
+                    f"Couldn't find a matching account group for {ws_group.display_name} group with regex matching"
                 )
-            else:
-                logger.info(f"Couldn't find a match for group {ws_group.display_name}")
+                continue
+            yield MigratedGroup(
+                id_in_workspace=ws_group.id,
+                name_in_workspace=ws_group.display_name,
+                name_in_account=account_group.display_name,
+                temporary_name=temporary_name,
+                external_id=account_group.external_id,
+                members=json.dumps([gg.as_dict() for gg in ws_group.members]) if ws_group.members else None,
+                roles=json.dumps([gg.as_dict() for gg in ws_group.roles]) if ws_group.roles else None,
+                entitlements=(
+                    json.dumps([gg.as_dict() for gg in ws_group.entitlements]) if ws_group.entitlements else None
+                ),
+            )
 
 
 class GroupManager(CrawlerBase[MigratedGroup]):

tests/unit/workspace_access/test_groups.py

@@ -1,4 +1,5 @@
 import json
+import logging
 from unittest.mock import create_autospec
 
 import pytest
@@ -636,6 +637,26 @@ def test_snapshot_with_group_matched_by_prefix():
     ]
 
 
+def test_snapshot_with_group_matched_by_prefix_not_found(caplog):
+    caplog.set_level(logging.INFO)
+    backend = MockBackend()
+    wsclient = create_autospec(WorkspaceClient)
+    group = Group(
+        id="1",
+        display_name="de_(1234)",
+        meta=ResourceMeta(resource_type="WorkspaceGroup"),
+    )
+    wsclient.groups.list.return_value = [group]
+    wsclient.groups.get.return_value = group
+    wsclient.api_client.do.return_value = {
+        "Resources": [],
+    }
+    GroupManager(
+        backend, wsclient, inventory_database="inv", workspace_group_regex="^", workspace_group_replace="px_"
+    ).snapshot()
+    assert "Couldn't find a matching account group for de_(1234) group with regex substitution" in caplog.text
+
+
 def test_snapshot_with_group_matched_by_subset():
     backend = MockBackend()
     wsclient = create_autospec(WorkspaceClient)
@@ -675,6 +696,26 @@ def test_snapshot_with_group_matched_by_subset():
     ]
 
 
+def test_snapshot_with_group_matched_by_subset_not_found(caplog):
+    caplog.set_level(logging.INFO)
+    backend = MockBackend()
+    wsclient = create_autospec(WorkspaceClient)
+    group = Group(
+        id="1",
+        display_name="de_(1234)",
+        meta=ResourceMeta(resource_type="WorkspaceGroup"),
+    )
+    wsclient.groups.list.return_value = [group]
+    wsclient.groups.get.return_value = group
+    wsclient.api_client.do.return_value = {
+        "Resources": [],
+    }
+    GroupManager(
+        backend, wsclient, inventory_database="inv", workspace_group_regex=r"\(([1-9]+)\)", account_group_regex="[1-9]+"
+    ).snapshot()
+    assert "Couldn't find a matching account group for de_(1234) group with regex matching" in caplog.text
+
+
 def test_snapshot_with_group_matched_by_external_id():
     backend = MockBackend()
     wsclient = create_autospec(WorkspaceClient)
@@ -712,6 +753,24 @@ def test_snapshot_with_group_matched_by_external_id():
     ]
 
 
+def test_snapshot_with_group_matched_by_external_id_not_found(caplog):
+    caplog.set_level(logging.INFO)
+    backend = MockBackend()
+    wsclient = create_autospec(WorkspaceClient)
+    group = Group(
+        id="1",
+        display_name="de_(1234)",
+        meta=ResourceMeta(resource_type="WorkspaceGroup"),
+    )
+    wsclient.groups.list.return_value = [group]
+    wsclient.groups.get.return_value = group
+    wsclient.api_client.do.return_value = {
+        "Resources": [],
+    }
+    GroupManager(backend, wsclient, inventory_database="inv", external_id_match=True).snapshot()
+    assert "Couldn't find a matching account group for de_(1234) group with external_id" in caplog.text
+
+
 def test_configure_include_groups():
     configure_groups = ConfigureGroups(
         MockPrompts(