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

databricks/labs/lakebridge/connections/credential_manager.py

@@ -0,0 +1,89 @@
+from pathlib import Path
+import logging
+from typing import Protocol
+
+import yaml
+
+from databricks.labs.lakebridge.connections.env_getter import EnvGetter
+
+
+logger = logging.getLogger(__name__)
+
+
+class SecretProvider(Protocol):
+    def get_secret(self, key: str) -> str:
+        pass
+
+
+class LocalSecretProvider:
+    def get_secret(self, key: str) -> str:
+        return key
+
+
+class EnvSecretProvider:
+    def __init__(self, env_getter: EnvGetter):
+        self._env_getter = env_getter
+
+    def get_secret(self, key: str) -> str:
+        try:
+            return self._env_getter.get(str(key))
+        except KeyError:
+            logger.debug(f"Environment variable {key} not found. Falling back to actual value")
+            return key
+
+
+class DatabricksSecretProvider:
+    def get_secret(self, key: str) -> str:
+        raise NotImplementedError("Databricks secret vault not implemented")
+
+
+class CredentialManager:
+    def __init__(self, credential_loader: dict, secret_providers: dict):
+        self._credentials = credential_loader
+        self._secret_providers = secret_providers
+        self._default_vault = self._credentials.get('secret_vault_type', 'local').lower()
+
+    def get_credentials(self, source: str) -> dict:
+        if source not in self._credentials:
+            raise KeyError(f"Source system: {source} credentials not found")
+
+        value = self._credentials[source]
+        if not isinstance(value, dict):
+            raise KeyError(f"Invalid credential format for source: {source}")
+
+        return {k: self._get_secret_value(v) for k, v in value.items()}
+
+    def _get_secret_value(self, key: str) -> str:
+        provider = self._secret_providers.get(self._default_vault)
+        if not provider:
+            raise ValueError(f"Unsupported secret vault type: {self._default_vault}")
+        return provider.get_secret(key)
+
+
+def _get_home() -> Path:
+    return Path(__file__).home()
+
+
+def cred_file(product_name) -> Path:
+    return Path(f"{_get_home()}/.databricks/labs/{product_name}/.credentials.yml")
+
+
+def _load_credentials(path: Path) -> dict:
+    try:
+        with open(path, encoding="utf-8") as f:
+            return yaml.safe_load(f)
+    except FileNotFoundError as e:
+        raise FileNotFoundError(f"Credentials file not found at {path}") from e
+
+
+def create_credential_manager(product_name: str, env_getter: EnvGetter):
+    file_path = Path(f"{_get_home()}/.databricks/labs/{product_name}/.credentials.yml")
+
+    secret_providers = {
+        'local': LocalSecretProvider(),
+        'env': EnvSecretProvider(env_getter),
+        'databricks': DatabricksSecretProvider(),
+    }
+
+    loader = _load_credentials(file_path)
+    return CredentialManager(loader, secret_providers)

databricks/labs/lakebridge/connections/database_manager.py

