Merge pull request #1074 from godatadriven/feature/service-principal-…

…resource

Add service principal resource

README.md

@@ -13,8 +13,8 @@ databricks-cdk is an open source library that allows you to deploy Databricks wo
 - AWS CDK v2
 - Databricks account
 - AWS Systems Manager (SSM) parameters:
-  - `/databricks/deploy/user`
-  - `/databricks/deploy/password` (secure parameter)
+  - `/databricks/deploy/client-id`
+  - `/databricks/deploy/client-secret` (secure parameter)
   - `/databricks/account-id` (not AWS account ID)
 
 ## Installation

aws-lambda/src/databricks_cdk/resources/handler.py

@@ -91,6 +91,16 @@
     create_or_update_secret_scope,
     delete_secret_scope,
 )
+from databricks_cdk.resources.service_principals.service_principal import (
+    ServicePrincipalProperties,
+    create_or_update_service_principal,
+    delete_service_principal,
+)
+from databricks_cdk.resources.service_principals.service_principal_secrets import (
+    ServicePrincipalSecretsProperties,
+    create_or_update_service_principal_secrets,
+    delete_service_principal_secrets,
+)
 from databricks_cdk.resources.sql_warehouses.sql_warehouses import (
     SQLWarehouseProperties,
     create_or_update_warehouse,
@@ -221,6 +231,13 @@ def create_or_update_resource(event: DatabricksEvent) -> CnfResponse:
         )
     elif action == "volume":
         return create_or_update_volume(VolumeProperties(**event.ResourceProperties), event.PhysicalResourceId)
+    elif action == "service-principal":
+        return create_or_update_service_principal(ServicePrincipalProperties(**event.ResourceProperties))
+    elif action == "service-principal-secrets":
+        return create_or_update_service_principal_secrets(
+            ServicePrincipalSecretsProperties(**event.ResourceProperties),
+            event.PhysicalResourceId,
+        )
     else:
         raise RuntimeError(f"Unknown action: {action}")
 
@@ -325,6 +342,14 @@ def delete_resource(event: DatabricksEvent) -> CnfResponse:
         )
     elif action == "volume":
         return delete_volume(VolumeProperties(**event.ResourceProperties), event.PhysicalResourceId)
+    elif action == "service-principal":
+        return delete_service_principal(
+            ServicePrincipalProperties(**event.ResourceProperties), event.PhysicalResourceId
+        )
+    elif action == "service-principal-secrets":
+        return delete_service_principal_secrets(
+            ServicePrincipalSecretsProperties(**event.ResourceProperties), event.PhysicalResourceId
+        )
     else:
         raise RuntimeError(f"Unknown action: {action}")
 

aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py

@@ -0,0 +1,114 @@
+import logging
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import NotFound
+from databricks.sdk.service.iam import ServicePrincipal
+from pydantic import BaseModel
+
+from databricks_cdk.utils import CnfResponse, get_account_client, get_workspace_client
+
+logger = logging.getLogger(__name__)
+
+
+class ServicePrincipalCreationError(Exception):
+    pass
+
+
+class ServicePrincipalNotFoundError(Exception):
+    pass
+
+
+class ServicePrincipalProperties(BaseModel):
+    workspace_url: str
+    service_principal: ServicePrincipal
+
+
+class ServicePrincipalResponse(CnfResponse):
+    name: str
+
+
+def create_or_update_service_principal(properties: ServicePrincipalProperties) -> ServicePrincipalResponse:
+    """
+    Create or update service principal on databricks. If service principal id is provided, it will update the existing service principal
+    else it will create a new one.
+    """
+    workspace_client = get_workspace_client(properties.workspace_url)
+
+    if properties.service_principal.id is None:
+        # service principal doesn't exist yet so create new one
+        return create_service_principal(properties.service_principal, workspace_client)
+
+    # check if service principal exists
+    get_service_principal(properties.service_principal.id, workspace_client)
+
+    # update existing service principal
+    return update_service_principal(properties.service_principal, workspace_client)
+
+
+def get_service_principal(physical_resource_id: str, workspace_client: WorkspaceClient) -> ServicePrincipal:
+    """Get service principal on databricks"""
+
+    try:
+        service_principal = workspace_client.service_principals.get(id=physical_resource_id)
+    except NotFound:
+        raise ServicePrincipalNotFoundError(f"Service principal with id {physical_resource_id} not found")
+
+    return service_principal
+
+
+def create_service_principal(
+    service_principal: ServicePrincipal, workspace_client: WorkspaceClient
+) -> ServicePrincipalResponse:
+    """Create service principal on databricks"""
+
+    created_service_principal = workspace_client.service_principals.create(
+        active=service_principal.active,
+        application_id=service_principal.application_id,
+        display_name=service_principal.display_name,
+        entitlements=service_principal.entitlements,
+        external_id=service_principal.external_id,
+        groups=service_principal.groups,
+        id=service_principal.id,
+        roles=service_principal.roles,
+        schemas=service_principal.schemas,
+    )
+
+    if created_service_principal.id is None:
+        raise ServicePrincipalCreationError("Service principal creation failed, there was no id found")
+
+    return ServicePrincipalResponse(
+        name=created_service_principal.display_name, physical_resource_id=created_service_principal.id
+    )
+
+
+def update_service_principal(
+    service_principal: ServicePrincipal,
+    workspace_client: WorkspaceClient,
+) -> ServicePrincipalResponse:
+    """Update service principal on databricks."""
+    workspace_client.service_principals.update(
+        id=service_principal.id,
+        active=service_principal.active,
+        application_id=service_principal.application_id,
+        display_name=service_principal.display_name,
+        entitlements=service_principal.entitlements,
+        external_id=service_principal.external_id,
+        groups=service_principal.groups,
+        roles=service_principal.roles,
+        schemas=service_principal.schemas,
+    )
+
+    return ServicePrincipalResponse(name=service_principal.display_name, physical_resource_id=service_principal.id)
+
+
+def delete_service_principal(properties: ServicePrincipalProperties, physical_resource_id: str) -> CnfResponse:
+    """Delete a service pricncipal on databricks. It will delete the service principal from workspace and account, because it's duplicated during creation."""
+    workspace_client = get_workspace_client(properties.workspace_url)
+    account_client = get_account_client()
+    try:
+        workspace_client.service_principals.delete(id=physical_resource_id)
+        account_client.service_principals.delete(id=physical_resource_id)
+    except NotFound:
+        logger.warning("Service principal not found, never existed or already removed")
+
+    return CnfResponse(physical_resource_id=physical_resource_id)

