Initial Spec For Credential Management and SQLAlchemy Database Connectors

.github/scripts/setup_mssql_odbc.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+
+set -xve
+#Repurposed from https://github.com/Yarden-zamir/install-mssql-odbc
+
+curl -sSL -O https://packages.microsoft.com/config/ubuntu/$(grep VERSION_ID /etc/os-release | cut -d '"' -f 2)/packages-microsoft-prod.deb
+
+sudo dpkg -i packages-microsoft-prod.deb
+#rm packages-microsoft-prod.deb
+
+sudo apt-get update
+sudo ACCEPT_EULA=Y apt-get install -y msodbcsql18
+

.github/workflows/acceptance.yml

@@ -0,0 +1,58 @@
+name: acceptance
+
+on:
+  pull_request:
+    types: [ opened, synchronize, ready_for_review ]
+  merge_group:
+    types: [ checks_requested ]
+  push:
+    branches:
+      - main
+
+permissions:
+  id-token: write
+  contents: read
+  pull-requests: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  integration:
+    if: github.event_name == 'pull_request' && github.event.pull_request.draft == false
+    environment: tool
+    runs-on: larger
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Install Python
+        uses: actions/setup-python@v5
+        with:
+          cache: 'pip'
+          cache-dependency-path: '**/pyproject.toml'
+          python-version: '3.10'
+
+      - name: Install hatch
+        run: pip install hatch==1.9.4
+
+      - name: Install MSSQL ODBC Driver
+        run: |
+          chmod +x $GITHUB_WORKSPACE/.github/scripts/setup_mssql_odbc.sh
+          $GITHUB_WORKSPACE/.github/scripts/setup_mssql_odbc.sh
+
+      - name: Run integration tests
+        uses: databrickslabs/sandbox/acceptance@acceptance/v0.4.2
+        with:
+          vault_uri: ${{ secrets.VAULT_URI }}
+          directory: ${{ github.workspace }}
+          timeout: 2h
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }}
+          ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }}
+          TEST_ENV: 'ACCEPTANCE'
+

Makefile

@@ -17,10 +17,10 @@ fmt:
 setup_spark_remote:
 	.github/scripts/setup_spark_remote.sh
 
-test: setup_spark_remote
+test:
 	hatch run test
 
-integration:
+integration: setup_spark_remote
 	hatch run integration
 
 coverage:

labs.yml

@@ -3,9 +3,7 @@ name: remorph
 description: Code Transpiler and Data Reconciliation tool for Accelerating Data onboarding to Databricks from EDW, CDW and other ETL sources.
 install:
   min_runtime_version: 13.3
-  require_running_cluster: false
-  require_databricks_connect: false
-  script: src/databricks/labs/remorph/install.py
+  script: src/databricks/labs/remorph/base_install.py
 uninstall:
   script: src/databricks/labs/remorph/uninstall.py
 entrypoint: src/databricks/labs/remorph/cli.py
@@ -66,3 +64,9 @@ commands:
     description: Utility to setup Scope and Secrets on Databricks Workspace
   - name: debug-me
     description: "[INTERNAL] Debug SDK connectivity"
+  - name: install-assessment
+    description: "Install Assessment"
+  - name: install-transpile
+    description: "Install Transpile"
+  - name: install-reconcile
+    description: "Install Reconcile"

pyproject.toml

@@ -18,7 +18,10 @@ dependencies = [
   "databricks-labs-blueprint[yaml]>=0.2.3",
   "databricks-labs-lsql>=0.7.5,<0.14.0",  # TODO: Limit the LSQL version until dependencies are correct.
   "cryptography>=41.0.3",
+  "pyodbc",
+  "SQLAlchemy",
   "pygls>=2.0.0a2",
+
 ]
 
 [project.urls]
