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

databricks/labs/lakebridge/reconcile/normalize_recon_config_service.py

@@ -0,0 +1,133 @@
+import dataclasses
+
+from databricks.labs.lakebridge.reconcile.connectors.data_source import DataSource
+from databricks.labs.lakebridge.reconcile.recon_config import (
+    Table,
+    Aggregate,
+    ColumnMapping,
+    Transformation,
+    ColumnThresholds,
+)
+
+
+class NormalizeReconConfigService:
+    def __init__(self, source: DataSource, target: DataSource):
+        self.source = source
+        self.target = target
+
+    def normalize_recon_table_config(self, table: Table) -> Table:
+        normalized_table = dataclasses.replace(table)
+
+        self._normalize_sampling(normalized_table)
+        self._normalize_aggs(normalized_table)
+        self._normalize_join_cols(normalized_table)
+        self._normalize_select_cols(normalized_table)
+        self._normalize_drop_cols(normalized_table)
+        self._normalize_col_mappings(normalized_table)
+        self._normalize_transformations(normalized_table)
+        self._normalize_col_thresholds(normalized_table)
+        self._normalize_jdbc_options(normalized_table)
+
+        return normalized_table
+
+    def _normalize_sampling(self, table: Table):
+        if table.sampling_options:
+            normalized_sampling = dataclasses.replace(table.sampling_options)
+            normalized_sampling.stratified_columns = (
+                [self.source.normalize_identifier(c).ansi_normalized for c in normalized_sampling.stratified_columns]
+                if normalized_sampling.stratified_columns
+                else None
+            )
+            table.sampling_options = normalized_sampling
+        return table
+
+    def _normalize_aggs(self, table: Table):
+        normalized = [self._normalize_agg(a) for a in table.aggregates] if table.aggregates else None
+        table.aggregates = normalized
+        return table
+
+    def _normalize_agg(self, agg: Aggregate) -> Aggregate:
+        normalized = dataclasses.replace(agg)
+        normalized.agg_columns = [self.source.normalize_identifier(c).ansi_normalized for c in normalized.agg_columns]
+        normalized.group_by_columns = (
+            [self.source.normalize_identifier(c).ansi_normalized for c in normalized.group_by_columns]
+            if normalized.group_by_columns
+            else None
+        )
+        return normalized
+
+    def _normalize_join_cols(self, table: Table):
+        table.join_columns = (
+            [self.source.normalize_identifier(c).ansi_normalized for c in table.join_columns]
+            if table.join_columns
+            else None
+        )
+        return table
+
+    def _normalize_select_cols(self, table: Table):
+        table.select_columns = (
+            [self.source.normalize_identifier(c).ansi_normalized for c in table.select_columns]
+            if table.select_columns
+            else None
+        )
+        return table
+
+    def _normalize_drop_cols(self, table: Table):
+        table.drop_columns = (
+            [self.source.normalize_identifier(c).ansi_normalized for c in table.drop_columns]
+            if table.drop_columns
+            else None
+        )
+        return table
+
+    def _normalize_col_mappings(self, table: Table):
+        table.column_mapping = (
+            [self._normalize_col_mapping(m) for m in table.column_mapping] if table.column_mapping else None
+        )
+        return table
+
+    def _normalize_col_mapping(self, mapping: ColumnMapping):
+        return ColumnMapping(
+            source_name=self.source.normalize_identifier(mapping.source_name).ansi_normalized,
+            target_name=self.target.normalize_identifier(mapping.target_name).ansi_normalized,
+        )
+
+    def _normalize_transformations(self, table: Table):
+        table.transformations = (
+            [self._normalize_transformation(t) for t in table.transformations] if table.transformations else None
+        )
+        return table
+
+    def _normalize_transformation(self, transform: Transformation):
+        """normalize user-configured transformations
+
+        The user configures the table column and passes SQL code to transform the source table and target table.
+        This is useful in scenarios when the data changes e.g. migrating `datetime`s. The SQL code is not normalized
+        and it is the user responsibility to pass valid SQL respecting source database and target database.
+        """
+        normalized = dataclasses.replace(transform)
+        normalized.column_name = self.source.normalize_identifier(transform.column_name).ansi_normalized
+        return normalized
+
+    def _normalize_col_thresholds(self, table: Table):
+        table.column_thresholds = (
+            [self._normalize_col_threshold(t) for t in table.column_thresholds] if table.column_thresholds else None
+        )
+        return table
+
+    def _normalize_col_threshold(self, threshold: ColumnThresholds):
+        normalized = dataclasses.replace(threshold)
+        normalized.column_name = self.source.normalize_identifier(threshold.column_name).ansi_normalized
+        return normalized
+
+    def _normalize_jdbc_options(self, table: Table):
+        if table.jdbc_reader_options:
+            normalized = dataclasses.replace(table.jdbc_reader_options)
+            normalized.partition_column = (
+                self.source.normalize_identifier(normalized.partition_column).ansi_normalized
+                if normalized.partition_column
+                else None
+            )
+            table.jdbc_reader_options = normalized
+
+        return table