aws-lambda/src/databricks_cdk/resources/service_principals/service_principal_secrets.py

@@ -0,0 +1,161 @@
+import json
+import logging
+from typing import Optional
+
+import boto3
+from databricks.sdk import AccountClient
+from databricks.sdk.errors import NotFound
+from databricks.sdk.service.iam import ServicePrincipal
+from databricks.sdk.service.oauth2 import SecretInfo
+from pydantic import BaseModel
+
+from databricks_cdk.resources.service_principals.service_principal import get_service_principal
+from databricks_cdk.utils import CnfResponse, get_account_client
+
+SECRETS_MANAGER_RESOURCE_PREFIX = "/databricks/service_principal/secrets"
+
+logger = logging.getLogger(__name__)
+
+
+class ServicePrincipalSecretsCreationError(Exception):
+    pass
+
+
+class ServicePrincipalSecretsProperties(BaseModel):
+    service_principal_id: int
+
+
+class ServicePrincipalSecretsResponse(CnfResponse):
+    secrets_manager_arn: str
+    secrets_manager_name: str
+
+
+def create_or_update_service_principal_secrets(
+    properties: ServicePrincipalSecretsProperties, physical_resource_id: Optional[str] = None
+) -> ServicePrincipalSecretsResponse:
+    """
+    Create or update service principal secrets on databricks.
+    If service principal secrets already exist, it will return the existing service principal secrets.
+    If service principal secrets doesn't exist, it will create a new one.
+    """
+    account_client = get_account_client()
+
+    if physical_resource_id:
+        return get_existing_service_principal_secrets_response(properties, physical_resource_id, account_client)
+
+    return create_service_principal_secrets(properties, account_client)
+
+
+def get_service_principal_secrets(
+    service_principal_id: int, physical_resource_id: str, account_client: AccountClient
+) -> SecretInfo:
+    """Get service principal secrets on databricks based on physical resource id and service principal id."""
+    existing_service_principal_secrets = account_client.service_principal_secrets.list(
+        service_principal_id=service_principal_id
+    )
+
+    for secret_info in existing_service_principal_secrets:
+        if secret_info is not None and secret_info.id == physical_resource_id:
+            return secret_info
+    else:
+        raise NotFound(f"Service principal secrets with id {physical_resource_id} not found")
+
+
+def create_service_principal_secrets(
+    properties: ServicePrincipalSecretsProperties, account_client: AccountClient
+) -> ServicePrincipalSecretsResponse:
+    """
+    Create service principal secrets on databricks.
+    It will create a new service principal secrets and store it in secrets manager.
+    """
+    service_principal = get_service_principal(properties.service_principal_id, account_client)
+    created_service_principal_secrets = account_client.service_principal_secrets.create(
+        service_principal_id=properties.service_principal_id
+    )
+
+    if created_service_principal_secrets.id is None:
+        raise ServicePrincipalSecretsCreationError("Failed to create service principal secrets")
+
+    secret_name = _construct_secret_name(service_principal, created_service_principal_secrets.id)
+    secrets_manager_resource = add_to_secrets_manager(
+        secret_name=secret_name,
+        client_id=service_principal.application_id,
+        client_secret=created_service_principal_secrets.secret,
+    )
+    return ServicePrincipalSecretsResponse(
+        physical_resource_id=created_service_principal_secrets.id,
+        secrets_manager_arn=secrets_manager_resource["ARN"],
+        secrets_manager_name=secrets_manager_resource["Name"],
+    )
+
+
+def get_existing_service_principal_secrets_response(
+    properties: ServicePrincipalSecretsProperties,
+    physical_resource_id: str,
+    account_client: AccountClient,
+) -> ServicePrincipalSecretsResponse:
+    """
+    Get existing service principal secrets response.
+    It will fetch the existing service principal secrets and secrets manager resource.
+    """
+    service_principal = get_service_principal(properties.service_principal_id, account_client)
+    existing_service_principal_secrets = get_service_principal_secrets(
+        service_principal_id=properties.service_principal_id,
+        physical_resource_id=physical_resource_id,
+        account_client=account_client,
+    )
+    secret_name = _construct_secret_name(service_principal, existing_service_principal_secrets.id)
+    secrets_manager_resource = get_from_secrets_manager(secret_name)
+    return ServicePrincipalSecretsResponse(
+        physical_resource_id=existing_service_principal_secrets.id,
+        secrets_manager_arn=secrets_manager_resource["ARN"],
+        secrets_manager_name=secrets_manager_resource["Name"],
+    )
+
+
+def delete_service_principal_secrets(
+    properties: ServicePrincipalSecretsProperties, physical_resource_id: str
+) -> CnfResponse:
+    """Delete service pricncipal secrets on databricks. It will delete the service principal secrets from databricks and secrets manager."""
+    account_client = get_account_client()
+
+    try:
+        account_client.service_principal_secrets.delete(
+            service_principal_id=properties.service_principal_id, secret_id=physical_resource_id
+        )
+    except NotFound:
+        logger.warning("Service principal secrets with id %s not found", physical_resource_id)
+
+    service_principal = get_service_principal(properties.service_principal_id, account_client)
+    secret_name = _construct_secret_name(service_principal, physical_resource_id)
+    delete_from_secrets_manager(secret_name)
+    return CnfResponse(physical_resource_id=physical_resource_id)
+
+
+def get_from_secrets_manager(secret_name: str) -> dict:
+    """
+    Get information from secrets manager at /{{SECRETS_MANAGER_RESOURCE_PREFIX}}/{{secret_name}}.
+    This will only retrieve general information about the secret, not the actual secret value.
+    """
+    client = boto3.client("secretsmanager")
+    return client.describe_secret(SecretId=secret_name)
+
+
+def add_to_secrets_manager(secret_name: str, client_id: str, client_secret: str) -> dict:
+    """Adds credentials to secrets manager at /{{SECRETS_MANAGER_RESOURCE_PREFIX}}/{{secret_name}}"""
+    client = boto3.client("secretsmanager")
+    secret_string = {"client_id": client_id, "client_secret": client_secret}
+    return client.create_secret(Name=secret_name, SecretString=json.dumps(secret_string))
+
+
+def delete_from_secrets_manager(secret_name: str) -> None:
+    """Removes credentials from secrets manager at /{{SECRETS_MANAGER_RESOURCE_PREFIX}}/{{secret_name}}"""
+    client = boto3.client("secretsmanager")
+    try:
+        client.delete_secret(SecretId=secret_name, ForceDeleteWithoutRecovery=True)
+    except client.exceptions.ResourceNotFoundException:
+        logger.warning("Secrets are not found in secrets manager")
+
+
+def _construct_secret_name(service_principal: ServicePrincipal, secrets_id: str) -> str:
+    return f"{SECRETS_MANAGER_RESOURCE_PREFIX}/{service_principal.display_name}/{service_principal.id}/{secrets_id}"

