Profiler CLI

labs.yml

@@ -61,3 +61,7 @@ commands:
         default: null
   - name: configure-reconcile
     description: "Configure Necessary Reconcile Dependencies"
+  - name: execute-databse-profiler
+    description: "[EXPERIMENTAL] Profile the source system database"
+
+

pyproject.toml

@@ -114,7 +114,7 @@ asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope="function"
 
 [tool.mypy]
-exclude = ["tests/resources/.*"]
+exclude = ["tests/resources/.*", "src/databricks/labs/remorph/resources/assessments/*"]
 
 [tool.black]
 target-version = ["py310"]
@@ -147,7 +147,8 @@ lint.ignore = [
   "UP038", # Use `X | Y` in `isinstance` call instead of `(X, Y)`
 ]
 extend-exclude = [
-  "notebooks/*.py"
+  "notebooks/*.py",
+  "src/databricks/labs/remorph/resources/assessments/*",
 ]
 
 
@@ -204,7 +205,7 @@ fail-under = 10.0
 ignore-patterns = ["^\\.#"]
 
 # Ignore files under tests/resources
-ignore-paths = ["tests/resources"]
+ignore-paths = ["tests/resources", "src/databricks/labs/remorph/resources/assessments/*" ]
 
 # List of module names for which member attributes should not be checked (useful
 # for modules/projects where namespaces are manipulated during runtime and thus

src/databricks/labs/lakebridge/assessments/configure_assessment.py

@@ -123,6 +123,9 @@ def _configure_credentials(self) -> str:
             "sql_user": self.prompts.question("Enter SQL user"),
             "sql_password": self.prompts.question("Enter SQL password"),
             "tz_info": self.prompts.question("Enter timezone (e.g. America/New_York)", default="UTC"),
+            "driver": self.prompts.question(
+                "Enter the ODBC driver installed locally", default="ODBC Driver 18 for SQL Server"
+            ),
         }
 
         # Azure API Access Settings

src/databricks/labs/lakebridge/assessments/pipeline.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from subprocess import run, CalledProcessError
+from subprocess import run, CalledProcessError, Popen, PIPE, STDOUT, DEVNULL
 from dataclasses import dataclass
 from enum import Enum
 
@@ -35,15 +35,17 @@ class StepExecutionResult:
 
 
 class PipelineClass:
-    def __init__(self, config: PipelineConfig, executor: DatabaseManager):
+    def __init__(self, config: PipelineConfig, executor: DatabaseManager | None):
         self.config = config
         self.executor = executor
         self.db_path_prefix = Path(config.extract_folder)
+        self._create_dir(self.db_path_prefix)
 
     def execute(self) -> list[StepExecutionResult]:
         logging.info(f"Pipeline initialized with config: {self.config.name}, version: {self.config.version}")
         execution_results: list[StepExecutionResult] = []
         for step in self.config.steps:
+            logging.info(f"Executing step: {step.name}")
             result = self._process_step(step)
             execution_results.append(result)
             logging.info(f"Step '{step.name}' completed with status: {result.status}")
@@ -55,7 +57,6 @@ def _process_step(self, step: Step) -> StepExecutionResult:
         if step.flag != "active":
             logging.info(f"Skipping step: {step.name} as it is not active")
             return StepExecutionResult(step_name=step.name, status=StepExecutionStatus.SKIPPED)
-
         logging.debug(f"Executing step: {step.name}")
         try:
             status = self._execute_step(step)
@@ -80,6 +81,10 @@ def _execute_sql_step(self, step: Step):
         with open(step.extract_source, 'r', encoding='utf-8') as file:
             query = file.read()
 
+        if self.executor is None:
+            logging.error("DatabaseManager executor is not set.")
+            raise RuntimeError("DatabaseManager executor is not set.")
+
         # Execute the query using the database manager
         logging.info(f"Executing query: {query}")
         try:
