Implement databricks creds manager

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

@@ -9,10 +9,8 @@
 from databricks.labs.lakebridge.connections.credential_manager import (
     cred_file as creds,
     CredentialManager,
-    create_credential_manager,
 )
 from databricks.labs.lakebridge.connections.database_manager import DatabaseManager
-from databricks.labs.lakebridge.connections.env_getter import EnvGetter
 from databricks.labs.lakebridge.assessments import CONNECTOR_REQUIRED
 
 logger = logging.getLogger(__name__)
@@ -44,8 +42,8 @@ def __init__(
     def _configure_credentials(self) -> str:
         pass
 
-    @staticmethod
-    def _test_connection(source: str, cred_manager: CredentialManager):
+    def _test_connection(self, source: str):
+        cred_manager = CredentialManager.from_file(self._credential_file)
         config = cred_manager.get_credentials(source)
 
         try:
@@ -67,9 +65,7 @@ def run(self):
         logger.info(f"{source.capitalize()} details and credentials received.")
         if CONNECTOR_REQUIRED.get(self._source_name, True):
             if self.prompts.confirm(f"Do you want to test the connection to {source}?"):
-                cred_manager = create_credential_manager("lakebridge", EnvGetter())
-                if cred_manager:
-                    self._test_connection(source, cred_manager)
+                self._test_connection(source)
         logger.info(f"{source.capitalize()} Assessment Configuration Completed")
 
 

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

@@ -5,9 +5,8 @@
 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,
+    CredentialManager,
 )
