databricks-labs-lakebridge 0.10.6__py3-none-any.whl → 0.10.7__py3-none-any.whl

databricks/labs/lakebridge/reconcile/connectors/data_source.py

@@ -3,6 +3,7 @@ from abc import ABC, abstractmethod
 
 from pyspark.sql import DataFrame
 
+from databricks.labs.lakebridge.reconcile.connectors.models import NormalizedIdentifier
 from databricks.labs.lakebridge.reconcile.exception import DataSourceRuntimeException
 from databricks.labs.lakebridge.reconcile.recon_config import JdbcReaderOptions, Schema
 
@@ -31,12 +32,27 @@ class DataSource(ABC):
     ) -> list[Schema]:
         return NotImplemented
 
+    @abstractmethod
+    def normalize_identifier(self, identifier: str) -> NormalizedIdentifier:
+        pass
+
     @classmethod
     def log_and_throw_exception(cls, exception: Exception, fetch_type: str, query: str):
         error_msg = f"Runtime exception occurred while fetching {fetch_type} using {query} : {exception}"
         logger.warning(error_msg)
         raise DataSourceRuntimeException(error_msg) from exception
 
+    def _map_meta_column(self, meta_column) -> Schema:
+        """Create a normalized Schema DTO from the database metadata
+
+        Used in the implementations of get_schema to build a Schema DTO from the `INFORMATION_SCHEMA` query result.
+        The returned Schema is normalized in case the database is having columns with special characters and standardize
+        """
+        name = meta_column.col_name
+        dtype = meta_column.data_type.strip().lower()
+        normalized = self.normalize_identifier(name)
+        return Schema(normalized.ansi_normalized, dtype, normalized.ansi_normalized, normalized.source_normalized)
+
 
 class MockDataSource(DataSource):
 
@@ -70,3 +86,6 @@ class MockDataSource(DataSource):
         if not mock_schema:
             return self.log_and_throw_exception(self._exception, "schema", f"({catalog}, {schema}, {table})")
         return mock_schema
+
+    def normalize_identifier(self, identifier: str) -> NormalizedIdentifier:
+        return NormalizedIdentifier(identifier, identifier)

databricks/labs/lakebridge/reconcile/connectors/databricks.py

@@ -8,7 +8,9 @@ from pyspark.sql.functions import col
 from sqlglot import Dialect
 
 from databricks.labs.lakebridge.reconcile.connectors.data_source import DataSource
+from databricks.labs.lakebridge.reconcile.connectors.models import NormalizedIdentifier
 from databricks.labs.lakebridge.reconcile.connectors.secrets import SecretsMixin
+from databricks.labs.lakebridge.reconcile.connectors.dialect_utils import DialectUtils
 from databricks.labs.lakebridge.reconcile.recon_config import JdbcReaderOptions, Schema
 from databricks.sdk import WorkspaceClient
 
@@ -35,6 +37,7 @@ def _get_schema_query(catalog: str, schema: str, table: str):
 
 
 class DatabricksDataSource(DataSource, SecretsMixin):