@@ -104,51 +109,65 @@ def _execute_python_step(self, step: Step):
             venv_python = venv_dir / "bin" / "python"
             venv_pip = venv_dir / "bin" / "pip"
 
-            logger.info(f"Creating a virtual environment for Python script execution: ${venv_dir}")
-            # Install dependencies in the virtual environment
+            logger.info(
+                f"Creating a virtual environment for Python script execution: ${venv_dir} for step: {step.name}"
+            )
             if step.dependencies:
-                logging.info(f"Installing dependencies: {', '.join(step.dependencies)}")
-                try:
-                    logging.debug("Upgrading local pip")
-                    run([str(venv_pip), "install", "--upgrade", "pip"], check=True, capture_output=True, text=True)
+                self._install_dependencies(venv_pip, step.dependencies)
 
-                    run([str(venv_pip), "install", *step.dependencies], check=True, capture_output=True, text=True)
-                except CalledProcessError as e:
-                    logging.error(f"Failed to install dependencies: {e.stderr}")
-                    raise RuntimeError(f"Failed to install dependencies: {e.stderr}") from e
+            self._run_python_script(venv_python, step.extract_source, db_path, credential_config)
+
+    @staticmethod
+    def _install_dependencies(venv_pip, dependencies):
+        logging.info(f"Installing dependencies: {', '.join(dependencies)}")
+        try:
+            logging.debug("Upgrading local pip")
+            run([str(venv_pip), "install", "--upgrade", "pip"], check=True, stdout=DEVNULL, stderr=DEVNULL)
+            run([str(venv_pip), "install", *dependencies], check=True, stdout=DEVNULL, stderr=DEVNULL)
+        except CalledProcessError as e:
+            logging.error(f"Failed to install dependencies: {e.stderr}")
+            raise RuntimeError(f"Failed to install dependencies: {e.stderr}") from e
 
-            # Execute the Python script using the virtual environment's Python interpreter
+    @staticmethod
+    def _run_python_script(venv_python, script_path, db_path, credential_config):
+        output_lines = []
+        try:
+            with Popen(
+                [
+                    str(venv_python),
+                    str(script_path),
+                    "--db-path",
+                    db_path,
+                    "--credential-config-path",
+                    credential_config,
+                ],
+                stdout=PIPE,
+                stderr=STDOUT,
+                text=True,
+                bufsize=1,
+            ) as process:
+                if process.stdout is not None:
+                    for line in process.stdout:
+                        logger.info(line.rstrip())
+                        output_lines.append(line)
+                process.wait()
+        except Exception as e:
+            logging.error(f"Python script failed: {str(e)}")
+            raise RuntimeError(f"Script execution failed: {str(e)}") from e
+
+        if output_lines:
             try:
-                result = run(
-                    [
-                        str(venv_python),
-                        str(step.extract_source),
-                        "--db-path",
-                        db_path,
-                        "--credential-config-path",
-                        credential_config,
-                    ],
-                    check=True,
-                    capture_output=True,
-                    text=True,
-                )
-
-                try:
-                    output = json.loads(result.stdout)
-                    if output["status"] == "success":
-                        logging.info(f"Python script completed: {output['message']}")
-                    else:
-                        raise RuntimeError(f"Script reported error: {output['message']}")
-                except json.JSONDecodeError:
-                    logging.info(f"Python script output: {result.stdout}")
-
-            except CalledProcessError as e:
-                error_msg = e.stderr
-                logging.error(f"Python script failed: {error_msg}")
-                raise RuntimeError(f"Script execution failed: {error_msg}") from e
+                output = json.loads(output_lines[-1])
+                if output.get("status") == "success":
+                    logging.info(f"Python script completed: {output['message']}")
+                else:
+                    raise RuntimeError(f"Script reported error: {output.get('message', 'Unknown error')}")
+            except json.JSONDecodeError:
+                logging.info("Could not parse script output as JSON.")
+        if process.returncode != 0:
+            raise RuntimeError(f"Script execution failed with exit code {process.returncode}")
 
     def _save_to_db(self, result, step_name: str, mode: str, batch_size: int = 1000):
