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

@@ -2,6 +2,7 @@
 from typing import Optional
 
 import cfnresponse
+from databricks_cdk.resources.service_principals.service_principal_secrets import ServicePrincipalSecretsProperties, create_or_update_service_principal_secrets, delete_service_principal_secrets
 from pydantic import BaseModel, ValidationError
 
 from databricks_cdk.resources.account.credentials import (
@@ -91,6 +92,11 @@
     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.sql_warehouses.sql_warehouses import (
     SQLWarehouseProperties,
     create_or_update_warehouse,
@@ -221,6 +227,12 @@ 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 +337,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_workspace_client, get_account_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"""
+    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,123 @@
+import json
+import logging
+from typing import Optional
+
+import boto3
+from databricks.sdk.service.oauth2 import SecretInfo
+from databricks.sdk import AccountClient
+from databricks.sdk.errors import NotFound
+from pydantic import BaseModel
+
+from databricks_cdk.utils import CnfResponse, get_account_client
+from databricks_cdk.resources.service_principals.service_principal import get_service_principal
+
+logger = logging.getLogger(__name__)
+
+
+class ServicePrincipalSecretsCreationError(Exception):
+    pass
+
+
+class ServicePrincipalSecretsProperties(BaseModel):
+    service_principal_id: int
+
+
+def create_or_update_service_principal_secrets(
+        properties: ServicePrincipalSecretsProperties,
+        physical_resource_id: Optional[str] = None
+    ) -> CnfResponse:
+    """
+    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:
+        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
+        )
+        return CnfResponse(
+            physical_resource_id=existing_service_principal_secrets.id
+        )
+
+    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) -> CnfResponse:
+    """
+    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 = f"{service_principal.display_name}/{service_principal.id}"
+    add_to_secrets_manager(
+        secret_name=secret_name,
+        client_id=service_principal.application_id,
+        client_secret=created_service_principal_secrets.secret
+    )
+    return CnfResponse(
+        physical_resource_id=created_service_principal_secrets.id
+    )
+
+
+def delete_service_principal_secrets(properties: ServicePrincipalSecretsProperties, physical_resource_id: str) -> CnfResponse:
+    """Delete service pricncipal secrets on databricks."""
+    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 = f"{service_principal.display_name}/{service_principal.id}"
+    delete_from_secrets_manager(secret_name)
+    return CnfResponse(physical_resource_id=physical_resource_id)
+
+
+def add_to_secrets_manager(secret_name: str, client_id: str, client_secret: str) -> None:
+    """Adds credentials to secrets manager at /databricks/service_principal/secrets/{{secret_name}}"""
+    client = boto3.client("secretsmanager")
+    secret_full_name = f"/databricks/service_principal/secrets/{secret_name}"
+    secret_string = {"client_id": client_id, "client_secret": client_secret}
+    client.create_secret(Name=secret_full_name, SecretString=json.dumps(secret_string))["ARN"]
+
+
+def delete_from_secrets_manager(secret_name: str) -> None:
+    """Removes credentials from secrets manager at /databricks/service_principal/secrets/{{secret_name}}"""
+    client = boto3.client("secretsmanager")
+    secret_full_name = f"/databricks/service_principal/secrets/{secret_name}"
+    try:
+        client.delete_secret(SecretId=secret_full_name, ForceDeleteWithoutRecovery=True)
+    except client.exceptions.ResourceNotFoundException:
+        logger.warning("Secret with name %s not found", secret_full_name)

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 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,6 @@ 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)
 
     return account_client