Added test

databrickslabs · Dec 2, 2024 · 6fae8ac · 6fae8ac
1 parent 7096ac0
commit 6fae8ac
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 13 deletions.
diff --git a/src/databricks/labs/ucx/hive_metastore/federation.py b/src/databricks/labs/ucx/hive_metastore/federation.py
@@ -78,7 +78,7 @@ def create_from_cli(self, prompts: Prompts):
             logger.info('Failed to get external Hive Metastore connection information')
 
         if ext_hms and prompts.confirm(
-            f'Identified a supported external Hive Metastore connection: {ext_hms.db_type}. Use this connection?'
+            f'A supported external Hive Metastore connection was identified: {ext_hms.db_type}. Use this connection?'
         ):
             connection_info = self._get_or_create_ext_connection(name, ext_hms)
         else:
@@ -91,8 +91,6 @@ def _get_ext_hms(self) -> ExtHms:
         if not config.spark_conf:
             raise ValueError('Spark config not found')
         spark_config = config.spark_conf
-        if not spark_config:
-            raise ValueError('Spark config not found')
         jdbc_url = self._get_value_from_config_key(spark_config, 'spark.hadoop.javax.jdo.option.ConnectionURL')
         if not jdbc_url:
             raise ValueError('JDBC URL not found')
@@ -168,7 +166,7 @@ def _get_or_create_int_connection(self, name: str) -> ConnectionInfo:
 
     def _get_or_create_ext_connection(self, name: str, ext_hms: ExtHms) -> ConnectionInfo:
         options: dict[str, str] = {
-            "builtin": "true",
+            "builtin": "false",
             "database": ext_hms.database,
             "db_type": ext_hms.db_type,
             "host": ext_hms.host,

diff --git a/tests/integration/hive_metastore/test_federation.py b/tests/integration/hive_metastore/test_federation.py
@@ -1,6 +1,7 @@
 from unittest.mock import create_autospec
 
 import pytest
+from databricks.labs.blueprint.tui import MockPrompts
 from databricks.sdk import WorkspaceClient
 
 from databricks.labs.ucx.account.workspaces import WorkspaceInfo
@@ -14,13 +15,15 @@ def ws():
 
 
 @pytest.mark.skip("needs to be enabled")
-def test_federation(ws, sql_backend):
+def test_federation(ws, ctx, sql_backend):
     schema = 'ucx'
+    installation = ctx.installation
     tables_crawler = TablesCrawler(sql_backend, schema)
     mounts_crawler = MountsCrawler(sql_backend, ws, schema)
     external_locations = ExternalLocations(ws, sql_backend, schema, tables_crawler, mounts_crawler)
     workspace_info = create_autospec(WorkspaceInfo)
     workspace_info.current.return_value = 'some_thing'
-    federation = HiveMetastoreFederation(ws, external_locations, workspace_info)
-    federation.register_internal_hms_as_federated_catalog()
+    federation = HiveMetastoreFederation(installation, ws, external_locations, workspace_info)
+    prompts = MockPrompts({})
+    federation.create_from_cli(prompts)
     workspace_info.current.assert_called_once()
diff --git a/tests/unit/hive_metastore/test_federation.py b/tests/unit/hive_metastore/test_federation.py
@@ -1,6 +1,8 @@
+import base64
 from unittest.mock import create_autospec, call
 
 from databricks.labs.blueprint.installation import MockInstallation
+from databricks.labs.blueprint.tui import MockPrompts
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import AlreadyExists
 from databricks.sdk.service.catalog import (
@@ -15,6 +17,7 @@
     ConnectionInfo,
 )
 from databricks.sdk.service.iam import User
+from databricks.sdk.service.workspace import GetSecretResponse
 
 from databricks.labs.ucx.account.workspaces import WorkspaceInfo
 from databricks.labs.ucx.config import WorkspaceConfig
@@ -23,7 +26,7 @@
 from databricks.labs.ucx.hive_metastore.locations import ExternalLocation
 
 
-def test_create_federated_catalog():
+def test_create_federated_catalog_int(mock_installation):
     workspace_client = create_autospec(WorkspaceClient)
     external_locations = create_autospec(ExternalLocations)
     workspace_info = create_autospec(WorkspaceInfo)
@@ -44,8 +47,11 @@ def test_create_federated_catalog():
         privilege_assignments=[PrivilegeAssignment(privileges=[Privilege.MANAGE], principal='any')]
     )
 