-        self._create_dir(self.db_path_prefix)
         db_path = str(self.db_path_prefix / DB_NAME)
 
         with duckdb.connect(db_path) as conn:
@@ -179,7 +198,7 @@ def _create_dir(dir_path: Path):
             dir_path.mkdir(parents=True, exist_ok=True)
 
     @staticmethod
-    def load_config_from_yaml(file_path: str) -> PipelineConfig:
+    def load_config_from_yaml(file_path: str | Path) -> PipelineConfig:
         with open(file_path, 'r', encoding='utf-8') as file:
             data = yaml.safe_load(file)
         steps = [Step(**step) for step in data['steps']]

src/databricks/labs/lakebridge/assessments/profiler.py

@@ -0,0 +1,77 @@
+import logging
+from pathlib import Path
+
+from databricks.labs.lakebridge.assessments.pipeline import PipelineClass
+from databricks.labs.lakebridge.assessments.profiler_config import PipelineConfig
+from databricks.labs.lakebridge.connections.database_manager import DatabaseManager
+from databricks.labs.lakebridge.connections.credential_manager import (
+    create_credential_manager,
+)
+from databricks.labs.lakebridge.connections.env_getter import EnvGetter
+
+_PLATFORM_TO_SOURCE_TECHNOLOGY = {
+    "Synapse": "src/databricks/labs/lakebridge/resources/assessments/synapse/pipeline_config.yml",
+}
+
+_CONNECTOR_REQUIRED = {
+    "Synapse": False,
+}
+
+PRODUCT_NAME = "lakebridge"
+PRODUCT_PATH_PREFIX = Path(__file__).home() / ".databricks" / "labs" / PRODUCT_NAME / "lib"
+
+logger = logging.getLogger(__name__)
+
+
+class Profiler:
+
+    @classmethod
+    def supported_source_technologies(cls) -> list[str]:
+        return list(_PLATFORM_TO_SOURCE_TECHNOLOGY.keys())
+
+    @staticmethod
+    def path_modifier(config_file: str | Path) -> PipelineConfig:
+        # TODO: Make this work install during developer mode
+        config = PipelineClass.load_config_from_yaml(config_file)
+        for step in config.steps:
+            step.extract_source = f"{PRODUCT_PATH_PREFIX}/{step.extract_source}"
+        return config
+
+    def profile(self, platform: str, extractor: DatabaseManager | None = None):
+        config_path = _PLATFORM_TO_SOURCE_TECHNOLOGY.get(platform, None)
+        if not config_path:
+            raise ValueError(f"Unsupported platform: {platform}")
+        self._execute(platform, config_path, extractor)
+
+    def _setup_extractor(self, platform: str) -> DatabaseManager | None:
+        if not _CONNECTOR_REQUIRED[platform]:
+            return None
+        cred_manager = create_credential_manager(PRODUCT_NAME, EnvGetter())
+        connect_config = cred_manager.get_credentials(platform)
+        return DatabaseManager(platform, connect_config)
+
+    def _execute(self, platform: str, config_path: str, extractor=None):
+        try:
+            config_full_path = self._locate_config(config_path)
+            config = Profiler.path_modifier(config_full_path)
+
+            if extractor is None:
+                extractor = self._setup_extractor(platform)
+
+            results = PipelineClass(config, extractor).execute()
+
+            for result in results:
+                logger.info(f"Step: {result.step_name}, Status: {result.status}, Error: {result.error_message}")
+
+        except FileNotFoundError as e:
+            logging.error(f"Configuration file not found for source {platform}: {e}")
+            raise FileNotFoundError(f"Configuration file not found for source {platform}: {e}") from e
+        except Exception as e:
+            logging.error(f"Error executing pipeline for source {platform}: {e}")
+            raise RuntimeError(f"Pipeline execution failed for source {platform} : {e}") from e
+
+    def _locate_config(self, config_path: str) -> Path:
+        config_file = PRODUCT_PATH_PREFIX / config_path
+        if not config_file.exists():
+            raise FileNotFoundError(f"Configuration file not found: {config_file}")
+        return config_file

