Databricks : Il se connecte aux API en mode Azure AD sans PAT tout le monde hallucine (PRANK)

De base pour se connecter à un workspace Databricks et utiliser les API, il faut utiliser un token appelé aussi PAT que l’on va passer dans le header des appels REST (cf. mon article sur le CI/CD et Databricks dans lequel j’utilise cette méthode : ici )

Voir ici pour s’authentifier avec un PAT de façon “classique” : https://docs.databricks.com/dev-tools/api/latest/authentication.html 

Seulement les problèmes avec ce PAT :

  • c’est d’abord qu’il est attaché à un utilisateur et à un workspace donc pour de l’administration et de la sécurité c’est pas top. Si l’utilisateur quitte l’entreprise par exemple on ne peut pas supprimer ou retirer les droits sur son compte sans tout casser.
  • de plus il n’est possible de le créer qu’à la main à partir de l’UI du workspace Databricks, ou via les API si on a déjà un PAT. Or, il n’y a pas de PAT disponible directement à la suite de la création d’un workspace. DONC c’est un gros problème si l’on cherche à automatiser une chaine complète dans un workflow CI/CD ! Par exemple, lors d’un déploiement d’un nouveau projet en production, nous ne pourrons pas créer le workspace et déployer dessus des notebooks ou configurer des clusters directement après, vu que nous n’avons pas de PAT Sad smile

Heureusement, il existe aujourd’hui un moyen de se connecter autrement qu’avec le PAT : directement avec une authentification Azure AD et un compte de service Azure qu’on appelle un service principal (SP). 

Dans cet article on va réaliser toutes les différentes actions en Powershell afin de montrer que cela est scriptable et intégrable dans Azure DevOps par exemple. Mais vu que nous sommes sur des techno REST, il est bien sûr possible de faire ça avec n’importe quel langage.

J’ajoute que dans cet article je me concentre sur l’authentification via Azure AD aux API, je pars du principe que la ressource Databricks existe déjà. Mais nous verrons dans de prochains articles comment automatiser la création des ressources et d’architectures complètes via des templates ARM ou Terraform.

C’est parti, il va donc falloir :

Créer un SP dans Azure AD et enregistrer son application ID et son secret key :

# Connect to Azure
Connect-AzAccount

# Create SP if not exists
if ($null -eq (Get-AzADServicePrincipal -DisplayName $ServicePrincipleName)) {
    $mySP = New-AzADServicePrincipal -DisplayName $servicePrincipleName -ErrorAction 'Stop'
    $ClientId = $mySP.ApplicationId
    $ClientSecret = [pscredential]::new($servicePrincipleName, $mySP.Secret).GetNetworkCredential().Password
}

Mettre le service principal en owner (propriétaire) du workspace Databricks :

# Set RBAC on Databricks workspace

New-AzRoleAssignment -ApplicationId $ClientId `
    -RoleDefinitionName "Owner" `
    -ResourceGroupName $ResourceGroupName `
    -ResourceName  $WorkspaceName `
    -ResourceType "Microsoft.Databricks/workspaces"

Se connecter au service au workspace Databricks avec le service principal, pour cela il va falloir dans un premier temps faire 2 appels REST au à l’entité d’authentification Azure avec votre SP : https://login.microsoftonline.com/<TenantId>/oauth2/token et récupérer 2 tokens:

– un token pour le service correspondant à l’application AD du service Databricks enregistré nativement dans l’Azure AD de votre tenant. Généralement, cette application a l’id 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d, mais il est aussi possible de retrouver cet id dans Azure AD en recherchant l’application AzureDatabricks :

VUw7BcTyBm

– un token pour le service de gestion Azure https://management.core.windows.net/ (qui sert d’ailleurs pour toutes les connexions classiques aux API Azure)

# Get AzureDatabricks app token
$RequestAccessTokenUri = "https://login.microsoftonline.com/$TenantId/oauth2/token"
$Resource = "https://management.core.windows.net/"
$DBXressource = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d" # CF. AzureDatabricks AzureAD application 

$body = "grant_type=client_credentials&client_id=$ClientId&client_secret=$ClientSecret&resource=$DBXressource"

$Token = Invoke-RestMethod -Method Post -Uri $RequestAccessTokenUri -Body $body -ContentType 'application/x-www-form-urlencoded'

Write-Host "Print Token" -ForegroundColor Green
Write-Output $Token.access_token
$apiKey = $Token.access_token

# Get Azure Management token

$bodyManagement = "grant_type=client_credentials&client_id=$ClientId&client_secret=$ClientSecret&resource=$Resource"

$Token = Invoke-RestMethod -Method Post -Uri $RequestAccessTokenUri -Body $bodyManagement -ContentType 'application/x-www-form-urlencoded'

Write-Host "Print Token" -ForegroundColor Green
Write-Output $Token.access_token
$apiKeyManagement = $Token.access_token

Enfin avec ces 2 tokens, vous allez pouvoir faire des appels aux API Databricks en passant ces tokens dans le header de vos appels :

# Call Azure Databricks API