-from databricks.labs.lakebridge.connections.env_getter import EnvGetter
 from databricks.labs.lakebridge.assessments import (
     PRODUCT_NAME,
     PRODUCT_PATH_PREFIX,
@@ -62,7 +61,7 @@ def profile(
     def _setup_extractor(platform: str) -> DatabaseManager | None:
         if not CONNECTOR_REQUIRED[platform]:
             return None
-        cred_manager = create_credential_manager(PRODUCT_NAME, EnvGetter())
+        cred_manager = CredentialManager.from_product_name(PRODUCT_NAME)
         connect_config = cred_manager.get_credentials(platform)
         return DatabaseManager(platform, connect_config)
 

src/databricks/labs/lakebridge/config.py

@@ -1,12 +1,14 @@
 import logging
 from collections.abc import Mapping, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, asdict
 from enum import Enum, auto
 from pathlib import Path
 from typing import Any, Literal, TypeVar, cast
 
 from databricks.labs.blueprint.installation import JsonValue
 from databricks.labs.blueprint.tui import Prompts
+
+from databricks.labs.lakebridge.reconcile.connectors.credentials import build_recon_creds, ReconcileCredentialsConfig
 from databricks.labs.lakebridge.transpiler.transpile_status import TranspileError
 from databricks.labs.lakebridge.reconcile.recon_config import Table
 
@@ -254,13 +256,24 @@ class ReconcileMetadataConfig:
 @dataclass
 class ReconcileConfig:
     __file__ = "reconcile.yml"
-    __version__ = 1
+    __version__ = 2
 
     data_source: str
     report_type: str
-    secret_scope: str
     database_config: DatabaseConfig
     metadata_config: ReconcileMetadataConfig
+    creds: ReconcileCredentialsConfig | None = None
+    # databricks does not require creds
+
+    @classmethod
+    def v1_migrate(cls, raw: dict[str, JsonValue]) -> dict[str, JsonValue]:
+        secret_scope = str(raw.pop("secret_scope"))
+        data_source = str(raw["data_source"])
+        maybe_creds = build_recon_creds(data_source, secret_scope)
+        if maybe_creds:
+            raw["creds"] = asdict(maybe_creds)
+        raw["version"] = 2
+        return raw
 
 
 @dataclass

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

@@ -1,9 +1,13 @@
 from pathlib import Path
 import logging
 from typing import Protocol
+import base64
 
 import yaml
 
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import NotFound
+
 from databricks.labs.lakebridge.connections.env_getter import EnvGetter
 
 
@@ -32,18 +36,69 @@ def get_secret(self, key: str) -> str:
             return key
 
 
-class DatabricksSecretProvider:
+class DatabricksSecretProvider(SecretProvider):
+    def __init__(self, ws: WorkspaceClient):
+        self._ws = ws
+
     def get_secret(self, key: str) -> str:
-        raise NotImplementedError("Databricks secret vault not implemented")
+        """Get the secret value given a secret scope & secret key.
+
+        Args:
+            key: key in the format 'scope/secret'
+        Returns:
+            The secret value.
+
+        Raises:
+          ValueError: The secret key must be in the format 'scope/secret'.
+          KeyError: The secret could not be found.
+        """
+        match key.split(sep="/", maxsplit=3):
+            case _scope, _key_only:
+                scope = _scope
+                key_only = _key_only
+            case _:
+                msg = f"Secret key must be in the format 'scope/secret': Got {key}"
+                raise ValueError(msg)
+
+        try:
+            secret = self._ws.secrets.get_secret(scope, key_only)
+            assert secret.value is not None
+            return base64.b64decode(secret.value).decode("utf-8")
+        except NotFound as e:
+            # TODO do not raise KeyError and standardize across all secret providers. Caller should handle missing secrets.
+            raise KeyError(f'Secret does not exist with scope: {scope} and key: {key_only}') from e
+        except UnicodeDecodeError as e:
+            msg = f"Secret {key} has Base64 bytes that cannot be decoded to UTF-8 string"
+            raise ValueError(msg) from e
 
 
 class CredentialManager:
     def __init__(self, credentials: dict, secret_providers: dict[str, SecretProvider]):
         self._credentials = credentials
         self._default_vault = self._credentials.get('secret_vault_type', 'local').lower()
         self._provider = secret_providers.get(self._default_vault)
-        if not self._provider:
-            raise ValueError(f"Unsupported secret vault type: {self._default_vault}")
+
+    @classmethod
+    def from_product_name(cls, product_name: str, ws: WorkspaceClient | None = None) -> "CredentialManager":
+        path = cred_file(product_name)
+        credentials = _load_credentials(path)
+        return cls.from_credentials(credentials, ws)
+
+    @classmethod
+    def from_file(cls, path: Path, ws: WorkspaceClient | None = None) -> "CredentialManager":
+        credentials = _load_credentials(path)
+        return cls.from_credentials(credentials, ws)
+
+    @classmethod
+    def from_credentials(cls, credentials: dict, ws: WorkspaceClient | None = None) -> "CredentialManager":
+        secret_providers: dict[str, SecretProvider] = {
+            'local': LocalSecretProvider(),
+            'env': EnvSecretProvider(EnvGetter()),
+        }
+
+        if ws:
+            secret_providers['databricks'] = DatabricksSecretProvider(ws)
+        return cls(credentials, secret_providers)
 
     def get_credentials(self, source: str) -> dict:
         if source not in self._credentials:
@@ -60,6 +115,23 @@ def _get_secret_value(self, key: str) -> str:
         return self._provider.get_secret(key)
 
 
+def build_credentials(vault_type: str, source: str, credentials: dict) -> dict:
+    """Build credentials dictionary with secret vault type included.
+
+    Args:
+        vault_type: The type of secret vault (e.g., 'local', 'databricks').
+        source: The source system name.
+        credentials: The original credentials dictionary.
+
+    Returns:
+        A new credentials dictionary including the secret vault type.
+    """
+    return {
+        source: credentials,
+        'secret_vault_type': vault_type.lower(),
+    }
+
+
 def _get_home() -> Path:
     return Path(__file__).home()
 
@@ -74,16 +146,3 @@ def _load_credentials(path: Path) -> dict:
             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) -> CredentialManager:
-    creds_path = cred_file(product_name)
-    creds = _load_credentials(creds_path)
-
-    secret_providers = {
-        'local': LocalSecretProvider(),
-        'env': EnvSecretProvider(env_getter),
-        'databricks': DatabricksSecretProvider(),
-    }
-
-    return CredentialManager(creds, secret_providers)

src/databricks/labs/lakebridge/deployment/recon.py

@@ -60,10 +60,11 @@ def uninstall(self, recon_config: ReconcileConfig | None):
             f"Won't remove reconcile metadata schema `{recon_config.metadata_config.schema}` "
             f"from catalog `{recon_config.metadata_config.catalog}`. Please remove it and the tables inside manually."
         )
-        logging.info(
-            f"Won't remove configured reconcile secret scope `{recon_config.secret_scope}`. "
-            f"Please remove it manually."
-        )
+        if recon_config.creds:
+            logging.info(
+                f"Won't remove configured reconcile credentials from `{recon_config.creds.vault_type}`. "
+                f"Please remove it manually."
+            )
 
     def _deploy_tables(self, recon_config: ReconcileConfig):
         logger.info("Deploying reconciliation metadata tables.")

src/databricks/labs/lakebridge/install.py

@@ -24,6 +24,7 @@
 from databricks.labs.lakebridge.contexts.application import ApplicationContext
 from databricks.labs.lakebridge.deployment.configurator import ResourceConfigurator
 from databricks.labs.lakebridge.deployment.installation import WorkspaceInstallation