src/databricks/labs/lakebridge/cli.py

@@ -22,6 +22,7 @@
     create_assessment_configurator,
     PROFILER_SOURCE_SYSTEM,
 )
+from databricks.labs.lakebridge.assessments.profiler import Profiler
 
 from databricks.labs.lakebridge.__about__ import __version__
 from databricks.labs.lakebridge.config import TranspileConfig, LSPConfigOptionV1
@@ -445,5 +446,20 @@ def analyze(w: WorkspaceClient, source_directory: str, report_file: str):
     Analyzer.analyze(Path(input_folder), Path(output_file), source_tech)
 
 
+@lakebridge.command()
+def database_profiler(w: WorkspaceClient):
+    """Run the Profiler"""
+    with_user_agent_extra("cmd", "profiler")
+    ctx = ApplicationContext(w)
+    prompts = ctx.prompts
+    source_tech = prompts.choice("Select the source technology", Profiler.supported_source_technologies())
+    with_user_agent_extra("profiler_source_tech", make_alphanum_or_semver(source_tech))
+    user = ctx.current_user
+    logger.debug(f"User: {user}")
+    profiler = Profiler()
+    # TODO: Add extractor logic to ApplicationContext instead of creating inside the Profiler class
+    profiler.profile(source_tech, None)
+
+
 if __name__ == "__main__":
     lakebridge()

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

@@ -43,7 +43,6 @@ def _create_connector(db_type: str, config: dict[str, Any]) -> DatabaseConnector
         "snowflake": SnowflakeConnector,
         "mssql": MSSQLConnector,
         "tsql": MSSQLConnector,
-        "synapse": MSSQLConnector,
     }
 
     connector_class = connectors.get(db_type.lower())

src/databricks/labs/remorph/resources/assessments/synapse/common/connector.py

@@ -0,0 +1,64 @@
+from pathlib import Path
+import yaml
+from databricks.labs.remorph.connections.credential_manager import (
+    CredentialManager,
+    LocalSecretProvider,
+    EnvSecretProvider,
+    DatabricksSecretProvider,
+)
+
+from databricks.labs.remorph.connections.env_getter import EnvGetter
+from databricks.labs.remorph.resources.assessments.synapse.common.functions import get_synapse_jdbc_settings
+from sqlalchemy import create_engine
+from sqlalchemy.engine import Engine, Result, URL
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy import text
+
+
+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(file_path: Path):
+    env_getter = EnvGetter()
+
+    secret_providers = {
+        'local': LocalSecretProvider(),
+        'env': EnvSecretProvider(env_getter),
+        'databricks': DatabricksSecretProvider(),
+    }
+
+    loader = _load_credentials(file_path)
+    return CredentialManager(loader, secret_providers)
+
+
+def get_sqlpool_reader(config: dict, db_name: str, endpoint_key='dedicated_sql_endpoint'):
+    """
+    :param config:
+    :param db_name:
+    :return: returns a sqlachemy reader for the given dedicated SQL Pool database
+    """
+
+    query_params = {
+        "driver": config['driver'],
+        "loginTimeout": "30",
+    }
+
+    connection_string = URL.create(
+        "mssql+pyodbc",
+        username=config['sql_user'],
+        password=config['sql_password'],
+        host=config[endpoint_key],
+        port=config.get('port', 1433),
+        database=db_name,
+        query=query_params,
+    )
+    engine = create_engine(connection_string)
+    session = sessionmaker(bind=engine)
+    connection = session()
+
+    return connection