aws-lambda/src/databricks_cdk/utils.py

@@ -66,7 +66,7 @@ def get_authorization_headers() -> Dict[str, str]:
     return get_authentication_config().authenticate()
 
 
-class CnfResponse(BaseModel):
+class CnfResponse(BaseModel):  # TODO: rename to CfnResponse
     physical_resource_id: str
 
 

aws-lambda/tests/conftest.py

@@ -3,6 +3,8 @@
 
 import pytest
 from databricks.sdk import AccountClient, CredentialsAPI, ExperimentsAPI, ModelRegistryAPI, VolumesAPI, WorkspaceClient
+from databricks.sdk.service.iam import AccountServicePrincipalsAPI, ServicePrincipalsAPI
+from databricks.sdk.service.oauth2 import ServicePrincipalSecretsAPI
 
 
 @pytest.fixture(scope="function", autouse=True)
@@ -24,6 +26,7 @@ def workspace_client():
     workspace_client.model_registry = MagicMock(spec=ModelRegistryAPI)
     workspace_client.experiments = MagicMock(spec=ExperimentsAPI)
     workspace_client.volumes = MagicMock(spec=VolumesAPI)
+    workspace_client.service_principals = MagicMock(spec=ServicePrincipalsAPI)
 
     return workspace_client
 
@@ -34,5 +37,7 @@ def account_client():
 
     # mock all of the underlying service api's
     account_client.credentials = MagicMock(spec=CredentialsAPI)
+    account_client.service_principal_secrets = MagicMock(spec=ServicePrincipalSecretsAPI)
+    account_client.service_principals = MagicMock(spec=AccountServicePrincipalsAPI)
 
     return account_client