databricks/labs/lakebridge/reconcile/query_builder/base.py

@@ -4,10 +4,13 @@ from abc import ABC
 import sqlglot.expressions as exp
 from sqlglot import Dialect, parse_one
 
+from databricks.labs.lakebridge.reconcile.connectors.data_source import DataSource
+from databricks.labs.lakebridge.reconcile.connectors.dialect_utils import DialectUtils
 from databricks.labs.lakebridge.reconcile.exception import InvalidInputException
 from databricks.labs.lakebridge.reconcile.query_builder.expression_generator import (
     DataType_transform_mapping,
     transform_expression,
+    build_column,
 )
 from databricks.labs.lakebridge.reconcile.recon_config import Schema, Table, Aggregate
 from databricks.labs.lakebridge.transpiler.sqlglot.dialect_utils import get_dialect, SQLGLOT_DIALECTS
@@ -16,21 +19,16 @@ logger = logging.getLogger(__name__)
 
 
 class QueryBuilder(ABC):
-    def __init__(
-        self,
-        table_conf: Table,
-        schema: list[Schema],
-        layer: str,
-        engine: Dialect,
-    ):
+    def __init__(self, table_conf: Table, schema: list[Schema], layer: str, engine: Dialect, data_source: DataSource):
         self._table_conf = table_conf
         self._schema = schema
         self._layer = layer
         self._engine = engine
+        self._data_source = data_source
 
     @property
     def engine(self) -> Dialect:
-        return self._engine
+        return self._engine if self.layer == "source" else get_dialect("databricks")
 
     @property
     def layer(self) -> str:
@@ -70,7 +68,25 @@ class QueryBuilder(ABC):
 
     @property
     def user_transformations(self) -> dict[str, str]:
-        return self._table_conf.get_transformation_dict(self._layer)
+        if self._table_conf.transformations:
+            if self.layer == "source":
+                return {
+                    trans.column_name: (
+                        trans.source
+                        if trans.source
+                        else self._data_source.normalize_identifier(trans.column_name).source_normalized
+                    )
+                    for trans in self._table_conf.transformations
+                }
+            return {
+                self._table_conf.get_layer_src_to_tgt_col_mapping(trans.column_name, self.layer): (
+                    trans.target
+                    if trans.target
+                    else self._table_conf.get_layer_src_to_tgt_col_mapping(trans.column_name, self.layer)
+                )
+                for trans in self._table_conf.transformations
+            }
+        return {}
 
     @property
     def aggregates(self) -> list[Aggregate] | None:
@@ -93,10 +109,12 @@ class QueryBuilder(ABC):
 
     def _user_transformer(self, node: exp.Expression, user_transformations: dict[str, str]) -> exp.Expression:
         if isinstance(node, exp.Column) and user_transformations:
-            dialect = self.engine if self.layer == "source" else get_dialect("databricks")
-            column_name = node.name
-            if column_name in user_transformations.keys():
-                return parse_one(user_transformations.get(column_name, column_name), read=dialect)
+            normalized_column = self._data_source.normalize_identifier(node.name)
+            ansi_name = normalized_column.ansi_normalized
+            if ansi_name in user_transformations.keys():
+                return parse_one(
+                    user_transformations.get(ansi_name, normalized_column.source_normalized), read=self.engine
+                )
         return node
 
     def _apply_default_transformation(
@@ -107,8 +125,7 @@ class QueryBuilder(ABC):
             with_transform.append(alias.transform(self._default_transformer, schema, source))
         return with_transform
 
-    @staticmethod
-    def _default_transformer(node: exp.Expression, schema: list[Schema], source: Dialect) -> exp.Expression:
+    def _default_transformer(self, node: exp.Expression, schema: list[Schema], source: Dialect) -> exp.Expression:
 
         def _get_transform(datatype: str):
             source_dialects = [source_key for source_key, dialect in SQLGLOT_DIALECTS.items() if dialect == source]
@@ -125,9 +142,10 @@ class QueryBuilder(ABC):
 
         schema_dict = {v.column_name: v.data_type for v in schema}
         if isinstance(node, exp.Column):
-            column_name = node.name
-            if column_name in schema_dict.keys():
-                transform = _get_transform(schema_dict.get(column_name, column_name))
+            normalized_column = self._data_source.normalize_identifier(node.name)
+            ansi_name = normalized_column.ansi_normalized
+            if ansi_name in schema_dict.keys():
+                transform = _get_transform(schema_dict.get(ansi_name, normalized_column.source_normalized))
                 return transform_expression(node, transform)
         return node
 
@@ -136,3 +154,20 @@ class QueryBuilder(ABC):
             message = f"Exception for {self.table_conf.target_name} target table in {self.layer} layer --> {message}"
             logger.error(message)
             raise InvalidInputException(message)
+
+    def _build_column_with_alias(self, column: str):
+        return build_column(
+            this=self._build_column_name_source_normalized(column),
+            alias=DialectUtils.unnormalize_identifier(
+                self.table_conf.get_layer_tgt_to_src_col_mapping(column, self.layer)
+            ),
+            quoted=True,
+        )
+
+    def _build_column_name_source_normalized(self, column: str):
+        return self._data_source.normalize_identifier(column).source_normalized
+
+    def _build_alias_source_normalized(self, column: str):
+        return self._data_source.normalize_identifier(
+            self.table_conf.get_layer_tgt_to_src_col_mapping(column, self.layer)
+        ).source_normalized

databricks/labs/lakebridge/reconcile/query_builder/expression_generator.py

@@ -125,6 +125,7 @@ def anonymous(expr: exp.Column, func: str, is_expr: bool = False, dialect=None)
     return new_expr
 
 
+# TODO Standardize impl and use quoted and Identifier/Column consistently
 def build_column(this: exp.ExpOrStr, table_name="", quoted=False, alias=None) -> exp.Expression:
     if alias:
         if isinstance(this, str):
@@ -135,6 +136,10 @@ def build_column(this: exp.ExpOrStr, table_name="", quoted=False, alias=None) ->
     return exp.Column(this=exp.Identifier(this=this, quoted=quoted), table=table_name)
 
 
+def build_column_no_alias(this: str, table_name="") -> exp.Expression:
+    return exp.Column(this=this, table=table_name)
+
+
 def build_literal(this: exp.ExpOrStr, alias=None, quoted=False, is_string=True, cast=None) -> exp.Expression:
     base_literal = exp.Literal(this=this, is_string=is_string)
     if not cast and not alias:
@@ -207,10 +212,11 @@ def build_sub(
     right_column_name: str,
     left_table_name: str | None = None,
     right_table_name: str | None = None,
+    quoted: bool = False,
 ) -> exp.Sub:
     return exp.Sub(
-        this=build_column(left_column_name, left_table_name),
-        expression=build_column(right_column_name, right_table_name),
+        this=build_column(left_column_name, left_table_name, quoted=quoted),
+        expression=build_column(right_column_name, right_table_name, quoted=quoted),
     )
 
 

databricks/labs/lakebridge/reconcile/query_builder/hash_query.py

@@ -11,8 +11,8 @@ from databricks.labs.lakebridge.reconcile.query_builder.expression_generator imp
     get_hash_transform,
     lower,
     transform_expression,
+    build_column_no_alias,
 )
-from databricks.labs.lakebridge.transpiler.sqlglot.dialect_utils import get_dialect
 
 logger = logging.getLogger(__name__)
 
@@ -41,15 +41,12 @@ class HashQueryBuilder(QueryBuilder):
 
         key_cols = hash_cols if report_type == "row" else sorted(_join_columns | self.partition_column)
 