-    hms_fed = HiveMetastoreFederation(workspace_client, external_locations, workspace_info, enable_hms_federation=True)
-    hms_fed.register_internal_hms_as_federated_catalog()
+    hms_fed = HiveMetastoreFederation(
+        mock_installation, workspace_client, external_locations, workspace_info, enable_hms_federation=True
+    )
+
+    hms_fed.create_from_cli(MockPrompts({}))
 
     workspace_client.connections.create.assert_called_with(
         name='a',
@@ -74,7 +80,75 @@ def test_create_federated_catalog():
     assert calls == workspace_client.grants.method_calls
 
 
-def test_already_existing_connection():
+def test_create_federated_catalog_ext(mock_installation):
+    workspace_client = create_autospec(WorkspaceClient)
+    external_locations = create_autospec(ExternalLocations)
+    workspace_info = create_autospec(WorkspaceInfo)
+
+    workspace_info.current.return_value = 'a'
+    external_locations.snapshot.return_value = [
+        ExternalLocation('s3://b/c/d', 1),
+    ]
+    workspace_client.current_user.me.return_value = User(user_name='serge')
+    workspace_client.connections.create.return_value = CatalogInfo(name='a')
+    workspace_client.secrets.get_secret.return_value = GetSecretResponse(
+        key='secret_key', value=base64.standard_b64encode('bar'.encode()).decode()
+    )
+    workspace_client.external_locations.list.return_value = [
+        ExternalLocationInfo(url='s3://b/c/d', name='b'),
+    ]
+    workspace_client.grants.get.return_value = PermissionsList(
+        privilege_assignments=[PrivilegeAssignment(privileges=[Privilege.MANAGE], principal='any')]
+    )
+    mock_installation.load = lambda _: WorkspaceConfig(
+        inventory_database='ucx',
+        spark_conf={
+            "spark.hadoop.javax.jdo.option.ConnectionDriverName": "org.mariadb.jdbc.Driver",
+            "spark.hadoop.javax.jdo.option.ConnectionPassword": "{{secrets/secret_scope/secret_key}}",
+            "spark.hadoop.javax.jdo.option.ConnectionURL": "jdbc:mysql://hostname.us-east-2.rds.amazonaws.com:3306/metastore",
+            "spark.hadoop.javax.jdo.option.ConnectionUserName": "foo",
+            "spark.sql.hive.metastore.jars": "maven",
+            "spark.sql.hive.metastore.version": "2.3.0",
+        },
+    )
+
+    hms_fed = HiveMetastoreFederation(
+        mock_installation, workspace_client, external_locations, workspace_info, enable_hms_federation=True
+    )
+
+    hms_fed.create_from_cli(MockPrompts({"A supported external Hive Metastore.*": "yes"}))
+
+    workspace_client.connections.create.assert_called_with(
+        name='a',
+        connection_type=ConnectionType.HIVE_METASTORE,
+        options={
+            'builtin': 'false',
+            'database': 'metastore',
+            'db_type': 'mysql',
+            'host': 'hostname.us-east-2.rds.amazonaws.com',
+            'password': 'bar',
+            'port': '3306',
+            'user': 'foo',
+            'version': '2.3.0',
+        },
+    )
+    workspace_client.catalogs.create.assert_called_with(
+        name='a',
+        connection_name='a',
+        options={"authorized_paths": 's3://b/c/d'},
+    )
+    calls = [
+        call.get(SecurableType.EXTERNAL_LOCATION, 'b'),
+        call.update(
+            SecurableType.EXTERNAL_LOCATION,
+            'b',
+            changes=[PermissionsChange(principal='serge', add=[Privilege.CREATE_FOREIGN_CATALOG])],
+        ),
+    ]
+    assert calls == workspace_client.grants.method_calls
+
+
+def test_already_existing_connection(mock_installation):
     workspace_client = create_autospec(WorkspaceClient)
     external_locations = create_autospec(ExternalLocations)
     workspace_info = create_autospec(WorkspaceInfo)
@@ -96,8 +170,13 @@ def test_already_existing_connection():
         privilege_assignments=[PrivilegeAssignment(privileges=[Privilege.MANAGE], principal='any')]
     )
 
-    hms_fed = HiveMetastoreFederation(workspace_client, external_locations, workspace_info, enable_hms_federation=True)
-    hms_fed.register_internal_hms_as_federated_catalog()
+    hms_fed = HiveMetastoreFederation(
+        mock_installation, workspace_client, external_locations, workspace_info, enable_hms_federation=True
+    )
+
+    prompts = MockPrompts({})
+
+    hms_fed.create_from_cli(prompts)
 
     workspace_client.connections.create.assert_called_with(
         name='a',