+    _IDENTIFIER_DELIMITER = "`"
 
     def __init__(
         self,
@@ -82,6 +85,13 @@ class DatabricksDataSource(DataSource, SecretsMixin):
             logger.info(f"Fetching Schema: Started at: {datetime.now()}")
             schema_metadata = self._spark.sql(schema_query).where("col_name not like '#%'").distinct().collect()
             logger.info(f"Schema fetched successfully. Completed at: {datetime.now()}")
-            return [Schema(field.col_name.lower(), field.data_type.lower()) for field in schema_metadata]
+            return [self._map_meta_column(field) for field in schema_metadata]
         except (RuntimeError, PySparkException) as e:
             return self.log_and_throw_exception(e, "schema", schema_query)
+
+    def normalize_identifier(self, identifier: str) -> NormalizedIdentifier:
+        return DialectUtils.normalize_identifier(
+            identifier,
+            source_start_delimiter=DatabricksDataSource._IDENTIFIER_DELIMITER,
+            source_end_delimiter=DatabricksDataSource._IDENTIFIER_DELIMITER,
+        )

databricks/labs/lakebridge/reconcile/connectors/dialect_utils.py

@@ -0,0 +1,126 @@
+from databricks.labs.lakebridge.reconcile.connectors.models import NormalizedIdentifier
+
+
+class DialectUtils:
+    _ANSI_IDENTIFIER_DELIMITER = "`"
+
+    @staticmethod
+    def unnormalize_identifier(identifier: str) -> str:
+        """Return an ansi identifier without the outer backticks.
+
+        Use this at your own risk as the missing outer backticks will result in bugs.
+        E.g. <`mary's lamb`> is returned <mary's lamb> so the outer backticks are needed.
+        This is useful for scenarios where the returned identifier will be part of another delimited identifier.
+
+        :param identifier: a database identifier
+        :return: ansi identifier without the outer backticks
+        """
+        ansi = DialectUtils.ansi_normalize_identifier(identifier)
+        unescape = (
+            DialectUtils._unescape_source_end_delimiter(ansi[1:-1], DialectUtils._ANSI_IDENTIFIER_DELIMITER)
+            if ansi
+            else ansi
+        )
+        return unescape
+
+    @staticmethod
+    def ansi_normalize_identifier(identifier: str) -> str:
+        return DialectUtils.normalize_identifier(
+            identifier, DialectUtils._ANSI_IDENTIFIER_DELIMITER, DialectUtils._ANSI_IDENTIFIER_DELIMITER
+        ).ansi_normalized
+
+    @staticmethod
+    def normalize_identifier(
+        identifier: str, source_start_delimiter: str, source_end_delimiter: str
+    ) -> NormalizedIdentifier:
+        identifier = identifier.strip().lower()
+
+        ansi = DialectUtils._normalize_identifier_source_agnostic(
+            identifier,
+            source_start_delimiter,
+            source_end_delimiter,
+            DialectUtils._ANSI_IDENTIFIER_DELIMITER,
+            DialectUtils._ANSI_IDENTIFIER_DELIMITER,
+        )
+
+        # Input was already ansi normalized
+        if ansi == identifier:
+            source = DialectUtils._normalize_identifier_source_agnostic(
+                identifier,
+                DialectUtils._ANSI_IDENTIFIER_DELIMITER,
+                DialectUtils._ANSI_IDENTIFIER_DELIMITER,
+                source_start_delimiter,
+                source_end_delimiter,
+            )
+
+            # Ansi has backticks escaped which has to be unescaped for other delimiters and escape source end delimiters
+            if source != ansi:
+                source = DialectUtils._unescape_source_end_delimiter(source, DialectUtils._ANSI_IDENTIFIER_DELIMITER)
+                source = (
+                    DialectUtils._escape_source_end_delimiter(source, source_start_delimiter, source_end_delimiter)
+                    if source
+                    else source
+                )
+        else:
+            # Make sure backticks are escaped properly for ansi and source end delimiters are unescaped
+            ansi = DialectUtils._unescape_source_end_delimiter(ansi, source_end_delimiter)
+            ansi = DialectUtils._escape_backticks(ansi) if ansi else ansi
+
+            if source_end_delimiter != DialectUtils._ANSI_IDENTIFIER_DELIMITER:
+                ansi = DialectUtils._unescape_source_end_delimiter(ansi, source_end_delimiter)
+
+            source = DialectUtils._normalize_identifier_source_agnostic(
+                identifier, source_start_delimiter, source_end_delimiter, source_start_delimiter, source_end_delimiter
+            )
+
+            # Make sure source end delimiter is escaped else nothing as it was already normalized
+            if source != identifier:
+                source = (
+                    DialectUtils._escape_source_end_delimiter(source, source_start_delimiter, source_end_delimiter)
+                    if source
+                    else source
+                )
+
+        return NormalizedIdentifier(ansi, source)
+
+    @staticmethod
+    def _normalize_identifier_source_agnostic(
+        identifier: str,
+        source_start_delimiter: str,
+        source_end_delimiter: str,
+        expected_source_start_delimiter: str,
+        expected_source_end_delimiter: str,
+    ) -> str:
+        if identifier == "" or identifier is None:
+            return ""
+
+        if DialectUtils.is_already_delimited(
+            identifier, expected_source_start_delimiter, expected_source_end_delimiter
+        ):
+            return identifier
+
+        if DialectUtils.is_already_delimited(identifier, source_start_delimiter, source_end_delimiter):
+            stripped_identifier = identifier.removeprefix(source_start_delimiter).removesuffix(source_end_delimiter)
+        else:
+            stripped_identifier = identifier
+        return f"{expected_source_start_delimiter}{stripped_identifier}{expected_source_end_delimiter}"
+
+    @staticmethod
+    def is_already_delimited(identifier: str, start_delimiter: str, end_delimiter: str) -> bool:
+        return identifier.startswith(start_delimiter) and identifier.endswith(end_delimiter)
+
+    @staticmethod
+    def _escape_backticks(identifier: str) -> str:
+        identifier = identifier[1:-1]
+        identifier = identifier.replace("`", "``")
+        return f"`{identifier}`"
+
+    @staticmethod
+    def _unescape_source_end_delimiter(identifier: str, source_end_delimiter: str) -> str:
+        return identifier.replace(f"{source_end_delimiter}{source_end_delimiter}", source_end_delimiter)
+
+    @staticmethod
+    def _escape_source_end_delimiter(identifier: str, start_end_delimiter, source_end_delimiter: str) -> str:
+        identifier = identifier[1:-1]
+        identifier = identifier.replace(source_end_delimiter, f"{source_end_delimiter}{source_end_delimiter}")
+        return f"{start_end_delimiter}{identifier}{source_end_delimiter}"

databricks/labs/lakebridge/reconcile/connectors/models.py

@@ -0,0 +1,7 @@
+import dataclasses
+
+
+@dataclasses.dataclass
+class NormalizedIdentifier:
+    ansi_normalized: str
+    source_normalized: str

databricks/labs/lakebridge/reconcile/connectors/oracle.py

@@ -9,7 +9,9 @@ from sqlglot import Dialect
 
 from databricks.labs.lakebridge.reconcile.connectors.data_source import DataSource
 from databricks.labs.lakebridge.reconcile.connectors.jdbc_reader import JDBCReaderMixin
+from databricks.labs.lakebridge.reconcile.connectors.models import NormalizedIdentifier
 from databricks.labs.lakebridge.reconcile.connectors.secrets import SecretsMixin
+from databricks.labs.lakebridge.reconcile.connectors.dialect_utils import DialectUtils
 from databricks.labs.lakebridge.reconcile.recon_config import JdbcReaderOptions, Schema
 from databricks.sdk import WorkspaceClient
 
@@ -18,6 +20,7 @@ logger = logging.getLogger(__name__)
 
 class OracleDataSource(DataSource, SecretsMixin, JDBCReaderMixin):
     _DRIVER = "oracle"
+    _IDENTIFIER_DELIMITER = "\""
     _SCHEMA_QUERY = """select column_name, case when (data_precision is not null
                                               and data_scale <> 0)
                                               then data_type || '(' || data_precision || ',' || data_scale || ')'
@@ -91,7 +94,7 @@ class OracleDataSource(DataSource, SecretsMixin, JDBCReaderMixin):
             schema_metadata = df.select([col(c).alias(c.lower()) for c in df.columns]).collect()
             logger.info(f"Schema fetched successfully. Completed at: {datetime.now()}")
             logger.debug(f"schema_metadata: ${schema_metadata}")
-            return [Schema(field.column_name.lower(), field.data_type.lower()) for field in schema_metadata]
+            return [self._map_meta_column(field) for field in schema_metadata]
         except (RuntimeError, PySparkException) as e:
             return self.log_and_throw_exception(e, "schema", schema_query)
 
@@ -106,3 +109,10 @@ class OracleDataSource(DataSource, SecretsMixin, JDBCReaderMixin):
 
     def reader(self, query: str) -> DataFrameReader:
         return self._get_jdbc_reader(query, self.get_jdbc_url, OracleDataSource._DRIVER)
+
+    def normalize_identifier(self, identifier: str) -> NormalizedIdentifier:
+        return DialectUtils.normalize_identifier(
+            identifier,
+            source_start_delimiter=OracleDataSource._IDENTIFIER_DELIMITER,
+            source_end_delimiter=OracleDataSource._IDENTIFIER_DELIMITER,
+        )

databricks/labs/lakebridge/reconcile/connectors/snowflake.py

@@ -11,7 +11,9 @@ from cryptography.hazmat.primitives import serialization
 
 from databricks.labs.lakebridge.reconcile.connectors.data_source import DataSource
 from databricks.labs.lakebridge.reconcile.connectors.jdbc_reader import JDBCReaderMixin
+from databricks.labs.lakebridge.reconcile.connectors.models import NormalizedIdentifier
 from databricks.labs.lakebridge.reconcile.connectors.secrets import SecretsMixin
+from databricks.labs.lakebridge.reconcile.connectors.dialect_utils import DialectUtils
 from databricks.labs.lakebridge.reconcile.exception import InvalidSnowflakePemPrivateKey
 from databricks.labs.lakebridge.reconcile.recon_config import JdbcReaderOptions, Schema
 from databricks.sdk import WorkspaceClient
@@ -22,6 +24,8 @@ logger = logging.getLogger(__name__)
 
 class SnowflakeDataSource(DataSource, SecretsMixin, JDBCReaderMixin):
     _DRIVER = "snowflake"
+    _IDENTIFIER_DELIMITER = "\""
+
     """
        * INFORMATION_SCHEMA:
           - see https://docs.snowflake.com/en/sql-reference/info-schema#considerations-for-replacing-show-commands-with-information-schema-views
@@ -144,9 +148,10 @@ class SnowflakeDataSource(DataSource, SecretsMixin, JDBCReaderMixin):
         try:
             logger.debug(f"Fetching schema using query: \n`{schema_query}`")
             logger.info(f"Fetching Schema: Started at: {datetime.now()}")
-            schema_metadata = self.reader(schema_query).load().collect()
+            df = self.reader(schema_query).load()
+            schema_metadata = df.select([col(c).alias(c.lower()) for c in df.columns]).collect()
             logger.info(f"Schema fetched successfully. Completed at: {datetime.now()}")
-            return [Schema(field.COLUMN_NAME.lower(), field.DATA_TYPE.lower()) for field in schema_metadata]
+            return [self._map_meta_column(field) for field in schema_metadata]
         except (RuntimeError, PySparkException) as e:
             return self.log_and_throw_exception(e, "schema", schema_query)
 
@@ -171,3 +176,10 @@ class SnowflakeDataSource(DataSource, SecretsMixin, JDBCReaderMixin):
                 raise NotFound(message) from e
 
         return self._spark.read.format("snowflake").option("dbtable", f"({query}) as tmp").options(**options)
+
+    def normalize_identifier(self, identifier: str) -> NormalizedIdentifier:
+        return DialectUtils.normalize_identifier(
+            identifier,
+            source_start_delimiter=SnowflakeDataSource._IDENTIFIER_DELIMITER,
+            source_end_delimiter=SnowflakeDataSource._IDENTIFIER_DELIMITER,
+        )

databricks/labs/lakebridge/reconcile/connectors/tsql.py

@@ -9,7 +9,9 @@ from sqlglot import Dialect
 
 from databricks.labs.lakebridge.reconcile.connectors.data_source import DataSource
 from databricks.labs.lakebridge.reconcile.connectors.jdbc_reader import JDBCReaderMixin
+from databricks.labs.lakebridge.reconcile.connectors.models import NormalizedIdentifier
 from databricks.labs.lakebridge.reconcile.connectors.secrets import SecretsMixin
+from databricks.labs.lakebridge.reconcile.connectors.dialect_utils import DialectUtils
 from databricks.labs.lakebridge.reconcile.recon_config import JdbcReaderOptions, Schema
 from databricks.sdk import WorkspaceClient
 
@@ -49,6 +51,7 @@ _SCHEMA_QUERY = """SELECT
 
 class TSQLServerDataSource(DataSource, SecretsMixin, JDBCReaderMixin):
     _DRIVER = "sqlserver"
+    _IDENTIFIER_DELIMITER = {"prefix": "[", "suffix": "]"}
 
     def __init__(
         self,
@@ -122,11 +125,33 @@ class TSQLServerDataSource(DataSource, SecretsMixin, JDBCReaderMixin):
         try:
             logger.debug(f"Fetching schema using query: \n`{schema_query}`")
             logger.info(f"Fetching Schema: Started at: {datetime.now()}")
-            schema_metadata = self.reader(schema_query).load().collect()
+            df = self.reader(schema_query).load()
+            schema_metadata = df.select([col(c).alias(c.lower()) for c in df.columns]).collect()
             logger.info(f"Schema fetched successfully. Completed at: {datetime.now()}")
-            return [Schema(field.COLUMN_NAME.lower(), field.DATA_TYPE.lower()) for field in schema_metadata]
+            return [self._map_meta_column(field) for field in schema_metadata]
         except (RuntimeError, PySparkException) as e:
             return self.log_and_throw_exception(e, "schema", schema_query)
 
     def reader(self, query: str, prepare_query_str="") -> DataFrameReader:
         return self._get_jdbc_reader(query, self.get_jdbc_url, self._DRIVER, prepare_query_str)
+
+    def normalize_identifier(self, identifier: str) -> NormalizedIdentifier:
+        return DialectUtils.normalize_identifier(
+            TSQLServerDataSource._normalize_quotes(identifier),
+            source_start_delimiter=TSQLServerDataSource._IDENTIFIER_DELIMITER["prefix"],
+            source_end_delimiter=TSQLServerDataSource._IDENTIFIER_DELIMITER["suffix"],
+        )
+
+    @staticmethod
+    def _normalize_quotes(identifier: str):
+        if DialectUtils.is_already_delimited(identifier, '"', '"'):
+            identifier = identifier[1:-1]
+            identifier = identifier.replace('""', '"')
+            identifier = (
+                TSQLServerDataSource._IDENTIFIER_DELIMITER["prefix"]
+                + identifier
+                + TSQLServerDataSource._IDENTIFIER_DELIMITER["suffix"]
+            )
+            return identifier
+
+        return identifier

databricks/labs/lakebridge/reconcile/constants.py

@@ -15,10 +15,11 @@ class AutoName(Enum):
 
 
 class ReconSourceType(AutoName):
-    SNOWFLAKE = auto()
-    ORACLE = auto()
     DATABRICKS = auto()
-    TSQL = auto()
+    MSSQL = auto()
+    ORACLE = auto()
+    SNOWFLAKE = auto()
+    SYNAPSE = auto()
 
 
 class ReconReportType(AutoName):