-        cols_with_alias = [
-            build_column(this=col, alias=self.table_conf.get_layer_tgt_to_src_col_mapping(col, self.layer))
-            for col in key_cols
-        ]
+        cols_with_alias = [self._build_column_with_alias(col) for col in key_cols]
 
         # in case if we have column mapping, we need to sort the target columns in the order of source columns to get
         # same hash value
         hash_cols_with_alias = [
-            {"this": col, "alias": self.table_conf.get_layer_tgt_to_src_col_mapping(col, self.layer)}
+            {"this": self._build_column_name_source_normalized(col), "alias": self._build_alias_source_normalized(col)}
             for col in hash_cols
         ]
         sorted_hash_cols_with_alias = sorted(hash_cols_with_alias, key=lambda column: column["alias"])
@@ -60,12 +57,11 @@ class HashQueryBuilder(QueryBuilder):
         )
         hash_col_with_transform = [self._generate_hash_algorithm(hashcols_sorted_as_src_seq, _HASH_COLUMN_NAME)]
 
-        dialect = self.engine if self.layer == "source" else get_dialect("databricks")
         res = (
             exp.select(*hash_col_with_transform + key_cols_with_transform)
             .from_(":tbl")
-            .where(self.filter)
-            .sql(dialect=dialect)
+            .where(self.filter, dialect=self.engine)
+            .sql(dialect=self.engine)
         )
 
         logger.info(f"Hash Query for {self.layer}: {res}")
@@ -76,10 +72,8 @@ class HashQueryBuilder(QueryBuilder):
         cols: list[str],
         column_alias: str,
     ) -> exp.Expression:
-        cols_with_alias = [build_column(this=col, alias=None) for col in cols]
-        cols_with_transform = self.add_transformations(
-            cols_with_alias, self.engine if self.layer == "source" else get_dialect("databricks")
-        )
+        cols_no_alias = [build_column_no_alias(this=col) for col in cols]
+        cols_with_transform = self.add_transformations(cols_no_alias, self.engine)
         col_exprs = exp.select(*cols_with_transform).iter_expressions()
         concat_expr = concat(list(col_exprs))
 

databricks/labs/lakebridge/reconcile/query_builder/sampling_query.py

@@ -4,6 +4,7 @@ import sqlglot.expressions as exp
 from pyspark.sql import DataFrame
 from sqlglot import select
 