+from databricks.labs.lakebridge.reconcile.connectors.credentials import build_recon_creds
 from databricks.labs.lakebridge.reconcile.constants import ReconReportType, ReconSourceType
 from databricks.labs.lakebridge.transpiler.installers import (
     BladebridgeInstaller,
@@ -325,18 +326,19 @@ def _prompt_for_new_reconcile_installation(self) -> ReconcileConfig:
         report_type = self._prompts.choice(
             "Select the report type:", [report_type.value for report_type in ReconReportType]
         )
-        scope_name = self._prompts.question(
+        scope_name = self._prompts.question(  # TODO deprecate
             f"Enter Secret scope name to store `{data_source.capitalize()}` connection details / secrets",
             default=f"remorph_{data_source}",
         )
+        creds = build_recon_creds(data_source, scope_name)
 
         db_config = self._prompt_for_reconcile_database_config(data_source)
         metadata_config = self._prompt_for_reconcile_metadata_config()
 
         return ReconcileConfig(
             data_source=data_source,
             report_type=report_type,
-            secret_scope=scope_name,
+            creds=creds,
             database_config=db_config,
             metadata_config=metadata_config,
         )

src/databricks/labs/lakebridge/reconcile/connectors/credentials.py

@@ -0,0 +1,90 @@
+import logging
+from dataclasses import dataclass
+
+from databricks.sdk import WorkspaceClient
+
+from databricks.labs.lakebridge.connections.credential_manager import build_credentials, CredentialManager
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ReconcileCredentialsConfig:
+    vault_type: str
+    vault_secret_names: dict[str, str]
+
+    def __post_init__(self):
+        if self.vault_type != "databricks":
+            raise ValueError(f"Unsupported vault_type: {self.vault_type}")
+
+
+_REQUIRED_JDBC_CREDS = [
+    "host",
+    "port",
+    "database",
+    "user",
+    "password",
+]
+
+_TSQL_REQUIRED_CREDS = [*_REQUIRED_JDBC_CREDS, "encrypt", "trustServerCertificate"]
+
+_ORACLE_REQUIRED_CREDS = [*_REQUIRED_JDBC_CREDS]
+
+_SNOWFLAKE_REQUIRED_CREDS = [
+    "sfUser",
+    "sfUrl",
+    "sfDatabase",
+    "sfSchema",
+    "sfWarehouse",
+    "sfRole",
+    # sfPassword is not required here; auth is validated separately
+]
+
+_SOURCE_CREDENTIALS_MAP = {
+    "databricks": [],
+    "snowflake": _SNOWFLAKE_REQUIRED_CREDS,
+    "oracle": _ORACLE_REQUIRED_CREDS,
+    "mssql": _TSQL_REQUIRED_CREDS,
+    "synapse": _TSQL_REQUIRED_CREDS,
+}
+
+
+def build_recon_creds(source: str, secret_scope: str) -> ReconcileCredentialsConfig | None:
+    if source == "databricks":
+        return None
+
+    keys = _SOURCE_CREDENTIALS_MAP.get(source)
+    if not keys:
+        raise ValueError(f"Unsupported source system: {source}")
+    parsed = {key: f"{secret_scope}/{key}" for key in keys}
+
+    if source == "snowflake":
+        logger.warning("Please specify the Snowflake authentication method in the credentials config.")
+        parsed["pem_private_key"] = f"{secret_scope}/pem_private_key"
+        parsed["sfPassword"] = f"{secret_scope}/sfPassword"
+
+    return ReconcileCredentialsConfig("databricks", parsed)
+
+
+def validate_creds(creds: ReconcileCredentialsConfig, source: str) -> None:
+    required_keys = _SOURCE_CREDENTIALS_MAP.get(source)
+    if not required_keys:
+        raise ValueError(f"Unsupported source system: {source}")
+
+    missing = [k for k in required_keys if not creds.vault_secret_names.get(k)]
+    if missing:
+        raise ValueError(
+            f"Missing mandatory {source} credentials. " f"Please configure all of {required_keys}. Missing: {missing}"
+        )
+
+
+def load_and_validate_credentials(
+    creds: ReconcileCredentialsConfig,
+    ws: WorkspaceClient,
+    source: str,
+) -> dict[str, str]:
+    validate_creds(creds, source)
+
+    parsed = build_credentials(creds.vault_type, source, creds.vault_secret_names)
+    resolved = CredentialManager.from_credentials(parsed, ws).get_credentials(source)
+    return resolved