Introduce table definition (#1450)

This pull request introduces significant enhancements to the table
discovery module by adding new data structures and services for geting
table definitions. The changes include the creation of data classes for
table and field information, the introduction of an abstract service for
fetching table definitions, and a concrete implementation for T-SQL
table definitions.

---------

Co-authored-by: sundar.shankar <[email protected]>
Co-authored-by: SundarShankar89 <[email protected]>

Author: sriram251-code

src/databricks/labs/remorph/discovery/table_definition.py

@@ -0,0 +1,23 @@
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from typing import Any
+
+from databricks.labs.remorph.discovery.table import TableDefinition
+
+
+class TableDefinitionService(ABC):
+
+    def __init__(self, connection: Any):
+        self.connection = connection
+
+    @abstractmethod
+    def get_table_definition(self, catalog_name: str) -> Iterable[TableDefinition]:
+        pass
+
+    @abstractmethod
+    def _get_table_definition_query(self, catalog_name: str) -> str:
+        pass
+
+    @abstractmethod
+    def get_all_catalog(self) -> Iterable[str]:
+        pass

src/databricks/labs/remorph/discovery/tsql_table_definition.py

@@ -0,0 +1,185 @@
+from collections.abc import Iterable
+
+from databricks.labs.remorph.connections.database_manager import DatabaseManager
+from databricks.labs.remorph.discovery.table import TableDefinition, TableFQN, FieldInfo
+from databricks.labs.remorph.discovery.table_definition import TableDefinitionService
+
+
+class TsqlTableDefinitionService(TableDefinitionService):
+
+    # Hexadecimal value of § is U+00A7.Hexadecimal value of ‡ (double dagger) is U+2021
+    def _get_table_definition_query(self, catalog_name: str) -> str:
+        query = f"""
+        WITH column_info AS (
+            SELECT
+                TABLE_CATALOG,
+                TABLE_SCHEMA,
+                TABLE_NAME,
+                STRING_AGG(
+                    CONCAT(
+                        column_name,
+                        '§',
+                        CASE
+                            WHEN numeric_precision IS NOT NULL AND numeric_scale IS NOT NULL THEN CONCAT(data_type, '(', numeric_precision, ',', numeric_scale, ')')
+                            WHEN LOWER(data_type) = 'text' THEN CONCAT('varchar', '(', CHARACTER_MAXIMUM_LENGTH, ')')
+                            ELSE data_type
+                        END,
+                        '§',
+                        CASE
+                            WHEN cis.IS_NULLABLE = 'YES' THEN 'true'
+                            ELSE 'false'
+                        END,
+                        '§',
+                        ISNULL(CAST(ep_col.value AS NVARCHAR(MAX)), '')
+                    ),
+                    '‡'
+                ) WITHIN GROUP (ORDER BY ordinal_position) AS DERIVED_SCHEMA
+            FROM
+                {catalog_name}.sys.tables t
+                INNER JOIN {catalog_name}.sys.columns c ON t.object_id = c.object_id
+                INNER JOIN {catalog_name}.INFORMATION_SCHEMA.COLUMNS cis ON t.name = cis.TABLE_NAME AND c.name = cis.COLUMN_NAME
+                OUTER APPLY (
+                    SELECT TOP 1 value
+                    FROM {catalog_name}.sys.extended_properties
+                    WHERE major_id = t.object_id AND minor_id = 0
+                    ORDER BY name DESC
+                ) ep_tbl
+                OUTER APPLY (
+                    SELECT TOP 1 value
+                    FROM {catalog_name}.sys.extended_properties
+                    WHERE major_id = c.object_id AND minor_id = c.column_id
+                    ORDER BY name DESC
+                ) ep_col
+            GROUP BY
+                TABLE_CATALOG,
+                TABLE_SCHEMA,
+                TABLE_NAME
+        ),
+        table_file_info AS (
+            SELECT
+                s.name AS TABLE_SCHEMA,
+                t.name AS TABLE_NAME,
+                f.physical_name AS location,
+                f.type_desc AS TABLE_FORMAT,
+                CAST(ROUND(SUM(a.used_pages) * 8.0 / 1024, 2) AS DECIMAL(18, 2)) AS SIZE_GB
+            FROM
+                {catalog_name}.sys.tables t
+                INNER JOIN {catalog_name}.sys.indexes i ON t.object_id = i.object_id
+                INNER JOIN {catalog_name}.sys.partitions p ON i.object_id = p.object_id AND i.index_id = p.index_id
+                INNER JOIN {catalog_name}.sys.allocation_units a ON p.partition_id = a.container_id
+                INNER JOIN {catalog_name}.sys.schemas s ON t.schema_id = s.schema_id
+                INNER JOIN {catalog_name}.sys.database_files f ON a.data_space_id = f.data_space_id
+                LEFT JOIN {catalog_name}.sys.extended_properties ep ON ep.major_id = t.object_id AND ep.minor_id = 0
+            GROUP BY
+                s.name,
+                t.name,
+                f.name,
+                f.physical_name,
+                f.type_desc
+        ),
+        table_comment_info AS (
+            SELECT
+                s.name AS TABLE_SCHEMA,
+                t.name AS TABLE_NAME,
+                CAST(ep.value AS NVARCHAR(MAX)) AS TABLE_COMMENT
+            FROM
+                {catalog_name}.sys.tables t
+                INNER JOIN {catalog_name}.sys.schemas s ON t.schema_id = s.schema_id
+                OUTER APPLY (
+                    SELECT TOP 1 value
+                    FROM {catalog_name}.sys.extended_properties
+                    WHERE major_id = t.object_id AND minor_id = 0
+                    ORDER BY name DESC
+                ) ep
+        ),
+        table_pk_info AS (
+            SELECT
+                TC.TABLE_CATALOG,
+                TC.TABLE_SCHEMA,
+                TC.TABLE_NAME,
+                STRING_AGG(KU.COLUMN_NAME,':') as PK_COLUMN_NAME
+                FROM {catalog_name}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS TC
+                JOIN {catalog_name}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS KU
+                ON TC.CONSTRAINT_NAME = KU.CONSTRAINT_NAME
+                AND TC.TABLE_NAME = KU.TABLE_NAME
+                WHERE TC.CONSTRAINT_TYPE = 'PRIMARY KEY' group by TC.TABLE_CATALOG, TC.TABLE_SCHEMA, TC.TABLE_NAME)
+        SELECT
+            sft.TABLE_CATALOG,
+            sft.TABLE_SCHEMA,
+            sft.TABLE_NAME,
+            tfi.location,
+            tfi.TABLE_FORMAT,
+            '' as view_definition,
+            column_info.DERIVED_SCHEMA,
+            tfi.SIZE_GB,
+            tci.TABLE_COMMENT,
+            tpK.PK_COLUMN_NAME
+        FROM
+            column_info
+            JOIN {catalog_name}.INFORMATION_SCHEMA.TABLES sft ON column_info.TABLE_CATALOG = sft.TABLE_CATALOG AND column_info.TABLE_SCHEMA = sft.TABLE_SCHEMA AND column_info.TABLE_NAME = sft.TABLE_NAME
+            LEFT JOIN table_file_info tfi ON column_info.TABLE_SCHEMA = tfi.TABLE_SCHEMA AND column_info.TABLE_NAME = tfi.TABLE_NAME
+            LEFT JOIN table_comment_info tci ON column_info.TABLE_SCHEMA = tci.TABLE_SCHEMA AND column_info.TABLE_NAME = tci.TABLE_NAME
+            LEFT JOIN table_pk_info tpK ON column_info.TABLE_SCHEMA = tpK.TABLE_SCHEMA AND column_info.TABLE_NAME = tpK.TABLE_NAME
+
+        UNION ALL
+        SELECT
+            sfv.TABLE_CATALOG,
+            sfv.TABLE_SCHEMA,
+            sfv.TABLE_NAME,
+            '' location,
+            '' TABLE_FORMAT,
+            sfv.view_definition,
+            '' DERIVED_SCHEMA,
+            0 SIZE_GB,
+            '' TABLE_COMMENT,
+            '' PK_COLUMN_NAME
+        FROM {catalog_name}.INFORMATION_SCHEMA.VIEWS sfv
+        """
+        return query
+
+    def get_table_definition(self, catalog_name: str) -> Iterable[TableDefinition]:
+        sql = self._get_table_definition_query(catalog_name)
+        tsql_connection = self.connection
+        result = tsql_connection.execute_query(sql)
+
+        column_names = list(result.keys())
+        table_definitions = []
+
+        for row in result:
+            result = dict(zip(column_names, row))
+            table_fqn = TableFQN(
+                catalog=result["TABLE_CATALOG"], schema=result["TABLE_SCHEMA"], name=result["TABLE_NAME"]
+            )
+            columns = result["DERIVED_SCHEMA"].split("‡") if result["DERIVED_SCHEMA"] else None
+            field_info = []
+            if columns is not None:
+                for column in columns:
+                    column_info = column.split("§")
+                    field = FieldInfo(
+                        name=column_info[0],
+                        data_type=column_info[1],
+                        nullable=column_info[2],
+                        comment=column_info[3],
+                    )
+                    field_info.append(field)
+
+            pks = result["PK_COLUMN_NAME"].split(":") if result["PK_COLUMN_NAME"] else None
+            table_definition = TableDefinition(
+                fqn=table_fqn,
+                location=result["location"],
+                table_format=result["TABLE_FORMAT"],
+                view_text=result["view_definition"],
+                columns=field_info,
+                size_gb=result["SIZE_GB"],
+                comment=result["TABLE_COMMENT"],
+                primary_keys=pks,
+            )
+            table_definitions.append(table_definition)
+        return table_definitions
+
+    def get_all_catalog(self) -> Iterable[str]:
+        cursor: DatabaseManager = self.connection
+        result = cursor.connector.execute_query("""select name from sys.databases""")
+        catalogs = [row[0] for row in result]
+        print(catalogs)
+        return catalogs

tests/integration/discovery/__init__.py

@@ -0,0 +1,22 @@
+import pytest
+
+from databricks.labs.remorph.discovery.tsql_table_definition import TsqlTableDefinitionService
+from ..connections.helpers import get_db_manager
+
+
+@pytest.fixture()
+def extractor(mock_credentials):
+    return get_db_manager("remorph", "mssql")
+
+
+def test_tsql_get_catalog(extractor):
+    tss = TsqlTableDefinitionService(extractor)
+    catalogs = list(tss.get_all_catalog())
+    assert catalogs is not None
+    assert len(catalogs) > 0
+
+
+def test_tsql_get_table_definition(extractor):
+    tss = TsqlTableDefinitionService(extractor)
+    table_def = tss.get_table_definition("labs_azure_sandbox_remorph")
+    assert table_def is not None

tests/integration/discovery/test_tsql_table_definition.py

@@ -0,0 +1,55 @@
+from unittest.mock import MagicMock
+from databricks.labs.remorph.discovery.table import TableDefinition
+from databricks.labs.remorph.discovery.tsql_table_definition import TsqlTableDefinitionService
+
+
+def test_get_table_definition_with_data():
+    db_manager = MagicMock()
+    mock_result = [
+        (
+            "catalog1",
+            "schema1",
+            "table1",
+            "col1§int§YES§Primary Column‡col2§string§NO§Description",
+            "col1:col2",
+            "/path/to/table",
+            "parquet",
+            "",
+            10.5,
+            "Table Comment",
+        ),
+    ]
+
+    mock_column_names = [
+        "TABLE_CATALOG",
+        "TABLE_SCHEMA",
+        "TABLE_NAME",
+        "DERIVED_SCHEMA",
+        "PK_COLUMN_NAME",
+        "location",
+        "TABLE_FORMAT",
+        "view_definition",
+        "SIZE_GB",
+        "TABLE_COMMENT",
+    ]
+
+    mock_query_result = MagicMock()
+    mock_query_result.keys.return_value = mock_column_names
+    mock_query_result.__iter__.return_value = iter(mock_result)
+    db_manager.execute_query.return_value = mock_query_result
+
+    tss = TsqlTableDefinitionService(db_manager)
+    result = list(tss.get_table_definition("test_catalog"))
+    assert result[0].primary_keys == ['col1', 'col2']
+    assert isinstance(result[0], TableDefinition)
+    assert result[0].fqn.catalog == 'catalog1'
+    assert result[0].fqn.schema == 'schema1'
+    assert result[0].fqn.name == 'table1'
+
+
+def test_get_catalogs():
+    db_manager = MagicMock()
+    db_manager.connector.execute_query.return_value = [('db1',), ('db2',)]
+    tss = TsqlTableDefinitionService(db_manager)
+    result = list(tss.get_all_catalog())
+    assert result == ['db1', 'db2']