Add SQL Server Profiler

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

@@ -5,10 +5,11 @@
 
 PLATFORM_TO_SOURCE_TECHNOLOGY_CFG = {
     "synapse": "src/databricks/labs/lakebridge/resources/assessments/synapse/pipeline_config.yml",
+    "mssql": "src/databricks/labs/lakebridge/resources/assessments/mssql/pipeline_config.yml",
 }
 
 # TODO modify this PLATFORM_TO_SOURCE_TECHNOLOGY.keys() once all platforms are supported
-PROFILER_SOURCE_SYSTEM = ["synapse"]
+PROFILER_SOURCE_SYSTEM = ["synapse", "mssql"]
 
 
 # This flag indicates whether a connector is required for the source system when pipeline is trigger

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

@@ -87,18 +87,21 @@ def _configure_credentials(self) -> str:
         secret_vault_type = str(self.prompts.choice("Enter secret vault type (local | env)", ["local", "env"])).lower()
         secret_vault_name = None
 
-        logger.info("Please refer to the documentation to understand the difference between local and env.")
-
         credential = {
             "secret_vault_type": secret_vault_type,
             "secret_vault_name": secret_vault_name,
             source: {
-                "database": self.prompts.question("Enter the database name"),
-                "driver": self.prompts.question("Enter the driver details"),
-                "server": self.prompts.question("Enter the server or host details"),
+                "auth_type": "sql_authentication",
+                "fetch_size": self.prompts.question("Enter fetch size", default="1000"),
+                "login_timeout": self.prompts.question("Enter login timeout (seconds)", default="30"),
+                "server": self.prompts.question("Enter the fully-qualified server name"),
                 "port": int(self.prompts.question("Enter the port details", valid_number=True)),
-                "user": self.prompts.question("Enter the user details"),
-                "password": self.prompts.password("Enter the password details"),
+                "user": self.prompts.question("Enter the SQL username"),
+                "password": self.prompts.password("Enter the 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"
+                ),
             },
         }
 

src/databricks/labs/lakebridge/resources/assessments/mssql/activity_extract.py

@@ -0,0 +1,72 @@
+import json
+import sys
+
+from databricks.labs.lakebridge.connections.credential_manager import create_credential_manager
+from databricks.labs.lakebridge.assessments import PRODUCT_NAME
+from databricks.labs.lakebridge.resources.assessments.mssql.common.connector import get_sqlserver_reader
+from databricks.labs.lakebridge.resources.assessments.mssql.common.queries import MSSQLQueries
+from databricks.labs.lakebridge.resources.assessments.synapse.common.duckdb_helpers import save_resultset_to_db
+from databricks.labs.lakebridge.resources.assessments.synapse.common.functions import arguments_loader, set_logger
+
+
+def execute():
+    logger = set_logger(__file__)
+
+    db_path, creds_file = arguments_loader(desc="MSSQL Server Activity Extract Script")
+    cred_manager = create_credential_manager(PRODUCT_NAME, creds_file)
+    mssql_settings = cred_manager.get_credentials("mssql")
+    auth_type = mssql_settings.get("auth_type", "sql_authentication")
+    server_name = mssql_settings.get("server", "")
+    try:
+
+        # TODO: get the last time the profiler was executed
+        # For now, we'll default to None, but this will eventually need
+        # input from a scheduler component.
+        last_execution_time = None
+        mode = "overwrite"
+
+        # Extract activity metrics
+        logger.info(f"Extracting activity metrics for: {server_name}")
+        print(f"Extracting activity metrics for: {server_name}")
+        connection = get_sqlserver_reader(
+            mssql_settings, db_name="master", server_name=server_name, auth_type=auth_type
+        )
+
+        # Query stats
+        table_name = "query_stats"
+        table_query = MSSQLQueries.get_query_stats(last_execution_time)
+        logger.info(f"Loading '{table_name}' for SQL server: {server_name}")
+        result = connection.fetch(table_query)
+        save_resultset_to_db(result, f"mssql_{table_name}", db_path, mode=mode)
+
+        # Stored procedure stats
+        table_name = "proc_stats"
+        table_query = MSSQLQueries.get_procedure_stats(last_execution_time)
+        logger.info(f"Loading '{table_name}' for SQL server: {server_name}")
+        result = connection.fetch(table_query)
+        save_resultset_to_db(result, f"mssql_{table_name}", db_path, mode=mode)
+
+        # Session info
+        table_name = "sessions"
+        table_query = MSSQLQueries.get_sessions(last_execution_time)
+        logger.info(f"Loading '{table_name}' for SQL server: {server_name}")
+        result = connection.fetch(table_query)
+        save_resultset_to_db(result, f"mssql_{table_name}", db_path, mode=mode)
+
+        # CPU Utilization
+        table_name = "cpu_utilization"
+        table_query = MSSQLQueries.get_cpu_utilization(last_execution_time)
+        logger.info(f"Loading '{table_name}' for SQL server: {server_name}")
+        result = connection.fetch(table_query)
+        save_resultset_to_db(result, f"mssql_{table_name}", db_path, mode=mode)
+
+        print(json.dumps({"status": "success", "message": "All data loaded successfully loaded successfully"}))
+
+    except Exception as e:
+        logger.error(f"Failed to extract activity info for SQL server: {str(e)}")
+        print(json.dumps({"status": "error", "message": str(e)}), file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    execute()

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

@@ -0,0 +1,22 @@
+from databricks.labs.lakebridge.connections.database_manager import DatabaseManager
+
+
+def get_sqlserver_reader(
+    input_cred: dict,
+    db_name: str,
+    *,
+    server_name: str,
+    auth_type: str = 'sql_authentication',
+) -> DatabaseManager:
+    config = {
+        "driver": input_cred['driver'],
+        "server": server_name,
+        "database": db_name,
+        "user": input_cred['user'],
+        "password": input_cred['password'],
+        "port": input_cred.get('port', 1433),
+        "auth_type": auth_type,
+    }
+    source = "mssql"
+
+    return DatabaseManager(source, config)

src/databricks/labs/lakebridge/resources/assessments/mssql/common/functions.py

@@ -0,0 +1,9 @@
+from azure.identity import DefaultAzureCredential
+from azure.mgmt.sql import SqlManagementClient
+
+
+def create_msql_sql_client(config: dict) -> SqlManagementClient:
+    """
+    Creates an Azure SQL management client for the provided subscription using the default Azure credential.
+    """
+    return SqlManagementClient(credential=DefaultAzureCredential(), subscription_id=config["subscription_id"])