+from databricks.labs.lakebridge.reconcile.connectors.dialect_utils import DialectUtils
 from databricks.labs.lakebridge.transpiler.sqlglot.dialect_utils import get_key_from_dialect
 from databricks.labs.lakebridge.reconcile.query_builder.base import QueryBuilder
 from databricks.labs.lakebridge.reconcile.query_builder.expression_generator import (
@@ -37,12 +38,9 @@ class SamplingQueryBuilder(QueryBuilder):
 
         cols = sorted((join_columns | self.select_columns) - self.threshold_columns - self.drop_columns)
 
-        cols_with_alias = [
-            build_column(this=col, alias=self.table_conf.get_layer_tgt_to_src_col_mapping(col, self.layer))
-            for col in cols
-        ]
+        cols_with_alias = [self._build_column_with_alias(col) for col in cols]
 
-        query = select(*cols_with_alias).from_(":tbl").where(self.filter).sql(dialect=self.engine)
+        query = select(*cols_with_alias).from_(":tbl").where(self.filter, dialect=self.engine).sql(dialect=self.engine)
 
         logger.info(f"Sampling Query with Alias for {self.layer}: {query}")
         return query
@@ -59,22 +57,22 @@ class SamplingQueryBuilder(QueryBuilder):
 
         cols = sorted((join_columns | self.select_columns) - self.threshold_columns - self.drop_columns)
 
-        cols_with_alias = [
-            build_column(this=col, alias=self.table_conf.get_layer_tgt_to_src_col_mapping(col, self.layer))
-            for col in cols
-        ]
+        cols_with_alias = [self._build_column_with_alias(col) for col in cols]
 
         sql_with_transforms = self.add_transformations(cols_with_alias, self.engine)
-        query_sql = select(*sql_with_transforms).from_(":tbl").where(self.filter)
+        query_sql = select(*sql_with_transforms).from_(":tbl").where(self.filter, dialect=self.engine)
         if self.layer == "source":
-            with_select = [build_column(this=col, table_name="src") for col in sorted(cols)]
+            with_select = [
+                build_column(this=DialectUtils.unnormalize_identifier(col), table_name="src", quoted=True)
+                for col in sorted(cols)
+            ]
         else:
             with_select = [
-                build_column(this=col, table_name="src")
+                build_column(this=DialectUtils.unnormalize_identifier(col), table_name="src", quoted=True)
                 for col in sorted(self.table_conf.get_tgt_to_src_col_mapping_list(cols))
             ]
 
-        join_clause = SamplingQueryBuilder._get_join_clause(key_cols)
+        join_clause = self._get_join_clause(key_cols)
 
         query = (
             with_clause.with_(alias="src", as_=query_sql)
@@ -86,10 +84,10 @@ class SamplingQueryBuilder(QueryBuilder):
         logger.info(f"Sampling Query for {self.layer}: {query}")
         return query
 
-    @classmethod
-    def _get_join_clause(cls, key_cols: list):
+    def _get_join_clause(self, key_cols: list):
+        normalized = [self._build_column_name_source_normalized(col) for col in key_cols]
         return build_join_clause(
-            "recon", key_cols, source_table_alias="src", target_table_alias="recon", kind="inner", func=exp.EQ
+            "recon", normalized, source_table_alias="src", target_table_alias="recon", kind="inner", func=exp.EQ
         )
 
     def _get_with_clause(self, df: DataFrame) -> exp.Select:
@@ -106,12 +104,13 @@ class SamplingQueryBuilder(QueryBuilder):
                 (
                     build_literal(
                         this=str(value),
-                        alias=col,
+                        alias=DialectUtils.unnormalize_identifier(col),
                         is_string=_get_is_string(column_types_dict, col),
-                        cast=orig_types_dict.get(col),
+                        cast=orig_types_dict.get(DialectUtils.ansi_normalize_identifier(col)),
+                        quoted=True,
                     )
                     if value is not None
-                    else exp.Alias(this=exp.Null(), alias=col)
+                    else exp.Alias(this=exp.Null(), alias=DialectUtils.unnormalize_identifier(col), quoted=True)
                 )
                 for col, value in zip(df.columns, row)
             ]

databricks/labs/lakebridge/reconcile/query_builder/threshold_query.py

@@ -3,6 +3,7 @@ import logging
 from sqlglot import expressions as exp
 from sqlglot import select
 
+from databricks.labs.lakebridge.reconcile.connectors.dialect_utils import DialectUtils
 from databricks.labs.lakebridge.reconcile.query_builder.base import QueryBuilder
 from databricks.labs.lakebridge.reconcile.query_builder.expression_generator import (
     anonymous,
@@ -54,6 +55,7 @@ class ThresholdQueryBuilder(QueryBuilder):
                     left_table_name="source",
                     right_column_name=column,
                     right_table_name="databricks",
+                    quoted=False,
                 )
             ).transform(coalesce)
 
@@ -62,7 +64,14 @@ class ThresholdQueryBuilder(QueryBuilder):
             where_clause.append(where)
         # join columns
         for column in sorted(join_columns):
-            select_clause.append(build_column(this=column, alias=f"{column}_source", table_name="source"))
+            select_clause.append(
+                build_column(
+                    this=column,
+                    alias=f"{DialectUtils.unnormalize_identifier(column)}_source",
+                    table_name="source",
+                    quoted=True,
+                )
+            )
         where = build_where_clause(where_clause)
 
         return select_clause, where
@@ -76,10 +85,20 @@ class ThresholdQueryBuilder(QueryBuilder):
         select_clause = []
         column = threshold.column_name
         select_clause.append(
-            build_column(this=column, alias=f"{column}_source", table_name="source").transform(coalesce)
+            build_column(
+                this=column,
+                alias=f"{DialectUtils.unnormalize_identifier(column)}_source",
+                table_name="source",
+                quoted=True,
+            ).transform(coalesce)
         )
         select_clause.append(
-            build_column(this=column, alias=f"{column}_databricks", table_name="databricks").transform(coalesce)
+            build_column(
+                this=column,
+                alias=f"{DialectUtils.unnormalize_identifier(column)}_databricks",
+                table_name="databricks",
+                quoted=True,
+            ).transform(coalesce)
         )
         where_clause = exp.NEQ(this=base, expression=exp.Literal(this="0", is_string=False))
         return select_clause, where_clause
