Repository: databrickslabs/ucx

Issue:
[FEATURE]: HMS Federation - add support for Postgres and SQL Server
### Is there an existing issue for this?

- [x] I have searched the existing issues

### Problem statement

Support for Postgres and SQL Server was added to HMS federation as well as support for HMS 3.1
The HMS federation feature in UCX has to be updated, accordingly.

### Proposed Solution

Add support for:
SQL Server
Postgres
HMS3.1

### Additional Context

_No response_

Discussion:

Solution patch:
diff --git a/src/databricks/labs/ucx/hive_metastore/federation.py b/src/databricks/labs/ucx/hive_metastore/federation.py
index 7768e07e..0a180713 100644
--- a/src/databricks/labs/ucx/hive_metastore/federation.py
+++ b/src/databricks/labs/ucx/hive_metastore/federation.py
@@ -78,9 +78,23 @@ class HiveMetastoreFederation(SecretsMixin):
         self._enable_hms_federation = enable_hms_federation
         self._config = config
 
-    # Supported databases and version for HMS Federation
-    supported_database_versions: ClassVar[dict[str, list[str]]] = {
-        "mysql": ["2.3", "0.13"],
+    # Supported databases and associated ports
+    # https://docs.databricks.com/en/data-governance/unity-catalog/hms-federation/hms-federation-external.html
+    # https://dev.mysql.com/doc/mysql-port-reference/en/mysql-port-reference-tables.html
+    # https://www.postgresql.org/docs/current/runtime-config-connection.html
+    # https://learn.microsoft.com/en-us/sql/connect/jdbc/building-the-connection-url?view=sql-server-ver15
+    supported_databases_port: ClassVar[dict[str, int]] = {
+        "mysql": 3306,
+        "postgresql": 5432,
+        "sqlserver": 1433,
+    }
+
+    # Supported HMS versions
+    # https://docs.databricks.com/en/data-governance/unity-catalog/hms-federation/hms-federation-external.html
+    supported_hms_versions: ClassVar[set[tuple[int, int]]] = {
+        (0, 13),
+        (2, 3),
+        (3, 1),
     }
 
     def create_from_cli(self, prompts: Prompts) -> None:
@@ -127,19 +141,24 @@ class HiveMetastoreFederation(SecretsMixin):
         if not version:
             logger.info('Hive Metastore version not found')
             return None
-        major_minor_match = re.match(r'(^\d+\.\d+)', version)
-        if not major_minor_match:
+        major_minor_match = re.match(r'(^(?P<major>\d+)\.(?P<minor>\d+))', version)
+        if not major_minor_match or not major_minor_match.group('major') or not major_minor_match.group('minor'):
             logger.info(f'Wrong Hive Metastore Database Version Format: {version}')
             return None
-        major_minor_version = major_minor_match.group(1)
-        external_hms = replace(self._split_jdbc_url(jdbc_url), version=major_minor_version)
-        supported_versions = self.supported_database_versions.get(external_hms.database_type)
-        if not supported_versions:
-            logger.info(f'Unsupported Hive Metastore: {external_hms.database_type}')
+        try:
+            major = int(major_minor_match.group('major'))
+            minor = int(major_minor_match.group('minor'))
+        except ValueError:
+            logger.info(f'Wrong Hive Metastore Database Version Format: {version}')
             return None
-        if major_minor_version not in supported_versions:
-            logger.info(f'Unsupported Hive Metastore Version: {external_hms.database_type} - {version}')
+
+        # Verify HMS version
+        if (major, minor) not in self.supported_hms_versions:
+            logger.info(
+                f'Unsupported Hive Metastore Version: {version}. We currently support: {self.supported_hms_versions}'
+            )
             return None
+        external_hms = replace(self._split_jdbc_url(jdbc_url), version=f'{major}.{minor}')
 
         if not external_hms.user:
             external_hms = replace(
@@ -158,19 +177,33 @@ class HiveMetastoreFederation(SecretsMixin):
     @classmethod
     def _split_jdbc_url(cls, jdbc_url: str) -> ExternalHmsInfo:
         # Define the regex pattern to match the JDBC URL components
+        # The regex supports the following JDBC URL formats:
+        # 1. jdbc:mysql://hostname:3306/metastore
+        # 2. jdbc:mysql://hostname/metastore
+        # 3. jdbc:mysql://hostname:3306/metastore?user=foo&password=bar
+        # 4. jdbc:mysql://hostname/metastore?user=foo&password=bar
+        # 5. jdbc:mssql://hostname:1433;database=database;user=foo;password=bar
         pattern = re.compile(
-            r'jdbc:(?P<db_type>[a-zA-Z0-9]+)://(?P<host>[^:/]+):(?P<port>\d+)/(?P<database>[^?]+)(\?user=(?P<user>[^&]+)&password=(?P<password>[^&]+))?'
+            r'jdbc:(?P<db_type>[a-zA-Z0-9]+)://(?P<host>[^:/?;]+)(:(?P<port>\d+))?(/(?P<database>[^?^;]+))?([?;](?P<parameters>.+))?'
         )
         match = pattern.match(jdbc_url)
         if not match:
             raise ValueError(f'Unsupported JDBC URL: {jdbc_url}')
 
+        params = {}
+        if match.group('parameters'):
+            params = dict(param.split('=') for param in re.split(r"[;&]", match.group('parameters')))
+
         db_type = match.group('db_type')
+        port = match.group('port') or str(cls.supported_databases_port.get(db_type))
+        if not port:
+            raise ValueError(f"Can't identify Port for {db_type}")
         host = match.group('host')
-        port = match.group('port')
-        database = match.group('database')
-        user = match.group('user')
-        password = match.group('password')
+        database = match.group('database') or params.get("database")
+        if not database or not isinstance(database, str):
+            raise ValueError(f"Can't identify Database for {db_type}")
+        user = params.get('user')
+        password = params.get('password')
 
         return ExternalHmsInfo(db_type, host, port, database, user, password, None)
 


Test patch:
diff --git a/tests/unit/hive_metastore/test_federation.py b/tests/unit/hive_metastore/test_federation.py
index a8c142a7..43abf232 100644
--- a/tests/unit/hive_metastore/test_federation.py
+++ b/tests/unit/hive_metastore/test_federation.py
@@ -1,5 +1,7 @@
 import base64
 from unittest.mock import create_autospec, call
+import pytest
+
 
 from databricks.labs.blueprint.installation import MockInstallation
 from databricks.labs.blueprint.tui import MockPrompts
@@ -87,7 +89,59 @@ def test_create_federated_catalog_int(mock_installation):
     assert calls == workspace_client.grants.method_calls
 
 
-def test_create_federated_catalog_ext(mock_installation):
+@pytest.mark.parametrize(
+    "config, expected",
+    [
+        (
+            {
+                "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.version": "2.3.0",
+            },
+            {
+                'database': 'metastore',
+                'db_type': 'mysql',
+                'host': 'hostname.us-east-2.rds.amazonaws.com',
+                'password': 'bar',
+                'port': '3306',
+                'user': 'foo',
+                'version': '2.3',
+            },
+        ),
+        (
+            {
+                "spark.hadoop.javax.jdo.option.ConnectionURL": "jdbc:mysql://hostname.us-east-2.rds.amazonaws.com:3306/metastore?user=foo&password=bar",
+                "spark.sql.hive.metastore.version": "3.1.2",
+            },
+            {
+                'database': 'metastore',
+                'db_type': 'mysql',
+                'host': 'hostname.us-east-2.rds.amazonaws.com',
+                'password': 'bar',
+                'port': '3306',
+                'user': 'foo',
+                'version': '3.1',
+            },
+        ),
+        (
+            {
+                "spark.hadoop.javax.jdo.option.ConnectionURL": "jdbc:sqlserver://teststableip.database.windows.net;database=teststableip;user=teststableip@teststableip;password=bar",
+                "spark.sql.hive.metastore.version": "3.1.2",
+            },
+            {
+                'database': 'teststableip',
+                'db_type': 'sqlserver',
+                'host': 'teststableip.database.windows.net',
+                'password': 'bar',
+                'port': '1433',
+                'user': 'teststableip@teststableip',
+                'version': '3.1',
+            },
+        ),
+    ],
+)
+def test_create_federated_catalog_ext(mock_installation, config, expected):
     workspace_client = create_autospec(WorkspaceClient)
     external_locations = create_autospec(ExternalLocations)
     workspace_info = create_autospec(WorkspaceInfo)
@@ -109,14 +163,7 @@ def test_create_federated_catalog_ext(mock_installation):
     )
     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",
-        },
+        spark_conf=config,
     )
 
     hms_fed = HiveMetastoreFederation(
@@ -134,15 +181,7 @@ def test_create_federated_catalog_ext(mock_installation):
     workspace_client.connections.create.assert_called_with(
         name='fed_source',
         connection_type=ConnectionType.HIVE_METASTORE,
-        options={
-            'database': 'metastore',
-            'db_type': 'mysql',
-            'host': 'hostname.us-east-2.rds.amazonaws.com',
-            'password': 'bar',
-            'port': '3306',
-            'user': 'foo',
-            'version': '2.3',
-        },
+        options=expected,
     )
     workspace_client.catalogs.create.assert_called_with(
         name='a',