From debe61a1e8ce3b8aa14725f2aa5d520a4107ddea Mon Sep 17 00:00:00 2001
From: Liran Bareket <liran.bareket@databricks.com>
Date: Thu, 5 Dec 2024 16:15:57 -0500
Subject: [PATCH] Addressed issue with version

---
 src/databricks/labs/ucx/hive_metastore/federation.py | 12 +++++++++---
 tests/unit/hive_metastore/test_federation.py         |  2 +-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/databricks/labs/ucx/hive_metastore/federation.py b/src/databricks/labs/ucx/hive_metastore/federation.py
index 566d785010..46f15f7e0a 100644
--- a/src/databricks/labs/ucx/hive_metastore/federation.py
+++ b/src/databricks/labs/ucx/hive_metastore/federation.py
@@ -62,8 +62,9 @@ def __init__(
         self._enable_hms_federation = enable_hms_federation
         self._installation = installation
 
+    # Supported databases and version for HMS Federation
     supported_db_vers: ClassVar[dict[str, list[str]]] = {
-        "mysql": ["2.3.0", "0.13"],
+        "mysql": ["2.3", "0.13"],
     }
 
     def create_from_cli(self, prompts: Prompts):
@@ -75,7 +76,7 @@ def create_from_cli(self, prompts: Prompts):
         try:
             ext_hms = self._get_ext_hms()
         except ValueError:
-            logger.info('Failed to get external Hive Metastore connection information')
+            logger.info('Failed to retrieve external Hive Metastore connection information')
 
         if ext_hms and prompts.confirm(
             f'A supported external Hive Metastore connection was identified: {ext_hms.db_type}. Use this connection?'
@@ -95,8 +96,13 @@ def _get_ext_hms(self) -> ExtHms:
         if not jdbc_url:
             raise ValueError('JDBC URL not found')
         version = self._get_value_from_config_key(spark_config, 'spark.sql.hive.metastore.version')
+        # extract major version from version using regex
         if not version:
             raise ValueError('Hive Metastore version not found')
+        major_version = re.match(r'(\d+\.\d+)', version)
+        if not major_version:
+            raise ValueError(f'Invalid Hive Metastore version: {version}')
+        version = major_version.group(1)
         ext_hms = replace(self._split_jdbc_url(jdbc_url), version=version)
         supported_versions = self.supported_db_vers.get(ext_hms.db_type)
         if not supported_versions:
@@ -166,7 +172,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": "false",
+            # TODO: Fix once the FEDERATION end point is fixed. Include "builtin": "false" in options
             "database": ext_hms.database,
             "db_type": ext_hms.db_type,
             "host": ext_hms.host,
diff --git a/tests/unit/hive_metastore/test_federation.py b/tests/unit/hive_metastore/test_federation.py
index 3b2fb62f43..0b233b7734 100644
--- a/tests/unit/hive_metastore/test_federation.py
+++ b/tests/unit/hive_metastore/test_federation.py
@@ -129,7 +129,7 @@ def test_create_federated_catalog_ext(mock_installation):
             'password': 'bar',
             'port': '3306',
             'user': 'foo',
-            'version': '2.3.0',
+            'version': '2.3',
         },
     )
     workspace_client.catalogs.create.assert_called_with(