@@ -0,0 +1,98 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine import Engine, Result, URL
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError
+
+logger = logging.getLogger(__name__)
+logger.setLevel("INFO")
+
+
+class DatabaseConnector(ABC):
+    @abstractmethod
+    def _connect(self) -> Engine:
+        pass
+
+    @abstractmethod
+    def execute_query(self, query: str) -> Result[Any]:
+        pass
+
+
+class _BaseConnector(DatabaseConnector):
+    def __init__(self, config: dict[str, Any]):
+        self.config = config
+        self.engine: Engine = self._connect()
+
+    def _connect(self) -> Engine:
+        raise NotImplementedError("Subclasses should implement this method")
+
+    def execute_query(self, query: str) -> Result[Any]:
+        if not self.engine:
+            raise ConnectionError("Not connected to the database.")
+        session = sessionmaker(bind=self.engine)
+        connection = session()
+        return connection.execute(text(query))
+
+
+def _create_connector(db_type: str, config: dict[str, Any]) -> DatabaseConnector:
+    connectors = {
+        "snowflake": SnowflakeConnector,
+        "mssql": MSSQLConnector,
+        "tsql": MSSQLConnector,
+        "synapse": MSSQLConnector,
+    }
+
+    connector_class = connectors.get(db_type.lower())
+
+    if connector_class is None:
+        raise ValueError(f"Unsupported database type: {db_type}")
+
+    return connector_class(config)
+
+
+class SnowflakeConnector(_BaseConnector):
+    def _connect(self) -> Engine:
+        raise NotImplementedError("Snowflake connector not implemented")
+
+
+class MSSQLConnector(_BaseConnector):
+    def _connect(self) -> Engine:
+        query_params = {"driver": self.config['driver']}
+
+        for key, value in self.config.items():
+            if key not in ["user", "password", "server", "database", "port"]:
+                query_params[key] = value
+        connection_string = URL.create(
+            "mssql+pyodbc",
+            username=self.config['user'],
+            password=self.config['password'],
+            host=self.config['server'],
+            port=self.config.get('port', 1433),
+            database=self.config['database'],
+            query=query_params,
+        )
+        return create_engine(connection_string)
+
+
+class DatabaseManager:
+    def __init__(self, db_type: str, config: dict[str, Any]):
+        self.connector = _create_connector(db_type, config)
+
+    def execute_query(self, query: str) -> Result[Any]:
+        try:
+            return self.connector.execute_query(query)
+        except OperationalError:
+            logger.error("Error connecting to the database check credentials")
+            raise ConnectionError("Error connecting to the database check credentials") from None
+
+    def check_connection(self) -> bool:
+        query = "SELECT 101 AS test_column"
+        result = self.execute_query(query)
+        row = result.fetchone()
+        if row is None:
+            return False
+        return row[0] == 101

databricks/labs/lakebridge/connections/env_getter.py

@@ -0,0 +1,13 @@
+import os
+
+
+class EnvGetter:
+    """Standardised inorder to support testing Capabilities, check debug_envgetter.py"""
+
+    def __init__(self):
+        self.env = dict(os.environ)
+
+    def get(self, key: str) -> str:
+        if key in self.env:
+            return self.env[key]
+        raise KeyError(f"not in env: {key}")

databricks/labs/lakebridge/contexts/application.py