$headers = @{
    "Authorization"="Bearer $apiKey";
    "X-Databricks-Azure-SP-Management-Token"=$apiKeyManagement;
    "X-Databricks-Azure-Workspace-Resource-Id"="/subscriptions/$SubscriptionId/resourceGroups/$ResourceGroupName/providers/Microsoft.Databricks/workspaces/$WorkspaceName"
}
$uri = "$uriroot/2.0/dbfs/list?path=/"
Invoke-RestMethod -Method 'Get' -Uri $uri -Headers $headers 

Voici le script Powershell complet :

# Variables
$dbxurl = "https://westeurope.azuredatabricks.net" # Modify the region : westeurope, northeurope...
$uriroot = "$dbxurl/api" 

$ServicePrincipleName = "MySP"
$TenantId = "" # Enter Tenant Id.
$ClientId = "" # Enter Client Id if exists
$ClientSecret = "" # Enter Client Secret if exists
$SubscriptionId = "" # Enter Subscription ID
$DBXressource = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d" # CF. AzureDatabricks AzureAD application 
$ResourceGroupName = "" # Enter the RG name where the Databricks ressource is
$WorkspaceName = "" # Enter the name of the Databricks ressource
$Resource = "https://management.core.windows.net/"

$RequestAccessTokenUri = "https://login.microsoftonline.com/$TenantId/oauth2/token"

# Connect to Azure
Connect-AzAccount

# Create SP if not exists
if ($null -eq (Get-AzADServicePrincipal -DisplayName $ServicePrincipleName)) {
    $mySP = New-AzADServicePrincipal -DisplayName $servicePrincipleName -ErrorAction 'Stop'
    $ClientId = $mySP.ApplicationId
    $ClientSecret = [pscredential]::new($servicePrincipleName, $mySP.Secret).GetNetworkCredential().Password
}

# Set RBAC on Databricks workspace

New-AzRoleAssignment -ApplicationId $ClientId `
    -RoleDefinitionName "Owner" `
    -ResourceGroupName $ResourceGroupName `
    -ResourceName  $WorkspaceName `
    -ResourceType "Microsoft.Databricks/workspaces"

# Get AzureDatabricks app token

$body = "grant_type=client_credentials&client_id=$ClientId&client_secret=$ClientSecret&resource=$DBXressource"

$Token = Invoke-RestMethod -Method Post -Uri $RequestAccessTokenUri -Body $body -ContentType 'application/x-www-form-urlencoded'

Write-Host "Print Token" -ForegroundColor Green
Write-Output $Token.access_token
$apiKey = $Token.access_token

# Get Azure Management token

$bodyManagement = "grant_type=client_credentials&client_id=$ClientId&client_secret=$ClientSecret&resource=$Resource"

$Token = Invoke-RestMethod -Method Post -Uri $RequestAccessTokenUri -Body $bodyManagement -ContentType 'application/x-www-form-urlencoded'

Write-Host "Print Token" -ForegroundColor Green
Write-Output $Token.access_token
$apiKeyManagement = $Token.access_token

# Call Azure Databricks API

$headers = @{
    "Authorization"="Bearer $apiKey";
    "X-Databricks-Azure-SP-Management-Token"=$apiKeyManagement;
    "X-Databricks-Azure-Workspace-Resource-Id"="/subscriptions/$SubscriptionId/resourceGroups/$ResourceGroupName/providers/Microsoft.Databricks/workspaces/$WorkspaceName"
}
$uri = "$uriroot/2.0/dbfs/list?path=/"
Invoke-RestMethod -Method 'Get' -Uri $uri -Headers $headers 

Et voila !

Une bonne pratique par la suite est d’aller générer un PAT Databricks classique avec cette méthode via l’API “api/2.0/token/create” (la doc ici) et l’enregistrer dans un secret Azure Key Vault pour que d’autre applicatifs l’utilisent comme Azure Data Factory par exemple.

Pour comprendre ce mode de connexion, j’ai du faire du reverse engineering sur le module Powershell azure.databricks.cicd.tools qui propose aussi une authentification Azure AD. Je vous recommande d’ailleurs vraiment ce module qui est très pratique : https://www.powershellgallery.com/packages/azure.databricks.cicd.tools/

FADATA

Fabien Adato est Consultant Data et BI chez AZEO. Après avoir intégré la société CGI Business Consulting où il rejoint une équipe dédiée à la Business Intelligence, il fait ses premières armes sur la solution Microsoft SQL Server. Pendant plus de 4 ans, il acquiert des compétences techniques et fonctionnelles sur toute la chaîne de valeur de la BI (ETL, Base de Données, Cube Olap et Rapport). Il a pu se spécialiser également sur de nouvelles technologies de la Data et notamment Hadoop, Pig et Hive, ainsi que des technologies self-service BI, le No-SQL comme la base de données MongoDB et le moteur d’indexation ElasticSearch. Il entre chez AZEO en 2014 pour y assurer le poste de Consultant DATA / BI, autour des technologies Microsoft SQL Server, Power BI et Azure Cortana Intelligence.

Laisser un commentaire

Votre adresse de messagerie ne sera pas publiée. Les champs obligatoires sont indiqués avec *