Added Configure Secrets support to `databricks labs remorph configu…

…re-secrets` cli command (#254)

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

@@ -5,6 +5,7 @@
 from databricks.labs.blueprint.entrypoint import get_logger
 from databricks.labs.blueprint.installation import Installation
 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
 from databricks.sdk import WorkspaceClient
@@ -86,5 +87,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/recon_config_utils.py

@@ -0,0 +1,178 @@
+import logging
+
+from databricks.labs.blueprint.tui import Prompts
+from databricks.labs.remorph.reconcile.constants import SourceType
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors.platform import ResourceDoesNotExist
+
+logger = logging.getLogger(__name__)
+
+recon_source_choices = [
+    SourceType.SNOWFLAKE.value,
+    SourceType.ORACLE.value,
+    SourceType.DATABRICKS.value,
+]
+
+
+class ReconConfigPrompts:
+    def __init__(self, ws: WorkspaceClient, prompts: Prompts = Prompts()):
+        self._source = None
+        self._prompts = prompts
+        self._ws = ws
+
+    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 _ensure_scope_exists(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 `{scope_name}`: {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 _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}*")
+            debug_op = "Storing"
+            info_op = "Stored"
+            if self._secret_key_exists(scope_name, secret_key):
+                overwrite_secret = self._prompts.confirm(f"Do you want to overwrite `{secret_key}`?")
+                if not overwrite_secret:
+                    continue
+                debug_op = "Overwriting"
+                info_op = "Overwritten"
+
+            logger.debug(f"{debug_op} Secret: *{secret_key}* in Scope: `{scope_name}`")
+            self._store_secret(scope_name, secret_key, value)
+            logger.info(f"{info_op} Secret: *{secret_key}* in Scope: `{scope_name}`")
+
+    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()
+
+    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._ensure_scope_exists(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.store_connection_secrets(scope_name, connection_details)

tests/unit/helpers/test_recon_config_utils.py

@@ -0,0 +1,145 @@
+from unittest.mock import patch
+
+import pytest
+
+from databricks.labs.blueprint.tui import MockPrompts
+from databricks.labs.remorph.helpers.recon_config_utils import ReconConfigPrompts
+from databricks.sdk.errors.platform import ResourceDoesNotExist
+from databricks.sdk.service.workspace import SecretScope
+
+SOURCE_DICT = {"databricks": "0", "oracle": "1", "snowflake": "2"}
+SCOPE_NAME = "dummy_scope"
+
+
+def test_configure_secrets_snowflake_overwrite(mock_workspace_client):
+    prompts = MockPrompts(
+        {
+            r"Select the source": SOURCE_DICT["snowflake"],
+            r"Enter Secret Scope name": SCOPE_NAME,
+            r"Enter Snowflake URL": "dummy",
+            r"Enter Account Name": "dummy",
+            r"Enter User": "dummy",
+            r"Enter Password": "dummy",
+            r"Enter Database": "dummy",
+            r"Enter Schema": "dummy",
+            r"Enter Snowflake Warehouse": "dummy",
+            r"Enter Role": "dummy",
+            r"Do you want to overwrite.*": "yes",
+        }
+    )
+    mock_workspace_client.secrets.list_scopes.side_effect = [[SecretScope(name=SCOPE_NAME)]]
+    recon_conf = ReconConfigPrompts(mock_workspace_client, prompts)
+    recon_conf.prompt_source()
+
+    recon_conf.prompt_and_save_connection_details()
+
+
+def test_configure_secrets_oracle_insert(mock_workspace_client):
+    # mock prompts for Oracle
+    prompts = MockPrompts(
+        {
+            r"Select the source": SOURCE_DICT["oracle"],
+            r"Enter Secret Scope name": SCOPE_NAME,
+            r"Do you want to create a new one?": "yes",
+            r"Enter User": "dummy",
+            r"Enter Password": "dummy",
+            r"Enter host": "dummy",
+            r"Enter port": "dummy",
+            r"Enter database/SID": "dummy",
+        }
+    )
+
+    mock_workspace_client.secrets.list_scopes.side_effect = [[SecretScope(name="scope_name")]]
+
+    with patch(
+        "databricks.labs.remorph.helpers.recon_config_utils.ReconConfigPrompts._secret_key_exists",
+        return_value=False,
+    ):
+        recon_conf = ReconConfigPrompts(mock_workspace_client, prompts)
+        recon_conf.prompt_source()
+
+        recon_conf.prompt_and_save_connection_details()
+
+
+def test_configure_secrets_invalid_source(mock_workspace_client):
+    prompts = MockPrompts(
+        {
+            r"Select the source": "3",
+            r"Enter Secret Scope name": SCOPE_NAME,
+        }
+    )
+
+    with patch(
+        "databricks.labs.remorph.helpers.recon_config_utils.ReconConfigPrompts._scope_exists",
+        return_value=True,
+    ):
+        recon_conf = ReconConfigPrompts(mock_workspace_client, prompts)
+        with pytest.raises(ValueError, match="cannot get answer within 10 attempt"):
+            recon_conf.prompt_source()
+
+
+def test_store_connection_secrets_exception(mock_workspace_client):
+    prompts = MockPrompts(
+        {
+            r"Do you want to overwrite `source_key`?": "no",
+        }
+    )
+
+    mock_workspace_client.secrets.get_secret.side_effect = ResourceDoesNotExist("Not Found")
+    mock_workspace_client.secrets.put_secret.side_effect = Exception("Timed out")
+
+    recon_conf = ReconConfigPrompts(mock_workspace_client, prompts)
+
+    with pytest.raises(Exception, match="Timed out"):
+        recon_conf.store_connection_secrets("scope_name", ("source", {"key": "value"}))
+
+
+def test_configure_secrets_no_scope(mock_workspace_client):
+    prompts = MockPrompts(
+        {
+            r"Select the source": SOURCE_DICT["snowflake"],
+            r"Enter Secret Scope name": SCOPE_NAME,
+            r"Do you want to create a new one?": "no",
+        }
+    )
+
+    mock_workspace_client.secrets.list_scopes.side_effect = [[SecretScope(name="scope_name")]]
+
+    recon_conf = ReconConfigPrompts(mock_workspace_client, prompts)
+    recon_conf.prompt_source()
+
+    with pytest.raises(SystemExit, match="Scope is needed to store Secrets in Databricks Workspace"):
+        recon_conf.prompt_and_save_connection_details()
+
+
+def test_configure_secrets_create_scope_exception(mock_workspace_client):
+    prompts = MockPrompts(
+        {
+            r"Select the source": SOURCE_DICT["snowflake"],
+            r"Enter Secret Scope name": SCOPE_NAME,
+            r"Do you want to create a new one?": "yes",
+        }
+    )
+
+    mock_workspace_client.secrets.list_scopes.side_effect = [[SecretScope(name="scope_name")]]
+    mock_workspace_client.secrets.create_scope.side_effect = Exception("Network Error")
+
+    recon_conf = ReconConfigPrompts(mock_workspace_client, prompts)
+    recon_conf.prompt_source()
+
+    with pytest.raises(Exception, match="Network Error"):
+        recon_conf.prompt_and_save_connection_details()
+
+
+def test_store_connection_secrets_overwrite(mock_workspace_client):
+    prompts = MockPrompts(
+        {
+            r"Do you want to overwrite `source_key`?": "no",
+        }
+    )
+
+    with patch(
+        "databricks.labs.remorph.helpers.recon_config_utils.ReconConfigPrompts._secret_key_exists", return_value=True
+    ):
+        recon_conf = ReconConfigPrompts(mock_workspace_client, prompts)
+        recon_conf.store_connection_secrets("scope_name", ("source", {"key": "value"}))

tests/unit/test_cli.py

@@ -4,8 +4,10 @@
 import pytest
 import yaml
 
+from databricks.labs.blueprint.tui import MockPrompts
 from databricks.labs.remorph import cli
 from databricks.labs.remorph.config import MorphConfig
+from databricks.labs.remorph.helpers.recon_config_utils import ReconConfigPrompts
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
 
@@ -283,3 +285,23 @@ def test_recon_with_valid_input(mock_workspace_client_cli):
     with patch("os.path.exists", return_value=True), patch("databricks.labs.remorph.cli.recon") as mock_recon:
         cli.reconcile(mock_workspace_client_cli, recon_conf, conn_profile, source, report)
         mock_recon.assert_called_once_with(recon_conf, conn_profile, source, report)
+
+
+def test_configure_secrets_databricks(mock_workspace_client):
+    source_dict = {"databricks": "0", "netezza": "1", "oracle": "2", "snowflake": "3"}
+    prompts = MockPrompts(
+        {
+            r"Select the source": source_dict["databricks"],
+        }
+    )
+
+    recon_conf = ReconConfigPrompts(mock_workspace_client, prompts)
+    recon_conf.prompt_source()
+
+    recon_conf.prompt_and_save_connection_details()
+
+
+def test_cli_configure_secrets_config(mock_workspace_client):
+    with patch("databricks.labs.remorph.cli.ReconConfigPrompts") as mock_recon_config:
+        cli.configure_secrets(mock_workspace_client)
+        mock_recon_config.assert_called_once_with(mock_workspace_client)