@@ -0,0 +1,133 @@
+import logging
+from functools import cached_property
+
+from databricks.labs.blueprint.installation import Installation
+from databricks.labs.blueprint.installer import InstallState
+from databricks.labs.blueprint.upgrades import Upgrades
+from databricks.labs.blueprint.tui import Prompts
+from databricks.labs.blueprint.wheels import ProductInfo
+from databricks.labs.lsql.backends import SqlBackend, StatementExecutionBackend
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.config import Config
+from databricks.sdk.errors import NotFound
+from databricks.sdk.service.iam import User
+
+from databricks.labs.lakebridge.config import TranspileConfig, ReconcileConfig, RemorphConfigs
+from databricks.labs.lakebridge.deployment.configurator import ResourceConfigurator
+from databricks.labs.lakebridge.deployment.dashboard import DashboardDeployment
+from databricks.labs.lakebridge.deployment.installation import WorkspaceInstallation
+from databricks.labs.lakebridge.deployment.recon import TableDeployment, JobDeployment, ReconDeployment
+from databricks.labs.lakebridge.helpers.metastore import CatalogOperations
+
+logger = logging.getLogger(__name__)
+
+
+class ApplicationContext:
+    def __init__(self, ws: WorkspaceClient):
+        self._ws = ws
+
+    def replace(self, **kwargs):
+        """Replace cached properties for unit testing purposes."""
+        for key, value in kwargs.items():
+            self.__dict__[key] = value
+        return self
+
+    @cached_property
+    def workspace_client(self) -> WorkspaceClient:
+        return self._ws
+
+    @cached_property
+    def current_user(self) -> User:
+        return self.workspace_client.current_user.me()
+
+    @cached_property
+    def product_info(self) -> ProductInfo:
+        return ProductInfo.from_class(RemorphConfigs)
+
+    @cached_property
+    def installation(self) -> Installation:
+        return Installation.assume_user_home(self.workspace_client, self.product_info.product_name())
+
+    @cached_property
+    def transpile_config(self) -> TranspileConfig | None:
+        try:
+            return self.installation.load(TranspileConfig)
+        except NotFound as err:
+            logger.debug(f"Couldn't find existing `transpile` installation: {err}")
+            return None
+
+    @cached_property
+    def recon_config(self) -> ReconcileConfig | None:
+        try:
+            return self.installation.load(ReconcileConfig)
+        except NotFound as err:
+            logger.debug(f"Couldn't find existing `reconcile` installation: {err}")
+            return None
+
+    @cached_property
+    def remorph_config(self) -> RemorphConfigs:
+        return RemorphConfigs(transpile=self.transpile_config, reconcile=self.recon_config)
+
+    @cached_property
+    def connect_config(self) -> Config:
+        return self.workspace_client.config
+
+    @cached_property
+    def install_state(self) -> InstallState:
+        return InstallState.from_installation(self.installation)
+
+    @cached_property
+    def sql_backend(self) -> SqlBackend:
+        # Installer to use only StatementExecutionBackend to eliminate the need for Databricks Connect
+        return StatementExecutionBackend(self.workspace_client, self.connect_config.warehouse_id)
+
+    @cached_property
+    def catalog_operations(self) -> CatalogOperations:
+        return CatalogOperations(self.workspace_client)
+
+    @cached_property
+    def prompts(self) -> Prompts:
+        return Prompts()
+
+    @cached_property
+    def resource_configurator(self) -> ResourceConfigurator:
+        return ResourceConfigurator(self.workspace_client, self.prompts, self.catalog_operations)
+
+    @cached_property
+    def table_deployment(self) -> TableDeployment:
+        return TableDeployment(self.sql_backend)
+
+    @cached_property
+    def job_deployment(self) -> JobDeployment:
+        return JobDeployment(self.workspace_client, self.installation, self.install_state, self.product_info)
+
+    @cached_property
+    def dashboard_deployment(self) -> DashboardDeployment:
+        return DashboardDeployment(self.workspace_client, self.installation, self.install_state)
+
+    @cached_property
+    def recon_deployment(self) -> ReconDeployment:
+        return ReconDeployment(
+            self.workspace_client,
+            self.installation,
+            self.install_state,
+            self.product_info,
+            self.table_deployment,
+            self.job_deployment,
+            self.dashboard_deployment,
+        )
+
+    @cached_property
+    def workspace_installation(self) -> WorkspaceInstallation:
+        return WorkspaceInstallation(
+            self.workspace_client,
+            self.prompts,
+            self.installation,
+            self.recon_deployment,
+            self.product_info,
+            self.upgrades,
+        )
+
+    @cached_property
+    def upgrades(self):
+        return Upgrades(self.product_info, self.installation)

databricks/labs/lakebridge/coverage/commons.py

