Introduced Profiler Skeleton

pyproject.toml

@@ -122,7 +122,7 @@ asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope="function"
 
 [tool.mypy]
-exclude = ["tests/resources/.*"]
+exclude = ["tests/resources/.*", "src/databricks/labs/lakebridge/resources/assessments/.*"]
 
 [tool.black]
 target-version = ["py310"]
@@ -211,7 +211,7 @@ fail-under = 10.0
 ignore-patterns = ["^\\.#"]
 
 # Ignore files under tests/resources
-ignore-paths = ["tests/resources"]
+ignore-paths = ["tests/resources", "src/databricks/labs/lakebridge/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/__init__.py

@@ -2,14 +2,14 @@
     PRODUCT_NAME,
     PRODUCT_PATH_PREFIX,
     PROFILER_SOURCE_SYSTEM,
-    PLATFORM_TO_SOURCE_TECHNOLOGY,
+    PLATFORM_TO_SOURCE_TECHNOLOGY_CFG,
     CONNECTOR_REQUIRED,
 )
 
 __all__ = [
     "PRODUCT_NAME",
     "PRODUCT_PATH_PREFIX",
     "PROFILER_SOURCE_SYSTEM",
-    "PLATFORM_TO_SOURCE_TECHNOLOGY",
+    "PLATFORM_TO_SOURCE_TECHNOLOGY_CFG",
     "CONNECTOR_REQUIRED",
 ]

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

@@ -3,13 +3,14 @@
 PRODUCT_NAME = "lakebridge"
 PRODUCT_PATH_PREFIX = Path.home() / ".databricks" / "labs" / PRODUCT_NAME / "lib"
 
-PLATFORM_TO_SOURCE_TECHNOLOGY = {
+PLATFORM_TO_SOURCE_TECHNOLOGY_CFG = {
     "synapse": "src/databricks/labs/lakebridge/resources/assessments/synapse/pipeline_config.yml",
 }
 
 # TODO modify this PLATFORM_TO_SOURCE_TECHNOLOGY.keys() once all platforms are supported
 PROFILER_SOURCE_SYSTEM = ["mssql", "synapse"]
 
+
 # This flag indicates whether a connector is required for the source system when pipeline is trigger
 # For example in the case of synapse no connector is required and the python scripts
 # manage the connection by directly reading the credentials files

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

@@ -0,0 +1,88 @@
+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
+from databricks.labs.lakebridge.assessments import (
+    PRODUCT_NAME,
+    PRODUCT_PATH_PREFIX,
+    PLATFORM_TO_SOURCE_TECHNOLOGY_CFG,
+    CONNECTOR_REQUIRED,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Profiler:
+
+    def __init__(self, platform: str, pipeline_configs: PipelineConfig | None = None):
+        self._platform = platform
+        self._pipeline_config = pipeline_configs
+
+    @classmethod
+    def create(cls, platform: str) -> "Profiler":
+        pipeline_config_path = PLATFORM_TO_SOURCE_TECHNOLOGY_CFG.get(platform, None)
+        pipeline_config = None
+        if pipeline_config_path:
+            pipeline_config_absolute_path = Profiler._locate_config(pipeline_config_path)
+            pipeline_config = Profiler.path_modifier(config_file=pipeline_config_absolute_path)
+        return cls(platform, pipeline_config)
+
+    @classmethod
+    def supported_platforms(cls) -> list[str]:
+        return list(PLATFORM_TO_SOURCE_TECHNOLOGY_CFG.keys())
+
+    @staticmethod
+    def path_modifier(*, config_file: str | Path, path_prefix: Path = PRODUCT_PATH_PREFIX) -> 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"{path_prefix}/{step.extract_source}"
+        return config
+
+    def profile(
+        self,
+        *,
+        extractor: DatabaseManager | None = None,
+        pipeline_config: PipelineConfig | None = None,
+    ) -> None:
+        platform = self._platform.lower()
+        if not pipeline_config:
+            if not self._pipeline_config:
+                raise ValueError(f"Cannot Proceed without a valid pipeline configuration for {platform}")
+            pipeline_config = self._pipeline_config
+        self._execute(platform, pipeline_config, extractor)
+
+    @staticmethod
+    def _setup_extractor(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, pipeline_config: PipelineConfig, extractor=None) -> None:
+        try:
+            if extractor is None:
+                extractor = Profiler._setup_extractor(platform)
+
+            result = PipelineClass(pipeline_config, extractor).execute()
+            logger.info(f"Profile execution has completed successfully for {platform} for more info check: {result}.")
+        except FileNotFoundError as e:
+            logger.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:
+            logger.error(f"Error executing pipeline for source {platform}: {e}")
+            raise RuntimeError(f"Pipeline execution failed for source {platform} : {e}") from e
+
+    @staticmethod
+    def _locate_config(config_path: str | Path) -> 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/connections/database_manager.py

@@ -70,18 +70,29 @@ def _connect(self) -> Engine:
 
 class MSSQLConnector(_BaseConnector):
     def _connect(self) -> Engine:
-        query_params = {"driver": self.config['driver']}
+        auth_type = self.config.get('auth_type', 'sql_authentication')
+        db_name = self.config.get('database')
+
+        query_params = {
+            "driver": self.config['driver'],
+            "loginTimeout": "30",
+        }
+
+        if auth_type == "ad_passwd_authentication":
+            query_params = {
+                **query_params,
+                "authentication": "ActiveDirectoryPassword",
+            }
+        elif auth_type == "spn_authentication":
+            raise NotImplementedError("SPN Authentication not implemented yet")
 
-        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",
+            drivername="mssql+pyodbc",
             username=self.config['user'],
             password=self.config['password'],
             host=self.config['server'],
             port=self.config.get('port', 1433),
-            database=self.config['database'],
+            database=db_name,
             query=query_params,
         )
         return create_engine(connection_string)

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

@@ -0,0 +1,23 @@
+from databricks.labs.lakebridge.connections.database_manager import DatabaseManager
+
+
+def get_sqlpool_reader(
+    input_cred: dict,
+    db_name: str,
+    *,
+    endpoint_key: str = 'dedicated_sql_endpoint',
+    auth_type: str = 'sql_authentication',
+) -> DatabaseManager:
+    config = {
+        "driver": input_cred['driver'],
+        "server": input_cred[endpoint_key],
+        "database": db_name,
+        "user": input_cred['sql_user'],
+        "password": input_cred['sql_password'],
+        "port": input_cred.get('port', 1433),
+        "auth_type": auth_type,
+    }
+    # synapse and mssql use the same connector
+    source = "mssql"
+
+    return DatabaseManager(source, config)