@@ -50,6 +53,7 @@ dependencies = [
   "pytest",
   "pytest-cov>=5.0.0,<6.0.0",
   "pytest-asyncio>=0.24.0",
+    "pytest-xdist~=3.5.0",
   "black>=23.1.0",
   "ruff>=0.0.243",
   "databricks-connect==15.1",
@@ -66,7 +70,7 @@ reconcile = "databricks.labs.remorph.reconcile.execute:main"
 [tool.hatch.envs.default.scripts]
 test        = "pytest --cov src --cov-report=xml tests/unit"
 coverage    = "pytest --cov src tests/unit --cov-report=html"
-integration = "pytest --cov src tests/integration --durations 20"
+integration = "pytest --cov src tests/integration --durations 20 --ignore=tests/integration/connections"
 fmt         = ["black .",
                "ruff  check . --fix",
                "mypy --disable-error-code 'annotation-unchecked' .",

src/databricks/labs/remorph/base_install.py

@@ -0,0 +1,9 @@
+from databricks.labs.blueprint.entrypoint import get_logger, is_in_debug
+
+if __name__ == "__main__":
+    logger = get_logger(__file__)
+    logger.setLevel("INFO")
+    if is_in_debug():
+        logger.getLogger("databricks").setLevel(logger.setLevel("DEBUG"))
+
+    logger.info("Successfully Setup Remorph Components Locally")

src/databricks/labs/remorph/cli.py

@@ -8,6 +8,8 @@
 from databricks.labs.remorph.config import TranspileConfig
 from databricks.labs.remorph.contexts.application import ApplicationContext
 from databricks.labs.remorph.helpers.recon_config_utils import ReconConfigPrompts
+from databricks.labs.remorph.__about__ import __version__
+from databricks.labs.remorph.install import WorkspaceInstaller
 from databricks.labs.remorph.reconcile.runner import ReconcileRunner
 from databricks.labs.remorph.lineage import lineage_generator
 from databricks.labs.remorph.transpiler.execute import transpile as do_transpile
@@ -34,6 +36,32 @@ def raise_validation_exception(msg: str) -> Exception:
 proxy_command(remorph, "debug-bundle")
 
 
+def _installer(ws: WorkspaceClient) -> WorkspaceInstaller:
+    app_context = ApplicationContext(_verify_workspace_client(ws))
+    return WorkspaceInstaller(
+        app_context.workspace_client,
+        app_context.prompts,
+        app_context.installation,
+        app_context.install_state,
+        app_context.product_info,
+        app_context.resource_configurator,
+        app_context.workspace_installation,
+    )
+
+
+def _verify_workspace_client(ws: WorkspaceClient) -> WorkspaceClient:
+    """
+    [Private] Verifies and updates the workspace client configuration.
+    """
+
+    # Using reflection to set right value for _product_info for telemetry
+    product_info = getattr(ws.config, '_product_info')
+    if product_info[0] != "remorph":
+        setattr(ws.config, '_product_info', ('remorph', __version__))
+
+    return ws
+
+
 @remorph.command
 def transpile(
     w: WorkspaceClient,
@@ -168,5 +196,25 @@ def configure_secrets(w: WorkspaceClient):
     recon_conf.prompt_and_save_connection_details()
 
 
+@remorph.command(is_unauthenticated=True)
+def install_assessment():
+    """Install the Remorph Assessment package"""
+    raise NotImplementedError("Assessment package is not available yet.")
+
+
+@remorph.command()
+def install_transpile(w: WorkspaceClient):
+    """Install the Remorph Transpile package"""
+    installer = _installer(w)
+    installer.run(module="transpile")
+
+
+@remorph.command(is_unauthenticated=False)
+def install_reconcile(w: WorkspaceClient):
+    """Install the Remorph Reconcile package"""
+    installer = _installer(w)
+    installer.run(module="reconcile")
+
+
 if __name__ == "__main__":
     remorph()

src/databricks/labs/remorph/connections/credential_manager.py

@@ -0,0 +1,49 @@
+from pathlib import Path
+import logging
+import yaml
+
+from databricks.labs.blueprint.wheels import ProductInfo
+from databricks.labs.remorph.connections.env_getter import EnvGetter
+
+logger = logging.getLogger(__name__)
+
+
+class Credentials:
+    def __init__(self, product_info: ProductInfo, env: EnvGetter) -> None:
+        self._product_info = product_info
+        self._env = env
+        self._credentials: dict[str, str] = self._load_credentials(self._get_local_version_file_path())
+
+    def _get_local_version_file_path(self) -> Path:
+        user_home = f"{Path(__file__).home()}"
+        return Path(f"{user_home}/.databricks/labs/{self._product_info.product_name()}/credentials.yml")
+
+    def _load_credentials(self, file_path: Path) -> dict[str, str]:
+        with open(file_path, encoding="utf-8") as f:
+            return yaml.safe_load(f)
+
+    def load(self, source: str) -> dict[str, str]:
+        error_msg = f"source system: {source} credentials not found in file credentials.yml"
+        if source in self._credentials:
+            value = self._credentials[source]
+            if isinstance(value, dict):
+                return {k: self._get_secret_value(v) for k, v in value.items()}
+            raise KeyError(error_msg)
+        raise KeyError(error_msg)
+
+    def _get_secret_value(self, key: str) -> str:
+        secret_vault_type = self._credentials.get('secret_vault_type', 'local').lower()
+        if secret_vault_type == 'local':
+            return key
+        if secret_vault_type == 'env':
+            try:
+                value = self._env.get(str(key))  # Port numbers can be int
+            except KeyError:
+                logger.debug(f"Environment variable {key} not found Failing back to actual string value")
+                return key
+            return value
+
+        if secret_vault_type == 'databricks':
+            raise NotImplementedError("Databricks secret vault not implemented")
+
+        raise ValueError(f"Unsupported secret vault type: {secret_vault_type}")

src/databricks/labs/remorph/connections/database_manager.py

@@ -0,0 +1,88 @@
+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__)
+
+
+class _ISourceSystemConnector(ABC):
+    @abstractmethod
+    def _connect(self) -> Engine:
+        pass
+
+    @abstractmethod
+    def execute_query(self, query: str) -> Result[Any]:
+        pass
+
+
+class _BaseConnector(_ISourceSystemConnector):
+    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.")
+        try:
+            session = sessionmaker(bind=self.engine)
+            connection = session()
+            return connection.execute(text(query))
+        except OperationalError:
+            raise ConnectionError("Error connecting to the database check credentials") from None
+
+
+def _create_connector(db_type: str, config: dict[str, Any]) -> _ISourceSystemConnector:
+    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]:
+        return self.connector.execute_query(query)

src/databricks/labs/remorph/connections/env_getter.py

@@ -0,0 +1,24 @@
+import os
+import json
+import logging
+
+
+class EnvGetter:
+    def __init__(self, is_debug: bool = False):
+        self.env = self._get_debug_env() if is_debug else 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}")
+
+    def _get_debug_env(self) -> dict:
+        try:
+            debug_env_file = f"{os.path.expanduser('~')}/.databricks/debug-env.json"
+            with open(debug_env_file, 'r', encoding='utf-8') as file:
+                contents = file.read()
+            logging.debug(f"Found debug env file: {debug_env_file}")
+            raw = json.loads(contents)
+            return raw.get("ucws", {})
+        except FileNotFoundError:
+            return dict(os.environ)