@@ -0,0 +1,223 @@
+# pylint: disable=all
+import collections
+import dataclasses
+import json
+import logging
+import os
+import subprocess
+import time
+from collections.abc import Generator
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TextIO, List
+
+import sqlglot
+from sqlglot.expressions import Expression
+from sqlglot.dialects.dialect import Dialect
+from sqlglot.dialects.databricks import Databricks
+from sqlglot.errors import ErrorLevel
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class ReportEntry:
+    project: str
+    commit_hash: str | None
+    version: str
+    timestamp: str
+    source_dialect: str
+    target_dialect: str
+    file: str
+    parsed: int = 0  # 1 for success, 0 for failure
+    statements: int = 0  # number of statements parsed
+    transpiled: int = 0  # 1 for success, 0 for failure
+    transpiled_statements: int = 0  # number of statements transpiled
+    failures: List[dict] = dataclasses.field(default_factory=lambda: [])
+
+
+def sqlglot_run_coverage(dialect, subfolder):
+    input_dir = get_env_var("INPUT_DIR_PARENT", required=True)
+    output_dir = get_env_var("OUTPUT_DIR", required=True)
+    sqlglot_version = sqlglot.__version__
+    SQLGLOT_COMMIT_HASH = ""  # C0103 pylint
+
+    if not input_dir:
+        raise ValueError("Environment variable `INPUT_DIR_PARENT` is required")
+    if not output_dir:
+        raise ValueError("Environment variable `OUTPUT_DIR` is required")
+
+    collect_transpilation_stats(
+        "SQLGlot",
+        SQLGLOT_COMMIT_HASH,
+        sqlglot_version,
+        dialect,
+        Databricks,
+        Path(input_dir) / subfolder,
+        Path(output_dir),
+    )
+
+
+def local_report(output_dir: Path):
+    all = collections.defaultdict(list)
+    for file in output_dir.rglob("*.json"):
+        with file.open("r", encoding="utf8") as f:
+            for line in f:
+                raw = json.loads(line)
+                entry = ReportEntry(**raw)
+                all[(entry.project, entry.source_dialect)].append(entry)
+    for (project, dialect), entries in sorted(all.items()):
+        total = len(entries)
+        parsed = sum(entry.parsed for entry in entries)
+        transpiled = sum(entry.transpiled for entry in entries)
+        parse_ratio = parsed / total
+        transpile_ratio = transpiled / total
+        print(
+            f"{project} -> {dialect}: {parse_ratio:.2%} parsed ({parsed}/{total}), "
+            f"{transpile_ratio:.2%} transpiled ({transpiled}/{total})"
+        )
+
+
+def get_supported_sql_files(input_dir: Path) -> Generator[Path, None, None]:
+    yield from filter(lambda item: item.is_file() and item.suffix.lower() in [".sql", ".ddl"], input_dir.rglob("*"))
+
+
+def write_json_line(file: TextIO, content: ReportEntry):
+    json.dump(dataclasses.asdict(content), file)
+    file.write("\n")
+
+
+def get_env_var(env_var: str, *, required: bool = False) -> str | None:
+    """
+    Get the value of an environment variable.
+
+    :param env_var: The name of the environment variable to get the value of.
+    :param required: Indicates if the environment variable is required and raises a ValueError if it's not set.
+    :return: Returns the environment variable's value, or None if it's not set and not required.
+    """
+    value = os.getenv(env_var)
+    if value is None and required:
+        message = f"Environment variable {env_var} is not set"
+        raise ValueError(message)
+    return value
+
+
+def get_current_commit_hash() -> str | None:
+    try:
+        return (
+            subprocess.check_output(
+                ["/usr/bin/git", "rev-parse", "--short", "HEAD"],
+                cwd=Path(__file__).resolve().parent,
+            )
+            .decode("ascii")
+            .strip()
+        )
+    except (subprocess.CalledProcessError, FileNotFoundError) as e:
+        logger.warning(f"Could not get the current commit hash. {e!s}")
+        return None
+
+
+def get_current_time_utc() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def parse_sql(sql: str, dialect: type[Dialect]) -> list[Expression]:
+    return [
+        expression for expression in sqlglot.parse(sql, read=dialect, error_level=ErrorLevel.IMMEDIATE) if expression
+    ]
+
+
+def generate_sql(expressions: list[Expression], dialect: type[Dialect]) -> list[str]:
+    generator_dialect = Dialect.get_or_raise(dialect)
+    return [generator_dialect.generate(expression, copy=False) for expression in expressions if expression]
+
+
+def _ensure_valid_io_paths(input_dir: Path, result_dir: Path):
+    if not input_dir.exists() or not input_dir.is_dir():
+        message = f"The input path {input_dir} doesn't exist or is not a directory"
+        raise NotADirectoryError(message)
+
+    if not result_dir.exists():
+        logger.info(f"Creating the output directory {result_dir}")
+        result_dir.mkdir(parents=True)
+    elif not result_dir.is_dir():
+        message = f"The output path {result_dir} exists but is not a directory"
+        raise NotADirectoryError(message)
+
+
+def _get_report_file_path(
+    project: str,
+    source_dialect: type[Dialect],
+    target_dialect: type[Dialect],
+    result_dir: Path,
+) -> Path:
+    source_dialect_name = source_dialect.__name__
+    target_dialect_name = target_dialect.__name__
+    current_time_ns = time.time_ns()
+    return result_dir / f"{project}_{source_dialect_name}_{target_dialect_name}_{current_time_ns}.json".lower()
+
+
+def _prepare_report_entry(
+    project: str,
+    commit_hash: str,
+    version: str,
+    source_dialect: type[Dialect],
+    target_dialect: type[Dialect],
+    file_path: str,
+    sql: str,
+) -> ReportEntry:
+    report_entry = ReportEntry(
+        project=project,
+        commit_hash=commit_hash,
+        version=version,
+        timestamp=get_current_time_utc().isoformat(),
+        source_dialect=source_dialect.__name__,
+        target_dialect=target_dialect.__name__,
+        file=file_path,
+    )
+    try:
+        expressions = parse_sql(sql, source_dialect)
+        report_entry.parsed = 1
+        report_entry.statements = len(expressions)
+    except Exception as pe:
+        report_entry.failures.append({'error_code': type(pe).__name__, 'error_message': repr(pe)})
+        return report_entry
+
+    try:
+        generated_sqls = generate_sql(expressions, target_dialect)
+        report_entry.transpiled = 1
+        report_entry.transpiled_statements = len([sql for sql in generated_sqls if sql.strip()])
+    except Exception as te:
+        report_entry.failures.append({'error_code': type(te).__name__, 'error_message': repr(te)})
+
+    return report_entry
+
+
+def collect_transpilation_stats(
+    project: str,
+    commit_hash: str,
+    version: str,
+    source_dialect: type[Dialect],
+    target_dialect: type[Dialect],
+    input_dir: Path,
+    result_dir: Path,
+):
+    _ensure_valid_io_paths(input_dir, result_dir)
+    report_file_path = _get_report_file_path(project, source_dialect, target_dialect, result_dir)
+
+    with report_file_path.open("w", encoding="utf8") as report_file:
+        for input_file in get_supported_sql_files(input_dir):
+            with input_file.open("r", encoding="utf-8-sig") as file:
+                sql = file.read()
+
+            file_path = str(input_file.absolute().relative_to(input_dir.parent.absolute()))
+            report_entry = _prepare_report_entry(
+                project,
+                commit_hash,
+                version,
+                source_dialect,
+                target_dialect,
+                file_path,
+                sql,
+            )
+            write_json_line(report_file, report_entry)

