Enable joining an existing installation to a collection (#1799)

HariGS-DB · web-flow · commit 2193127903ab · 2024-06-01T11:34:25.000+02:00
diff --git a/labs.yml b/labs.yml
@@ -45,7 +45,7 @@ commands:
 
   - name: report-account-compatibility
     is_account_level: true
-    description: aggregation of UCX output of multiple workspaces in the account. 
+    description: aggregation of UCX output of multiple workspaces in the account.
       If --workspace-ids is not provided, it will use all workspaces present in the account.
     flags:
       - name: workspace-ids
@@ -95,7 +95,7 @@ commands:
 
   - name: alias
     description: |
-      alias tables across schema/catalog withing a UC metastore 
+      alias tables across schema/catalog withing a UC metastore
       create a view pointing to the "from" table
       if a view is aliased, recreates the same view in the target schema/catalog
     flags:
@@ -111,9 +111,9 @@ commands:
         description: target schema to migrate tables to
 
   - name: principal-prefix-access
-    description: For azure cloud, identifies all storage account used by tables in the workspace, identify spn and its 
-      permission on each storage accounts. For aws, identifies all the Instance Profiles configured in the workspace and 
-      its access to all the S3 buckets, along with AWS roles that are set with UC access and its access to S3 buckets. 
+    description: For azure cloud, identifies all storage account used by tables in the workspace, identify spn and its
+      permission on each storage accounts. For aws, identifies all the Instance Profiles configured in the workspace and
+      its access to all the S3 buckets, along with AWS roles that are set with UC access and its access to S3 buckets.
       The output is stored in the workspace install folder.
     flags:
       - name: subscription-id
@@ -122,8 +122,8 @@ commands:
         description: AWS Profile to use for authentication
 
   - name: create-missing-principals
-    description: For AWS, this command identifies all the S3 locations that are missing a UC compatible role and 
-      creates them. It accepts a number of optional parameters, i.e. KMS Key, Role Name, Policy Name, and whether to 
+    description: For AWS, this command identifies all the S3 locations that are missing a UC compatible role and
+      creates them. It accepts a number of optional parameters, i.e. KMS Key, Role Name, Policy Name, and whether to
       create a single role for all the S3 locations.
     flags:
       - name: aws-profile
@@ -138,7 +138,7 @@ commands:
         description: (Optional) Create a single role for all S3 locations. (default:False)
 
   - name: create-uber-principal
-    description: For azure cloud, creates a service principal and gives STORAGE BLOB READER access on all the storage account 
+    description: For azure cloud, creates a service principal and gives STORAGE BLOB READER access on all the storage account
       used by tables in the workspace and stores the spn info in the UCX cluster policy. For aws,
       it identifies all s3 buckets used by the Instance Profiles configured in the workspace.
     flags:
@@ -215,7 +215,7 @@ commands:
 
   - name: migrate-tables
     description: |
-      Trigger the migrate-tables workflow and, optionally, migrate-external-hiveserde-tables-in-place-experimental 
+      Trigger the migrate-tables workflow and, optionally, migrate-external-hiveserde-tables-in-place-experimental
       workflow and migrate-external-tables-ctas workflow.
 
   - name: migrate-dbsql-dashboards
@@ -229,3 +229,12 @@ commands:
     flags:
       - name: dashboard-id
         description: (Optional) DBSQL dashboard ID to revert. If no dashboard ID is provided, all migrated DBSQL dashboards in the workspace will be reverted.
+
+  - name: join-collection
+    is_account_level: true
+    description: workspace_id to join a collection.
+    flags:
+      - name: workspace-id
+        description: workspace_id which should join a collection.
+      - name: target_workspace-id
+        description: (Optional) id of a workspace in the target collection. If not specified, ucx will prompt to select from a list
diff --git a/src/databricks/labs/ucx/account/workspaces.py b/src/databricks/labs/ucx/account/workspaces.py
@@ -93,11 +93,11 @@ def get_accessible_workspaces(self) -> list[Workspace]:
         """
         accessible_workspaces = []
         for workspace in self._ac.workspaces.list():
-            if self._can_administer(workspace):
+            if self.can_administer(workspace):
                 accessible_workspaces.append(workspace)
         return accessible_workspaces
 
-    def _can_administer(self, workspace: Workspace) -> bool:
+    def can_administer(self, workspace: Workspace) -> bool:
         try:
             # check if user has access to workspace
             ws = self.client_for(workspace)
@@ -252,7 +252,7 @@ def __init__(self, installation: Installation, ws: WorkspaceClient):
         self._installation = installation
         self._ws = ws
 
-    def _load_workspace_info(self) -> dict[int, Workspace]:
+    def load_workspace_info(self) -> dict[int, Workspace]:
         try:
             id_to_workspace = {}
             for workspace in self._installation.load(list[Workspace], filename=AccountWorkspaces.SYNC_FILE_NAME):
@@ -265,7 +265,7 @@ def _load_workspace_info(self) -> dict[int, Workspace]:
 
     def current(self) -> str:
         workspace_id = self._ws.get_workspace_id()
-        workspaces = self._load_workspace_info()
+        workspaces = self.load_workspace_info()
         if workspace_id not in workspaces:
             msg = f"Current workspace is not known: {workspace_id}"
             raise KeyError(msg) from None
diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py
@@ -13,6 +13,7 @@
 from databricks.labs.ucx.contexts.account_cli import AccountContext
 from databricks.labs.ucx.contexts.workspace_cli import WorkspaceContext, LocalCheckoutContext
 from databricks.labs.ucx.hive_metastore.tables import What
+from databricks.labs.ucx.install import AccountInstaller
 
 ucx = App(__file__)
 logger = get_logger(__file__)
@@ -471,6 +472,13 @@ def revert_dbsql_dashboards(w: WorkspaceClient, dashboard_id: str | None = None)
     ctx.redash.revert_dashboards(dashboard_id)
 
 
+@ucx.command(is_account=True)
+def join_collection(a: AccountClient, workspace_id: int, join_workspace_id: int):
+    """joins the workspace to an existing collection"""
+    account_installer = AccountInstaller(a)
+    account_installer.join_collection(workspace_id, join_workspace_id)
+
+
 @ucx.command
 def lint_local_code(
     w: WorkspaceClient, prompts: Prompts, path: str | None = None, ctx: LocalCheckoutContext | None = None
diff --git a/src/databricks/labs/ucx/install.py b/src/databricks/labs/ucx/install.py
@@ -40,6 +40,7 @@
 )
 
 from databricks.labs.ucx.__about__ import __version__
+from databricks.labs.ucx.account.workspaces import WorkspaceInfo
 from databricks.labs.ucx.assessment.azure import AzureServicePrincipalInfo
 from databricks.labs.ucx.assessment.clusters import ClusterInfo, PolicyInfo
 from databricks.labs.ucx.assessment.init_scripts import GlobalInitScriptInfo
@@ -663,34 +664,94 @@ def install_on_account(self):
         # upload the json dump of workspace info in the .ucx folder
         ctx.account_workspaces.sync_workspace_info(installed_workspaces)
 
+    @staticmethod
+    def _get_workspace(workspace_id: int, ids_to_workspace: dict[int, Workspace]) -> Workspace:
+        try:
+            workspace = ids_to_workspace[workspace_id]
+            return workspace
+        except KeyError:
+            msg = "Current workspace is not known, Please run as account-admin: databricks labs ucx sync-workspace-info"
+            raise KeyError(msg) from None
+
+    def _get_workspace_info(self, current_workspace_id: int):
+        account_client = self._get_safe_account_client()
+        workspace = account_client.workspaces.get(current_workspace_id)
+        current_workspace_client = account_client.get_workspace_client(workspace)
+        installation = Installation.current(current_workspace_client, self.product_info.product_name())
+        workspace_info = WorkspaceInfo(installation, current_workspace_client)
+        return workspace_info.load_workspace_info()
+
     def join_collection(
         self,
         current_workspace_id: int,
+        target_workspace_id: int | None = None,
     ):
-        if not self.is_account_install and self.prompts.confirm(
-            "Do you want to join the current installation to an existing collection?"
-        ):
-
-            installed_workspaces: list[Workspace] | None = []
-            accessible_workspaces: list[Workspace] = []
-            account_client = self._get_safe_account_client()
-            ctx = AccountContext(account_client)
-            try:
-                accessible_workspaces = ctx.account_workspaces.get_accessible_workspaces()
-            except PermissionDenied:
-                logger.warning("User doesnt have account admin permission, cant join a collection, skipping...")
-            collection_workspace = self._get_collection_workspace(accessible_workspaces, account_client)
-            if collection_workspace is not None:
-                installed_workspaces = self._sync_collection(collection_workspace, current_workspace_id, account_client)
-            if installed_workspaces is not None:
-                ctx.account_workspaces.sync_workspace_info(installed_workspaces)
+        if self.is_account_install:
+            # skip joining collection when installer is running for all account workspaces
+            return None
+        collection_workspace: Workspace
+        account_client = self._get_safe_account_client()
+        ctx = AccountContext(account_client)
+        ids_to_workspace = self._get_workspace_info(current_workspace_id)
+        if target_workspace_id is None:
+            if self.prompts.confirm("Do you want to join the current installation to an existing collection?"):
+                # If joining a collection as part of the installation then collection_workspace_id would be empty
+                try:
+                    # if user is account admin list and show available workspaces to select from
+                    accessible_workspaces = ctx.account_workspaces.get_accessible_workspaces()
+                    target_workspace = self._get_collection_workspace(
+                        accessible_workspaces,
+                        account_client,
+                    )
+                    assert target_workspace is not None
+                    target_workspace_id = target_workspace.workspace_id
+
+                except PermissionDenied:
+                    # if the user is not account admin, allow user to enter the workspace_id to join as collection.
+                    # if no workspace_id is entered, then exit
+                    logger.warning("User doesnt have account admin permission, cant list workspaces")
+                    target_workspace_id = int(
+                        self.prompts.question(
+                            "Please enter, the workspace id to join as a collection (enter 0 to skip it)",
+                            valid_number=True,
+                            default="0",
+                        )
+                    )
+            else:
+                return None
+        if target_workspace_id == 0 or target_workspace_id is None:
+            # if user didn't enter workspace id
+            logger.info("Skipping joining collection...")
+            return None
+        # below code is executed if either joining an existing collection (through the cli)
+        # or selecting one while installing
+        collection_workspace = AccountInstaller._get_workspace(
+            target_workspace_id,
+            ids_to_workspace,
+        )
+        if not ctx.account_workspaces.can_administer(collection_workspace):
+            # if user is not workspace admin on the workspace to join as collection then exit
+            logger.error(
+                f"User doesnt have admin access on the workspace {target_workspace_id}, " f"cant join collection."
+            )
+            return None
+        if collection_workspace is not None:
+            self._sync_collection(
+                collection_workspace,
+                current_workspace_id,
+                ids_to_workspace,
+            )
+        return None
 
     def _sync_collection(
         self,
         collection_workspace: Workspace,
         current_workspace_id: int,
-        account_client: AccountClient,
+        ids_to_workspace: dict[int, Workspace],
     ) -> list[Workspace] | None:
+        # gets the list of existing collection of workspace from the config
+        # checks if user is workspace admin on all the workspace id, if yes
+        # then joins the collection and syncs the config with all the affected workspaces
         installer = self._get_installer(collection_workspace)
         installed_workspace_ids = installer.config.installed_workspace_ids
         if installed_workspace_ids is None:
@@ -701,9 +762,20 @@ def _sync_collection(
             )
         installed_workspace_ids.append(current_workspace_id)
         installed_workspaces = []
-        for account_workspace in account_client.workspaces.list():
-            if account_workspace.workspace_id in installed_workspace_ids:
-                installed_workspaces.append(account_workspace)
+        ctx = AccountContext(self._get_safe_account_client())
+        for workspace_id in installed_workspace_ids:
+            workspace = AccountInstaller._get_workspace(
+                workspace_id,
+                ids_to_workspace,
+            )
+            installed_workspaces.append(workspace)
+            if not ctx.account_workspaces.can_administer(workspace):
+
+                logger.error(
+                    f"User doesnt have admin access on the workspace {workspace_id} in the collection, "
+                    f"cant join collection."
+                )
+                return None
 
         for installed_workspace in installed_workspaces:
             installer = self._get_installer(installed_workspace)
@@ -715,10 +787,15 @@ def _get_collection_workspace(
         accessible_workspaces: list[Workspace],
         account_client: AccountClient,
     ) -> Workspace | None:
+        # iterate through each workspace and select workspace which has existing ucx installation
+        # allow user to select from the eligible workspace to join as collection
         installed_workspaces = []
         for workspace in accessible_workspaces:
             workspace_client = account_client.get_workspace_client(workspace)
-            workspace_installation = Installation.existing(workspace_client, self.product_info.product_name())
+            workspace_installation = Installation.existing(
+                workspace_client,
+                self.product_info.product_name(),
+            )
             if len(workspace_installation) > 0:
                 installed_workspaces.append(workspace)
 
@@ -730,12 +807,12 @@ def _get_collection_workspace(
             for workspace in installed_workspaces
             if workspace.deployment_name is not None
         }
-        workspace = self.prompts.choice_from_dict(
+        target_workspace = self.prompts.choice_from_dict(
             "Please select a workspace, the current installation of ucx will be grouped as a "
             "collection with the selected workspace",
             workspaces,
         )
-        return workspace
+        return target_workspace
 
 
 if __name__ == "__main__":
@@ -749,5 +826,8 @@ def _get_collection_workspace(
         account_installer.install_on_account()
     else:
         workspace_installer = WorkspaceInstaller(WorkspaceClient(product="ucx", product_version=__version__))
+
         workspace_installer.run()
-        account_installer.join_collection(workspace_installer.workspace_client.get_workspace_id())
+        account_installer.join_collection(
+            workspace_installer.workspace_client.get_workspace_id(),
+        )
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -19,6 +19,7 @@
 from databricks.sdk.service.iam import Group
 
 from databricks.labs.ucx.__about__ import __version__
+from databricks.labs.ucx.account.workspaces import AccountWorkspaces
 from databricks.labs.ucx.assessment.aws import AWSRoleAction, run_command
 from databricks.labs.ucx.assessment.azure import (
     AzureServicePrincipalCrawler,
@@ -273,6 +274,9 @@ def __init__(
     def with_table_mapping_rules(self, rules):
         self.installation.save(rules, filename=TableMapping.FILENAME)
 
+    def with_workspace_info(self, workspace_info):
+        self.installation.save(workspace_info, filename=AccountWorkspaces.SYNC_FILE_NAME)
+
     def make_schema(self, **kwargs):
         schema_info = self._make_schema(**kwargs)
         self._schemas.append(schema_info)
@@ -764,18 +768,18 @@ def prepare_hiveserde_tables(context, random, schema, table_base_dir) -> dict[st
     )
 
     avro_table_name = f"avro_serde_{random}"
-    avro_ddl = f"""CREATE TABLE hive_metastore.{schema.name}.{avro_table_name} (id INT, region STRING) 
+    avro_ddl = f"""CREATE TABLE hive_metastore.{schema.name}.{avro_table_name} (id INT, region STRING)
                         ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.avro.AvroSerDe'
                         STORED AS INPUTFORMAT 'org.apache.hadoop.hive.ql.io.avro.AvroContainerInputFormat'
                                   OUTPUTFORMAT 'org.apache.hadoop.hive.ql.io.avro.AvroContainerOutputFormat'
                         TBLPROPERTIES ('avro.schema.literal'='{{
-                            "namespace": "org.apache.hive", 
-                            "name": "first_schema", 
+                            "namespace": "org.apache.hive",
+                            "name": "first_schema",
                             "type": "record",
                             "fields": [
                                 {{ "name":"id", "type":"int" }},
                                 {{ "name":"region", "type":"string" }}
-                            ] }}') 
+                            ] }}')
                         LOCATION '{table_base_dir}/{avro_table_name}'
                     """
     tables[avro_table_name] = context.make_table(
diff --git a/tests/integration/test_installation.py b/tests/integration/test_installation.py
@@ -7,17 +7,21 @@
 from datetime import timedelta
 
 import pytest  # pylint: disable=wrong-import-order
+from databricks.labs.ucx.__about__ import __version__
+
 from databricks.labs.blueprint.installation import Installation
 from databricks.labs.blueprint.installer import RawState
 from databricks.labs.blueprint.parallel import ManyError
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.labs.blueprint.wheels import ProductInfo
+from databricks.sdk import AccountClient
 from databricks.labs.lsql.backends import StatementExecutionBackend
 from databricks.sdk.errors import (
     AlreadyExists,
     InvalidParameterValue,
     NotFound,
 )
+
 from databricks.sdk.retries import retried
 from databricks.sdk.service import compute
 from databricks.sdk.service.iam import PermissionLevel
@@ -453,7 +457,13 @@ def test_compare_remote_local_install_versions(ws, installation_ctx):
 
 
 def test_new_collection(ws, sql_backend, installation_ctx, env_or_skip):
-    installation_ctx.workspace_installation.run()
+    host = ws.config.environment.deployment_url("accounts")
+    acc_client = AccountClient(
+        host=host, account_id=env_or_skip("DATABRICKS_ACCOUNT_ID"), product='ucx', product_version=__version__
+    )
+    workspace = acc_client.workspaces.get(ws.get_workspace_id())
+    installation_ctx.with_workspace_info([workspace])
+    installation_ctx.workspace_installer.run()
     workspace_id = installation_ctx.workspace_installer.workspace_client.get_workspace_id()
     acc_installer = installation_ctx.account_installer
     prompts = MockPrompts(
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py
diff --git a/tests/unit/test_collection.py b/tests/unit/test_collection.py