@@ -110,7 +129,13 @@ class ThresholdQueryBuilder(QueryBuilder):
             logger.error(error_message)
             raise ValueError(error_message)
 
-        select_clause.append(build_column(this=func(base=base, threshold=threshold), alias=f"{column}_match"))
+        select_clause.append(
+            build_column(
+                this=func(base=base, threshold=threshold),
+                alias=f"{DialectUtils.unnormalize_identifier(column)}_match",
+                quoted=True,
+            )
+        )
 
         return select_clause, where_clause
 
@@ -170,8 +195,8 @@ class ThresholdQueryBuilder(QueryBuilder):
                 ),
                 expression=exp.Is(
                     this=exp.Column(
-                        this=exp.Identifier(this=threshold.column_name, quoted=False),
-                        table=exp.Identifier(this='databricks'),
+                        this=threshold.column_name,
+                        table="databricks",
                     ),
                     expression=exp.Null(),
                 ),
@@ -211,21 +236,17 @@ class ThresholdQueryBuilder(QueryBuilder):
         self._validate(self.join_columns, "Join Columns are compulsory for threshold query")
         join_columns = self.join_columns if self.join_columns else set()
         keys: list[str] = sorted(self.partition_column.union(join_columns))
-        keys_select_alias = [
-            build_column(this=col, alias=self.table_conf.get_layer_tgt_to_src_col_mapping(col, self.layer))
-            for col in keys
-        ]
+        keys_select_alias = [self._build_column_with_alias(col) for col in keys]
         keys_expr = self._apply_user_transformation(keys_select_alias)
 
         # threshold column expression
-        threshold_alias = [
-            build_column(this=col, alias=self.table_conf.get_layer_tgt_to_src_col_mapping(col, self.layer))
-            for col in sorted(self.threshold_columns)
-        ]
+        threshold_alias = [self._build_column_with_alias(col) for col in sorted(self.threshold_columns)]
         thresholds_expr = threshold_alias
         if self.user_transformations:
             thresholds_expr = self._apply_user_transformation(threshold_alias)
 
-        query = (select(*keys_expr + thresholds_expr).from_(":tbl").where(self.filter)).sql(dialect=self.engine)
+        query = (select(*keys_expr + thresholds_expr).from_(":tbl").where(self.filter, dialect=self.engine)).sql(
+            dialect=self.engine
+        )
         logger.info(f"Threshold Query for {self.layer}: {query}")
         return query

databricks/labs/lakebridge/reconcile/recon_config.py

@@ -257,21 +257,6 @@ class Table:
             return set()
         return {self.get_layer_src_to_tgt_col_mapping(col, layer) for col in self.drop_columns}
 
-    def get_transformation_dict(self, layer: str) -> dict[str, str]:
-        if self.transformations:
-            if layer == "source":
-                return {
-                    trans.column_name: (trans.source if trans.source else trans.column_name)
-                    for trans in self.transformations
-                }
-            return {
-                self.get_layer_src_to_tgt_col_mapping(trans.column_name, layer): (
-                    trans.target if trans.target else self.get_layer_src_to_tgt_col_mapping(trans.column_name, layer)
-                )
-                for trans in self.transformations
-            }
-        return {}
-
     def get_partition_column(self, layer: str) -> set[str]:
         if self.jdbc_reader_options and layer == "source":
             if self.jdbc_reader_options.partition_column:
@@ -288,8 +273,11 @@ class Table:
 
 @dataclass
 class Schema:
+    # TODO remove: This will have the value of ansi_normalized_column_name. Kept for backwards compatibility.
     column_name: str
     data_type: str
+    ansi_normalized_column_name: str
+    source_normalized_column_name: str
 
 
 @dataclass

databricks/labs/lakebridge/reconcile/recon_output_config.py

@@ -31,7 +31,8 @@ class DataReconcileOutput:
 
 @dataclass
 class SchemaMatchResult:
-    source_column: str
+    source_column_normalized: str
+    source_column_normalized_ansi: str
     source_datatype: str
     databricks_column: str
     databricks_datatype: str