From fda3baa46de318e14f8e99ac2cdd381e58cc1725 Mon Sep 17 00:00:00 2001 From: Anton Shchederkin Date: Fri, 13 Dec 2024 14:30:21 +0100 Subject: [PATCH] Add service principal resource --- README.md | 4 +- .../resources/service_principals/__init__.py | 0 .../service_principals/service_principal.py | 108 ++++++++++++++++++ 3 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 aws-lambda/src/databricks_cdk/resources/service_principals/__init__.py create mode 100644 aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py diff --git a/README.md b/README.md index 60e29e32..153cdaa1 100644 --- a/README.md +++ b/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 diff --git a/aws-lambda/src/databricks_cdk/resources/service_principals/__init__.py b/aws-lambda/src/databricks_cdk/resources/service_principals/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py b/aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py new file mode 100644 index 00000000..bfe55399 --- /dev/null +++ b/aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py @@ -0,0 +1,108 @@ +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 + +logger = logging.getLogger(__name__) + + +class ServicePrincipalCreationError(Exception): + pass + + +class ServicePrincipalNotFoundError(Exception): + pass + + +class ServicePrincipalProperties(BaseModel): + workspace_url: str + service_principal: ServicePrincipal + + +class ServicePrincipalResponce(CnfResponse): + name: str + + +def create_or_update_service_principal(properties: ServicePrincipalProperties) -> ServicePrincipalResponce: + """ + 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) -> ServicePrincipalResponce: + """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 ServicePrincipalResponce(name=created_service_principal.display_name, physical_resource_id=created_service_principal.id) + + +def update_service_principal( + service_principal: ServicePrincipal, + workspace_client: WorkspaceClient, +) -> ServicePrincipalResponce: + """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 ServicePrincipalResponce(name=service_principal.display_name, physical_resource_id=service_principal.id) + + +def delete_service_principal(physical_resource_id: str, workspace_url: str) -> CnfResponse: + """Delete a service pricncipal on databricks""" + workspace_client = get_workspace_client(workspace_url) + try: + workspace_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)