databricks-labs-lakebridge 0.10.0__py3-none-any.whl

databricks/labs/lakebridge/transpiler/sqlglot/sqlglot_engine.py

@@ -0,0 +1,203 @@
+import logging
+import typing as t
+from collections.abc import Iterable
+from dataclasses import dataclass
+from pathlib import Path
+
+from sqlglot import expressions as exp, parse, transpile, Dialect
+from sqlglot.errors import ErrorLevel, ParseError, TokenError, UnsupportedError
+from sqlglot.expressions import Expression
+from sqlglot.tokens import Token, TokenType
+
+from databricks.labs.lakebridge.config import TranspileResult, TranspileConfig
+from databricks.labs.lakebridge.helpers.file_utils import is_sql_file
+from databricks.labs.lakebridge.helpers.string_utils import format_error_message
+from databricks.labs.lakebridge.transpiler.sqlglot import lca_utils
+from databricks.labs.lakebridge.transpiler.sqlglot.dialect_utils import get_dialect
+from databricks.labs.lakebridge.transpiler.sqlglot.dialect_utils import SQLGLOT_DIALECTS
+from databricks.labs.lakebridge.transpiler.transpile_status import TranspileError, ErrorKind, ErrorSeverity
+from databricks.labs.lakebridge.transpiler.transpile_engine import TranspileEngine
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ParsedExpression:
+    original_sql: str
+    parsed_expression: Expression
+
+
+@dataclass
+class ParserProblem:
+    original_sql: str
+    transpile_error: TranspileError
+
+
+class SqlglotEngine(TranspileEngine):
+
+    @property
+    def supported_dialects(self) -> list[str]:
+        return sorted(SQLGLOT_DIALECTS.keys())
+
+    def _partial_transpile(
+        self,
+        read_dialect: Dialect,
+        write_dialect: Dialect,
+        source_code: str,
+        file_path: Path,
+    ) -> tuple[list[str], list[ParserProblem]]:
+        transpiled_sqls: list[str] = []
+        parsed_expressions, problem_list = self.safe_parse(read_dialect, source_code, file_path)
+        for parsed_expression in parsed_expressions:
+            try:
+                transpiled_sql = write_dialect.generate(parsed_expression.parsed_expression, pretty=True)
+                # Checking if the transpiled SQL is a comment and raise an error
+                if transpiled_sql.startswith("--"):
+                    raise UnsupportedError("Unsupported SQL")
+                transpiled_sqls.append(transpiled_sql)
+            except TokenError as e:
+                error_msg = format_error_message("Token Error", e, parsed_expression.original_sql)
+                error = TranspileError("TOKEN_ERROR", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+                problem_list.append(ParserProblem(parsed_expression.original_sql, error))
+            except ParseError as e:
+                error_msg = format_error_message("Parsing Error", e, parsed_expression.original_sql)
+                error = TranspileError("PARSE_ERROR", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+                problem_list.append(ParserProblem(parsed_expression.original_sql, error))
+            except UnsupportedError as e:
+                error_msg = format_error_message("Unsupported SQL Error", e, parsed_expression.original_sql)
+                error = TranspileError("UNSUPPORTED_SQL", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+                problem_list.append(ParserProblem(parsed_expression.original_sql, error))
+        return transpiled_sqls, problem_list
+
+    async def initialize(self, config: TranspileConfig) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def transpile(
+        self, source_dialect: str, target_dialect: str, source_code: str, file_path: Path
+    ) -> TranspileResult:
+        read_dialect = get_dialect(source_dialect)
+        error: TranspileError | None = self._check_supported(read_dialect, source_code, file_path)
+        if error:
+            return TranspileResult(source_code, 1, [error])
+        write_dialect = get_dialect(target_dialect)
+        try:
+            transpiled_expressions = transpile(
+                source_code, read=read_dialect, write=write_dialect, pretty=True, error_level=ErrorLevel.RAISE
+            )
+            transpiled_code = "\n".join(transpiled_expressions)
+            return TranspileResult(transpiled_code, len(transpiled_expressions), [])
+        except (ParseError, TokenError, UnsupportedError) as e:
+            logger.error(f"Exception caught for file {file_path!s}: {e}")
+            transpiled_expressions, problems = self._partial_transpile(
+                read_dialect, write_dialect, source_code, file_path
+            )
+            transpiled_code = "\n".join(transpiled_expressions)
+            return TranspileResult(transpiled_code, 1, [problem.transpile_error for problem in problems])
+
+    def parse(
+        self, source_dialect: str, source_sql: str, file_path: Path
+    ) -> tuple[list[Expression | None] | None, TranspileError | None]:
+        expression = None
+        error = None
+        try:
+            expression = parse(source_sql, read=source_dialect, error_level=ErrorLevel.IMMEDIATE)
+        except TokenError as e:
+            error_msg = format_error_message("Token Error", e, source_sql)
+            error = TranspileError("TOKEN_ERROR", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+        except ParseError as e:
+            error_msg = format_error_message("Parsing Error", e, source_sql)
+            error = TranspileError("PARSE_ERROR", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+        except UnsupportedError as e:
+            error_msg = format_error_message("Unsupported SQL Error", e, source_sql)
+            error = TranspileError("UNSUPPORTED_SQL", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+        return expression, error
+
+    def analyse_table_lineage(
+        self, source_dialect: str, source_code: str, file_path: Path
+    ) -> Iterable[tuple[str, str]]:
+        parsed_expression, _ = self.parse(source_dialect, source_code, file_path)
+        if parsed_expression is not None:
+            for expr in parsed_expression:
+                child: str = str(file_path)
+                if expr is not None:
+                    # TODO: fix possible issue where the file reference is lost (if we have a 'create')
+                    for change in expr.find_all(exp.Create, exp.Insert, exp.Merge, bfs=False):
+                        child = self._find_root_table(change)
+
+                    for query in expr.find_all(exp.Select, exp.Join, exp.With, bfs=False):
+                        table = self._find_root_table(query)
+                        if table:
+                            yield table, child
+
+    def safe_parse(
+        self, read_dialect: Dialect, source_code: str, file_path: Path
+    ) -> tuple[list[ParsedExpression], list[ParserProblem]]:
+        try:
+            tokens = read_dialect.tokenize(sql=source_code)
+            return self._safe_parse(read_dialect, tokens, file_path)
+        except TokenError as e:
+            error_msg = format_error_message("Token error", e, source_code)
+            error = TranspileError("TOKEN_ERROR", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+            return [], [ParserProblem(source_code, error)]
+
+    def _safe_parse(
+        self, read_dialect: Dialect, all_tokens: list[Token], file_path: Path
+    ) -> tuple[list[ParsedExpression], list[ParserProblem]]:
+        chunks = self._make_chunks(all_tokens)
+        parsed_expressions: list[ParsedExpression] = []
+        problems: list[ParserProblem] = []
+        parser_opts = {"error_level": ErrorLevel.RAISE}
+        parser = read_dialect.parser(**parser_opts)
+        for sql, tokens in chunks:
+            try:
+                expressions = parser.parse(tokens)
+                expression = t.cast(Expression, expressions[0])
+                parsed_expressions.append(ParsedExpression(sql, expression))
+            except TokenError as e:
+                error_msg = format_error_message("Token error", e, sql)
+                error = TranspileError("TOKEN_ERROR", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+                problems.append(ParserProblem(sql, error))
+            except ParseError as e:
+                error_msg = format_error_message("Parsing error", e, sql)
+                error = TranspileError("PARSE_ERROR", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+                problems.append(ParserProblem(sql, error))
+            except UnsupportedError as e:
+                error_msg = format_error_message("Unsupported SQL error", e, sql)
+                error = TranspileError("UNSUPPORTED_SQL", ErrorKind.PARSING, ErrorSeverity.ERROR, file_path, error_msg)
+                problems.append(ParserProblem(sql, error))
+            finally:
+                parser.reset()
+        return parsed_expressions, problems
+
+    @staticmethod
+    def _make_chunks(tokens: list[Token]) -> list[tuple[str, list[Token]]]:
+        chunks: list[tuple[str, list[Token]]] = []
+        current_chunk: list[Token] = []
+        # Split tokens into chunks based on semicolons(or other separators)
+        # Need to define the separator in Class Tokenizer
+        for token in tokens:
+            current_chunk.append(token)
+            if token.token_type in {TokenType.SEMICOLON}:
+                original_sql = " ".join([token.text for token in current_chunk]).strip()
+                chunks.append((original_sql, current_chunk))
+                # reset
+                current_chunk = []
+        # don't forget the last chunk
+        if current_chunk:
+            original_sql = " ".join([token.text for token in current_chunk]).strip()
+            chunks.append((original_sql, current_chunk))
+        return chunks
+
+    @staticmethod
+    def _find_root_table(expression) -> str:
+        table = expression.find(exp.Table, bfs=False)
+        return table.name if table else ""
+
+    def _check_supported(self, source_dialect: Dialect, source_code: str, file_path: Path) -> TranspileError | None:
+        return lca_utils.check_for_unsupported_lca(source_dialect, source_code, file_path)
+
+    def is_supported_file(self, file: Path) -> bool:
+        return is_sql_file(file)

databricks/labs/lakebridge/transpiler/transpile_engine.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+import abc
+from pathlib import Path
+
+from databricks.labs.lakebridge.config import TranspileResult, TranspileConfig
+
+
+class TranspileEngine(abc.ABC):
+
+    @classmethod
+    def load_engine(cls, transpiler_config_path: Path) -> TranspileEngine:
+        # TODO remove this once sqlglot transpiler is pluggable
+        if str(transpiler_config_path) == "sqlglot":
+            # pylint: disable=import-outside-toplevel, cyclic-import
+            from databricks.labs.lakebridge.transpiler.sqlglot.sqlglot_engine import SqlglotEngine
+
+            return SqlglotEngine()
+        if not transpiler_config_path.exists():
+            raise ValueError(
+                f"Error: Invalid value for '--transpiler-config-path': '{str(transpiler_config_path)}', file does not exist."
+            )
+        # pylint: disable=import-outside-toplevel, cyclic-import
+        from databricks.labs.lakebridge.transpiler.lsp.lsp_engine import LSPEngine
+
+        return LSPEngine.from_config_path(transpiler_config_path)
+
+    @abc.abstractmethod
+    async def initialize(self, config: TranspileConfig) -> None: ...
+
+    @abc.abstractmethod
+    async def shutdown(self) -> None: ...
+
+    @abc.abstractmethod
+    async def transpile(
+        self, source_dialect: str, target_dialect: str, source_code: str, file_path: Path
+    ) -> TranspileResult: ...
+
+    @property
+    @abc.abstractmethod
+    def supported_dialects(self) -> list[str]: ...
+
+    def check_source_dialect(self, source_dialect: str | None) -> None:
+        if source_dialect not in self.supported_dialects:
+            raise ValueError(
+                f"Invalid value for '--source-dialect': '{source_dialect}' is not one of {self.supported_dialects}."
+            )
+
+    @abc.abstractmethod
+    def is_supported_file(self, file: Path) -> bool: ...

databricks/labs/lakebridge/transpiler/transpile_status.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+
+# not using StrEnum because they only appear with Python 3.11
+class ErrorSeverity(Enum):
+    INFO = "INFO"
+    WARNING = "WARNING"
+    ERROR = "ERROR"
+
+
+class ErrorKind(Enum):
+    ANALYSIS = "ANALYSIS"
+    PARSING = "PARSING"
+    GENERATION = "GENERATION"
+    VALIDATION = "VALIDATION"
+    INTERNAL = "INTERNAL"
+
+
+@dataclass
+class CodePosition:
+    line: int  # 0-based line number
+    character: int  # 0-based character number
+
+
+@dataclass
+class CodeRange:
+    start: CodePosition
+    end: CodePosition
+
+
+@dataclass
+class TranspileError:
+    code: str
+    kind: ErrorKind
+    severity: ErrorSeverity
+    path: Path
+    message: str
+    range: CodeRange | None = None
+
+    def __str__(self):
+        return f"{type(self).__name__}(code={self.code}, kind={self.kind.name}, severity={self.severity.name}, path='{self.path!s}', message='{self.message}')"
+
+
+@dataclass
+class TranspileStatus:
+    file_list: list[Path]
+    no_of_transpiled_queries: int
+    error_list: list[TranspileError]
+
+    @property
+    def analysis_error_count(self) -> int:
+        return len([error for error in self.error_list if error.kind == ErrorKind.ANALYSIS])
+
+    @property
+    def parsing_error_count(self) -> int:
+        return len([error for error in self.error_list if error.kind == ErrorKind.PARSING])
+
+    @property
+    def generation_error_count(self) -> int:
+        return len([error for error in self.error_list if error.kind == ErrorKind.GENERATION])
+
+    @property
+    def validation_error_count(self) -> int:
+        return len([error for error in self.error_list if error.kind == ErrorKind.VALIDATION])

databricks/labs/lakebridge/uninstall.py

@@ -0,0 +1,28 @@
+import logging
+
+from databricks.labs.blueprint.entrypoint import is_in_debug
+from databricks.sdk import WorkspaceClient
+
+from databricks.labs.lakebridge.__about__ import __version__
+from databricks.labs.lakebridge.contexts.application import ApplicationContext
+
+logger = logging.getLogger("databricks.labs.lakebridge.install")
+
+
+def run(context: ApplicationContext):
+    context.workspace_installation.uninstall(context.remorph_config)
+
+
+if __name__ == "__main__":
+    logger.setLevel("INFO")
+    if is_in_debug():
+        logging.getLogger("databricks").setLevel(logging.DEBUG)
+
+    run(
+        ApplicationContext(
+            WorkspaceClient(
+                product="remorph",
+                product_version=__version__,
+            )
+        )
+    )

databricks/labs/lakebridge/upgrades/v0.4.0_add_main_table_operation_name_column.py

@@ -0,0 +1,80 @@
+# pylint: disable=invalid-name
+import logging
+
+
+from databricks.labs.blueprint.installation import Installation
+from databricks.sdk import WorkspaceClient
+
+from databricks.labs.lakebridge.contexts.application import ApplicationContext
+from databricks.labs.lakebridge.deployment.recon import RECON_JOB_NAME
+from databricks.labs.lakebridge.helpers import db_sql
+
+from databricks.labs.lakebridge.deployment.upgrade_common import (
+    current_table_columns,
+    installed_table_columns,
+    recreate_table_sql,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _check_table_mismatch(
+    installed_table,
+    current_table,
+) -> bool:
+    current_table = [x for x in current_table if x != "operation_name"]
+    # Compare the current main table columns with the installed main table columns
+    if "operation_name" in installed_table and len(sorted(installed_table)) != len(sorted(current_table)):
+        return True
+    return False
+
+
+def _upgrade_reconcile_metadata_main_table(
+    installation: Installation,
+    ws: WorkspaceClient,
+    app_context: ApplicationContext,
+):
+    """
+    Add operation_name column to the main table as part of the upgrade process.
+    - Compare the current main table columns with the installed main table columns. If there is any mismatch:
+       * Verify all the current main table columns are present in the installed main table and then use CTAS to recreate the main table
+       * If any of the current main table columns are missing in the installed main table, prompt the user to recreate the main table:
+            - If the user confirms, recreate the main table using the main DDL file, else log an error message and exit
+    :param installation:
+    :param ws:
+    :param app_context:
+    """
+    reconcile_config = app_context.recon_config
+    assert reconcile_config, "Reconcile config must be present to upgrade the reconcile metadata main table"
+    table_name = "main"
+    table_identifier = (
+        f"{reconcile_config.metadata_config.catalog}.{reconcile_config.metadata_config.schema}.{table_name}"
+    )
+    installed_columns = installed_table_columns(ws, table_identifier)
+    current_columns = current_table_columns(table_name, table_identifier)
+    sql: str | None = f"ALTER TABLE {table_identifier} ADD COLUMN operation_name  STRING AFTER report_type"
+    if _check_table_mismatch(installed_columns, current_columns):
+        logger.info("Recreating main table")
+        sql = recreate_table_sql(table_identifier, installed_columns, current_columns, app_context.prompts)
+    if sql:
+        logger.debug(f"Executing SQL to upgrade main table: \n{sql}")
+        db_sql.get_sql_backend(ws).execute(sql)
+        installation.save(reconcile_config)
+        logger.debug("Upgraded Reconcile main table")
+
+
+def _upgrade_reconcile_workflow(app_context: ApplicationContext):
+    if app_context.recon_config:
+        logger.info("Upgrading reconcile workflow")
+        wheels = app_context.product_info.wheels(app_context.workspace_client)
+        with wheels as wheel_builder:
+            wheel_path = f"/Workspace{wheel_builder.upload_to_wsfs()}"
+        app_context.job_deployment.deploy_recon_job(RECON_JOB_NAME, app_context.recon_config, wheel_path)
+        logger.debug("Upgraded reconcile workflow")
+
+
+def upgrade(installation: Installation, ws: WorkspaceClient):
+    app_context = ApplicationContext(ws)
+    if app_context.recon_config is not None:
+        _upgrade_reconcile_metadata_main_table(installation, ws, app_context)
+        _upgrade_reconcile_workflow(app_context)

databricks/labs/lakebridge/upgrades/v0.6.0_alter_metrics_datatype.py

@@ -0,0 +1,51 @@
+# pylint: disable=invalid-name
+import logging
+
+from databricks.labs.blueprint.installation import Installation
+from databricks.sdk import WorkspaceClient
+
+from databricks.labs.lakebridge.contexts.application import ApplicationContext
+from databricks.labs.lakebridge.deployment.upgrade_common import (
+    current_table_columns,
+    installed_table_columns,
+    check_table_mismatch,
+    recreate_table_sql,
+)
+from databricks.labs.lakebridge.helpers import db_sql
+
+logger = logging.getLogger(__name__)
+
+
+def _upgrade_reconcile_metadata_metrics_table(
+    installation: Installation, ws: WorkspaceClient, app_context: ApplicationContext
+):
+    reconcile_config = app_context.recon_config
+    assert reconcile_config, "Reconcile config must be present to upgrade the reconcile metadata main table"
+    table_name = "metrics"
+    table_identifier = (
+        f"{reconcile_config.metadata_config.catalog}.{reconcile_config.metadata_config.schema}.{table_name}"
+    )
+    installed_columns = installed_table_columns(ws, table_identifier)
+    current_columns = current_table_columns(table_name, table_identifier)
+    sqls: list | None = [
+        f"ALTER TABLE {table_identifier} SET TBLPROPERTIES ('delta.enableTypeWidening' = 'true')",
+        f"ALTER TABLE {table_identifier} ALTER COLUMN recon_metrics.row_comparison.missing_in_source TYPE BIGINT",
+        f"ALTER TABLE {table_identifier} ALTER COLUMN recon_metrics.row_comparison.missing_in_target TYPE BIGINT",
+        f"ALTER TABLE {table_identifier} ALTER COLUMN recon_metrics.column_comparison.absolute_mismatch TYPE BIGINT",
+        f"ALTER TABLE {table_identifier} ALTER COLUMN recon_metrics.column_comparison.threshold_mismatch TYPE BIGINT",
+    ]
+    if check_table_mismatch(installed_columns, current_columns):
+        logger.info("Recreating main table")
+        sqls = [recreate_table_sql(table_identifier, installed_columns, current_columns, app_context.prompts)]
+    if sqls:
+        for sql in sqls:
+            logger.debug(f"Executing SQL to upgrade metrics table: \n{sql}")
+            db_sql.get_sql_backend(ws).execute(sql)
+        installation.save(reconcile_config)
+        logger.debug("Upgraded Reconcile metrics table")
+
+
+def upgrade(installation: Installation, ws: WorkspaceClient):
+    app_context = ApplicationContext(ws)
+    if app_context.recon_config is not None:
+        _upgrade_reconcile_metadata_metrics_table(installation, ws, app_context)

databricks_labs_lakebridge-0.10.0.dist-info/METADATA

@@ -0,0 +1,58 @@
+Metadata-Version: 2.4
+Name: databricks-labs-lakebridge
+Version: 0.10.0
+Summary: Fast and predictable migrations to Databricks Lakehouse Platform. This tool is designed to help you migrate your data and workloads to the Databricks Lakehouse Platform in a fast, predictable, and reliable way. It provides a set of tools and utilities to help you reconcile your data and workloads, assess your current state, and plan your migration.
+Project-URL: Documentation, https://github.com/databrickslabs/lakebridge
+Project-URL: Issues, https://github.com/databrickslabs/lakebridge/issues
+Project-URL: Source, https://github.com/databrickslabs/lakebridge
+Maintainer-email: Databricks Labs <labs-oss@databricks.com>
+License-File: LICENSE
+License-File: NOTICE
+Keywords: Databricks
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: Framework :: Pytest
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: System Administrators
+Classifier: License :: Other/Proprietary License
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Utilities
+Requires-Python: >=3.10
+Requires-Dist: cryptography<45.1.0,>=44.0.2
+Requires-Dist: databricks-bb-analyzer~=0.1.6
+Requires-Dist: databricks-labs-blueprint[yaml]<0.12.0,>=0.11.0
+Requires-Dist: databricks-labs-lsql==0.16.0
+Requires-Dist: databricks-sdk~=0.51.0
+Requires-Dist: duckdb~=1.2.2
+Requires-Dist: pygls~=2.0.0a2
+Requires-Dist: pyodbc~=5.2.0
+Requires-Dist: sqlalchemy~=2.0.40
+Requires-Dist: sqlglot==26.1.3
+Requires-Dist: standard-distutils~=3.11.9; python_version >= '3.11'
+Description-Content-Type: text/markdown
+
+Databricks Labs Lakebridge
+---
+![Databricks Labs Lakebridge White](/docs/lakebridge/static/img/lakebridge-lockup-white-background.svg)
+
+
+[![build](https://github.com/databrickslabs/remorph/actions/workflows/push.yml/badge.svg)](https://github.com/databrickslabs/remorph/actions/workflows/push.yml)
+![PyPI - Downloads](https://img.shields.io/pypi/dm/databricks-labs-remorph?cacheSeconds=3600)
+
+-----
+Documentation
+The complete documentation is available at: https://databrickslabs.github.io/lakebridge/
+
+Contribution
+Please see the contribution guidance here on how to contribute to the project (build, test, and submit a PR).
+
+Project Support
+Please note that this project is provided for your exploration only and is not formally supported by Databricks with Service Level Agreements (SLAs). They are provided AS-IS, and we do not make any guarantees. Please do not submit a support ticket relating to any issues arising from the use of this project.
+
+Any issues discovered through the use of this project should be filed as GitHub Issues on this repository. They will be reviewed as time permits, but no formal SLAs for support exist.