Added Configure Secrets support to databricks labs remorph configure-secrets cli command

labs.yml

@@ -50,3 +50,5 @@ commands:
       - name: report
         default: all
         description: Report type
+  - name: configure-secrets
+    description: Utility to setup Scope and Secrets on Databricks Workspace

src/databricks/labs/remorph/cli.py

@@ -7,6 +7,7 @@
 from databricks.sdk import WorkspaceClient
 
 from databricks.labs.remorph.config import MorphConfig
+from databricks.labs.remorph.helpers.recon_config_utils import ReconConfigPrompts
 from databricks.labs.remorph.reconcile.execute import recon
 from databricks.labs.remorph.transpiler.execute import morph
 
@@ -87,5 +88,17 @@ def reconcile(w: WorkspaceClient, recon_conf: str, conn_profile: str, source: st
     recon(recon_conf, conn_profile, source, report)
 
 
+@remorph.command
+def configure_secrets(w: WorkspaceClient):
+    """Setup reconciliation connection profile details as Secrets on Databricks Workspace"""
+    recon_conf = ReconConfigPrompts(w)
+
+    # Prompt for source
+    source = recon_conf.prompt_source()
+
+    logger.info(f"Setting up Scope, Secrets for `{source}` reconciliation")
+    recon_conf.prompt_and_save_connection_details()
+
+
 if __name__ == "__main__":
     remorph()

src/databricks/labs/remorph/helpers/db_workspace_utils.py

@@ -0,0 +1,96 @@
+import logging
+
+from databricks.labs.blueprint.tui import Prompts
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors.platform import ResourceDoesNotExist
+
+logger = logging.getLogger(__name__)
+
+
+class DatabricksSecretsClient:
+    def __init__(self, ws: WorkspaceClient, prompts: Prompts):
+        self._ws = ws
+        self._prompts = prompts
+
+    def _scope_exists(self, scope_name: str) -> bool:
+        scope_exists = scope_name in [scope.name for scope in self._ws.secrets.list_scopes()]
+
+        if not scope_exists:
+            logger.error(
+                f"Error: Cannot find Secret Scope: `{scope_name}` in Databricks Workspace."
+                f"\nUse `remorph configure-secrets` to setup Scope and Secrets"
+            )
+            return False
+        logger.debug(f"Found Scope: `{scope_name}` in Databricks Workspace")
+        return True
+
+    def get_or_create_scope(self, scope_name: str):
+        """
+        Get or Create a new Scope in Databricks Workspace
+        :param scope_name:
+        """
+        scope_exists = self._scope_exists(scope_name)
+        if not scope_exists:
+            allow_scope_creation = self._prompts.confirm("Do you want to create a new one?")
+            if not allow_scope_creation:
+                msg = "Scope is needed to store Secrets in Databricks Workspace"
+                raise SystemExit(msg)
+
+            try:
+                logger.debug(f" Creating a new Scope: `{scope_name}`")
+                self._ws.secrets.create_scope(scope_name)
+            except Exception as ex:
+                logger.error(f"Exception while creating Scope: {ex}")
+                raise ex
+
+            logger.info(f" Created a new Scope: `{scope_name}`")
+        logger.info(f" Using Scope: `{scope_name}`...")
+
+    def secret_key_exists(self, scope_name: str, secret_key: str) -> bool:
+        try:
+            self._ws.secrets.get_secret(scope_name, secret_key)
+            logger.info(f"Found Secret key `{secret_key}` in Scope `{scope_name}`")
+            return True
+        except ResourceDoesNotExist:
+            logger.debug(f"Secret key `{secret_key}` not found in Scope `{scope_name}`")
+            return False
+
+    def delete_secret(self, scope_name: str, secret_key: str):
+        try:
+            logger.debug(f"Deleting Secret: *{secret_key}* in Scope: `{scope_name}`")
+            self._ws.secrets.delete_secret(scope=scope_name, key=secret_key)
+        except Exception as ex:
+            logger.error(f"Exception while deleting Secret `{secret_key}`: {ex}")
+            raise ex
+
+    def store_secret(self, scope_name: str, secret_key: str, secret_value: str):
+        try:
+            logger.debug(f"Storing Secret: *{secret_key}* in Scope: `{scope_name}`")
+            self._ws.secrets.put_secret(scope=scope_name, key=secret_key, string_value=secret_value)
+        except Exception as ex:
+            logger.error(f"Exception while storing Secret `{secret_key}`: {ex}")
+            raise ex
+
+    def store_connection_secrets(self, scope_name: str, conn_details: tuple[str, dict[str, str]]):
+        engine = conn_details[0]
+        secrets = conn_details[1]
+
+        logger.debug(f"Storing `{engine}` Connection Secrets in Scope: `{scope_name}`")
+
+        for key, value in secrets.items():
+            secret_key = engine + '_' + key
+            logger.debug(f"Processing Secret: *{secret_key}*")
+            if self.secret_key_exists(scope_name, secret_key):
+                overwrite_secret = self._prompts.confirm(f"Do you want to overwrite `{secret_key}`?")
+                if overwrite_secret:
+                    self.delete_secret(scope_name, secret_key)
+                    logger.debug(f"Deleted Secret: *{secret_key}* in Scope: `{scope_name}`")
+                    self.store_secret(scope_name, secret_key, value)
+                    logger.info(f"Overwritten Secret: *{secret_key}* in Scope: `{scope_name}`")
+            else:
+                self.store_secret(scope_name, secret_key, value)
+                logger.info(f"Stored Secret: *{secret_key}* in Scope: `{scope_name}`")
+
+    @property
+    def ws(self):
+        return self._ws

src/databricks/labs/remorph/helpers/recon_config_utils.py

@@ -0,0 +1,109 @@
+import logging
+
+from databricks.labs.blueprint.tui import Prompts
+from databricks.sdk import WorkspaceClient
+
+from databricks.labs.remorph.helpers.db_workspace_utils import DatabricksSecretsClient
+from databricks.labs.remorph.reconcile.constants import SourceType
+
+logger = logging.getLogger(__name__)
+
+recon_source_choices = [
+    SourceType.SNOWFLAKE.value,
+    SourceType.ORACLE.value,
+    SourceType.DATABRICKS.value,
+    SourceType.NETEZZA.value,
+]
+
+
+class ReconConfigPrompts:
+    def __init__(self, ws: WorkspaceClient, prompts: Prompts = Prompts()):
+        self._source = None
+        self._prompts = prompts
+        self._db_secrets = DatabricksSecretsClient(ws, prompts)
+
+    def prompt_source(self):
+        source = self._prompts.choice("Select the source", recon_source_choices)
+        self._source = source
+        return source
+
+    def _prompt_snowflake_connection_details(self) -> tuple[str, dict[str, str]]:
+        """
+        Prompt for Snowflake connection details
+        :return: tuple[str, dict[str, str]]
+        """
+        logger.info(
+            f"Please answer a couple of questions to configure `{SourceType.SNOWFLAKE.value}` Connection profile"
+        )
+
+        sf_url = self._prompts.question("Enter Snowflake URL")
+        account = self._prompts.question("Enter Account Name")
+        sf_user = self._prompts.question("Enter User")
+        sf_password = self._prompts.question("Enter Password")
+        sf_db = self._prompts.question("Enter Database")
+        sf_schema = self._prompts.question("Enter Schema")
+        sf_warehouse = self._prompts.question("Enter Snowflake Warehouse")
+        sf_role = self._prompts.question("Enter Role", default=" ")
+
+        sf_conn_details = {
+            "sfUrl": sf_url,
+            "account": account,
+            "sfUser": sf_user,
+            "sfPassword": sf_password,
+            "sfDatabase": sf_db,
+            "sfSchema": sf_schema,
+            "sfWarehouse": sf_warehouse,
+            "sfRole": sf_role,
+        }
+
+        sf_conn_dict = (SourceType.SNOWFLAKE.value, sf_conn_details)
+        return sf_conn_dict
+
+    def _prompt_oracle_connection_details(self) -> tuple[str, dict[str, str]]:
+        """
+        Prompt for Oracle connection details
+        :return: tuple[str, dict[str, str]]
+        """
+        logger.info(f"Please answer a couple of questions to configure `{SourceType.ORACLE.value}` Connection profile")
+        user = self._prompts.question("Enter User")
+        password = self._prompts.question("Enter Password")
+        host = self._prompts.question("Enter host")
+        port = self._prompts.question("Enter port")
+        database = self._prompts.question("Enter database/SID")
+
+        oracle_conn_details = {"user": user, "password": password, "host": host, "port": port, "database": database}
+
+        oracle_conn_dict = (SourceType.ORACLE.value, oracle_conn_details)
+        return oracle_conn_dict
+
+    def _connection_details(self):
+        """
+        Prompt for connection details based on the source
+        :return: None
+        """
+        logger.debug(f"Prompting for `{self._source}` connection details")
+        match self._source:
+            case SourceType.SNOWFLAKE.value:
+                return self._prompt_snowflake_connection_details()
+            case SourceType.ORACLE.value:
+                return self._prompt_oracle_connection_details()
+            case _:
+                raise ValueError(f"Source {self._source} is not yet configured...")
+
+    def prompt_and_save_connection_details(self):
+        """
+        Prompt for connection details and save them as Secrets in Databricks Workspace
+        """
+        # prompt for connection_details only if source is other than Databricks
+        if self._source == SourceType.DATABRICKS.value:
+            logger.info("*Databricks* as a source is supported only for **Hive MetaStore (HMS) setup**")
+            return
+
+        # Prompt for secret scope
+        scope_name = self._prompts.question("Enter Secret Scope name")
+        self._db_secrets.get_or_create_scope(scope_name)
+
+        # Prompt for connection details
+        connection_details = self._connection_details()
+        logger.debug(f"Storing `{self._source}` connection details as Secrets in Databricks Workspace...")
+        self._db_secrets.store_connection_secrets(scope_name, connection_details)

tests/unit/helpers/test_db_workspace_utils.py

@@ -0,0 +1,80 @@
+# pylint: disable=protected-access
+from unittest.mock import patch
+
+import pytest
+from databricks.labs.blueprint.tui import MockPrompts
+from databricks.sdk.errors.platform import ResourceDoesNotExist
+from databricks.sdk.service.workspace import SecretScope
+
+from databricks.labs.remorph.helpers.db_workspace_utils import DatabricksSecretsClient
+
+
+@pytest.fixture
+def db_secrets_client(mock_workspace_client):
+    return DatabricksSecretsClient(
+        mock_workspace_client,
+        prompts=MockPrompts(
+            {
+                r"Do you want to create a new one?": "yes",
+                r".*": "",
+            }
+        ),
+    )
+
+
+def test_scope_exists(db_secrets_client):
+    with patch.object(db_secrets_client._ws.secrets, 'list_scopes', return_value=[SecretScope(name="scope_name")]):
+        assert db_secrets_client._scope_exists("scope_name")
+
+
+def test_get_or_create_scope(db_secrets_client):
+    with patch.object(db_secrets_client._ws.secrets, 'create_scope', return_value={}):
+        db_secrets_client.get_or_create_scope("scope_name")
+
+
+def test_get_or_create_scope_exception(db_secrets_client):
+    with patch.object(db_secrets_client._ws.secrets, 'create_scope', side_effect=Exception()):
+        with pytest.raises(Exception):
+            db_secrets_client.get_or_create_scope("scope_name")
+
+
+def test_get_or_create_scope_no_exit(mock_workspace_client):
+    db_secrets_client = DatabricksSecretsClient(
+        mock_workspace_client,
+        prompts=MockPrompts(
+            {
+                r"Do you want to create a new one?": "no",
+                r".*": "",
+            }
+        ),
+    )
+    with patch.object(db_secrets_client._ws.secrets, 'list_scopes', return_value=[]):
+        with pytest.raises(SystemExit, match="Scope is needed to store Secrets in Databricks Workspace"):
+            db_secrets_client.get_or_create_scope("scope_name")
+
+
+def test_secret_key_exists(db_secrets_client):
+    with patch.object(db_secrets_client._ws.secrets, 'get_secret', side_effect=ResourceDoesNotExist()):
+        assert not db_secrets_client.secret_key_exists("scope_name", "secret_key")
+
+
+def test_delete_secret(db_secrets_client):
+    with patch.object(db_secrets_client._ws.secrets, 'delete_secret', side_effect=Exception()):
+        with pytest.raises(Exception):
+            db_secrets_client.delete_secret("scope_name", "secret_key")
+
+
+def test_store_secret(db_secrets_client):
+    with patch.object(db_secrets_client._ws.secrets, 'put_secret', side_effect=Exception()):
+        with pytest.raises(Exception):
+            db_secrets_client.store_secret("scope_name", "secret_key", "secret_value")
+
+
+def test_store_connection_secrets(db_secrets_client):
+    with patch.object(db_secrets_client, 'store_secret', return_value=None):
+        db_secrets_client.store_connection_secrets("scope_name", ("source", {"key": "value"}))
+
+
+def test_get_ws(db_secrets_client):
+    # Assert that the workspace client is returned
+    assert db_secrets_client.ws() is not None