databricks/labs/lakebridge/coverage/lakebridge_snow_transpilation_coverage.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+from databricks.labs.blueprint.wheels import ProductInfo
+from databricks.labs.lakebridge.coverage import commons
+from databricks.labs.lakebridge.transpiler.sqlglot.generator.databricks import Databricks
+from databricks.labs.lakebridge.transpiler.sqlglot.parsers.snowflake import Snowflake
+
+if __name__ == "__main__":
+    input_dir = commons.get_env_var("INPUT_DIR_PARENT", required=True)
+    output_dir = commons.get_env_var("OUTPUT_DIR", required=True)
+
+    REMORPH_COMMIT_HASH = commons.get_current_commit_hash() or ""  # C0103 pylint
+    product_info = ProductInfo(__file__)
+    remorph_version = product_info.unreleased_version()
+
+    if not input_dir:
+        raise ValueError("Environment variable `INPUT_DIR_PARENT` is required")
+    if not output_dir:
+        raise ValueError("Environment variable `OUTPUT_DIR` is required")
+
+    commons.collect_transpilation_stats(
+        "Remorph",
+        REMORPH_COMMIT_HASH,
+        remorph_version,
+        Snowflake,
+        Databricks,
+        Path(input_dir) / 'snowflake',
+        Path(output_dir),
+    )

databricks/labs/lakebridge/coverage/local_report.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+
+from databricks.labs.lakebridge.coverage import commons
+
+if __name__ == "__main__":
+    output_dir = commons.get_env_var("OUTPUT_DIR", required=True)
+    if not output_dir:
+        raise ValueError("Environment variable `OUTPUT_DIR` is required")
+    commons.local_report(Path(output_dir))

databricks/labs/lakebridge/coverage/sqlglot_snow_transpilation_coverage.py

@@ -0,0 +1,5 @@
+from sqlglot.dialects.snowflake import Snowflake
+from databricks.labs.lakebridge.coverage.commons import sqlglot_run_coverage
+
+if __name__ == "__main__":
+    sqlglot_run_coverage(Snowflake, "snowflake")

databricks/labs/lakebridge/coverage/sqlglot_tsql_transpilation_coverage.py

@@ -0,0 +1,5 @@
+from sqlglot.dialects.tsql import TSQL
+from databricks.labs.lakebridge.coverage.commons import sqlglot_run_coverage
+
+if __name__ == "__main__":
+    sqlglot_run_coverage(TSQL, "tsql")