Public/Connect-Databricks.ps1

<#
 
.SYNOPSIS
    Connects your current PowerShell session to Azure Databricks.
 
.DESCRIPTION
    Connects your current PowerShell session to Azure Databricks.
    Supports Service Princial AAD authenication or via Databricks Bearer Token
 
.PARAMETER BearerToken
    Your Databricks Bearer token to authenticate to your workspace (see User Settings in Databricks WebUI)
 
.PARAMETER Region
    Azure Region - must match the URL of your Databricks workspace, example northeurope
 
.PARAMETER ApplicationId
    Azure Active Directory Service Principal Client ID (also known as Application ID)
     
.PARAMETER Secret
    Secret for given Client ID
 
.PARAMETER DatabricksOrgId
    Databricks OrganisationID this is found in the URL of your Worksapce as the o parameters (example o=123456789). Note the first time a service principal connects it must use the MANAGEMENT method (ie provide the Resource GRoup Name and Workspace Name - as this provisions the user)
 
.PARAMETER TenantId
    Tenant Id (Directory ID) for the AAD owning the ApplicationId
 
.PARAMETER SubscriptionId
    Subscription ID for the Workspace
 
.PARAMETER ResourceGroupName
    Resource Group Name for the Workspace
 
.PARAMETER WorkspaceName
    Workspace Name
 
.PARAMETER Force
    Removes any cached credentials and reconnects
 
.EXAMPLE
    C:\PS> Connect-Databricks -Region "westeurope" -ApplicationId "8a686772-0e5b-4cdb-ad19-bf1d1e7f89f3" -Secret "myPrivateSecret" -DatabricksOrgId 1234567 -TenantId "8a686772-0e5b-4cdb-ad19-bf1d1e7f89f3"
 
    This example of a DIRECT connection (using the Databricks organisation Id)
 
.EXAMPLE
    C:\PS> Connect-Databricks -Region "westeurope" -ApplicationId "8a686772-0e5b-4cdb-ad19-bf1d1e7f89f3" -Secret "myPrivateSecret" -ResourceGroupName "MyResourceGroup" -SubscriptionId "9a686882-0e5b-4edb-cd49-cf1f1e7f34d9" -WorkspaceName "workspaceName" -TenantId "8a686772-0e5b-4cdb-ad19-bf1d1e7f89f3"
 
    This example of a MANAGMENT connection (using the Azure resource identifiers to connect)
 
.EXAMPLE
    C:\PS> Connect-Databricks -BearerToken "dapi1234567890" -Region "westeurope"
 
    This example of a BEARER connection (using the Databricks Bearer token from the Web UI to login as a person)
 
.NOTES
    Author: Simon D'Morias / Data Thirst Ltd
 
#>


Function Connect-Databricks {  
    [cmdletbinding(DefaultParameterSetName = 'Bearer')]
    param (
        [parameter(Mandatory = $true, ParameterSetName = 'Bearer')]
        [string]$BearerToken,

        [parameter(Mandatory = $true, ParameterSetName = 'Bearer')]
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithOrgId')]
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithResource')]
        [string]$Region,
        [parameter(Mandatory = $false, ParameterSetName = 'Bearer')]
        [parameter(Mandatory = $false, ParameterSetName = 'AADwithOrgId')]
        [parameter(Mandatory = $false, ParameterSetName = 'AADwithResource')]
        [string]$DatabricksURISuffix = "azuredatabricks.net" ,
        [parameter(Mandatory = $false, ParameterSetName = 'AADwithOrgId')]
        [parameter(Mandatory = $false, ParameterSetName = 'AADwithResource')]
        [string]$oauthLogin = "login.microsoftonline.com" ,
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithOrgId')]
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithResource')]
        [string]$ApplicationId,
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithOrgId')]
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithResource')]
        [string]$Secret,

        [parameter(Mandatory = $true, ParameterSetName = 'AADwithOrgId')]
        [string]$DatabricksOrgId,

        [parameter(Mandatory = $true, ParameterSetName = 'AADwithOrgId')]
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithResource')]
        [string]$TenantId,
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithResource')]
        [string]$SubscriptionId,
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithResource')]
        [string]$ResourceGroupName,
        [parameter(Mandatory = $true, ParameterSetName = 'AADwithResource')]
        [string]$WorkspaceName,

        [parameter(Mandatory = $false, ParameterSetName = 'AADwithOrgId')]
        [parameter(Mandatory = $false, ParameterSetName = 'AADwithResource')]
        [switch]$Force,
        [switch]$TestConnectDatabricks
    ) 

    Write-Verbose "Globals at start of Connect:" 
    Write-Globals

    if ($Force) {
        Write-Verbose "-Force set - clearing global variables"
        Set-GlobalsNull
    }

    [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12
    $AzureRegion = $Region.Replace(" ", "")
    $AzureDatabricksURISuffix = $DatabricksURISuffix.Trim(".", " ").Replace(" ", "")
    $AzureOauthLogin = $oauthLogin.Trim("/", " ").Replace(" ", "")
    $URI = "https://$AzureOauthLogin/$tenantId/oauth2/token/"
    if ($PSCmdlet.ParameterSetName -eq "Bearer") {
        Set-GlobalsNull
        # Use Databricks Bearer Token Method
        $global:DatabricksAccessToken = "Bearer $BearerToken"
        # Basically do not expire the token
        $global:DatabricksTokenExpires = (Get-Date).AddDays(90)
        $global:Headers = @{"Authorization" = "$global:DatabricksAccessToken" }
    }
    elseif ($PSCmdlet.ParameterSetName -eq "AADwithOrgId") {
        Get-AADDatabricksToken
        $global:Headers = @{"Authorization" = "Bearer $DatabricksAccessToken";
            "X-Databricks-Org-Id"           = "$DatabricksOrgId"
        }
        $global:DatabricksOrgId = $DatabricksOrgId
    }
    elseif ($PSCmdlet.ParameterSetName -eq "AADwithResource") {
        Get-AADManagementToken
        Get-AADDatabricksToken
        $global:Headers = @{"Authorization"            = "Bearer $global:DatabricksAccessToken";
            "X-Databricks-Azure-SP-Management-Token"   = $global:ManagementAccessToken;
            "X-Databricks-Azure-Workspace-Resource-Id" = "/subscriptions/$SubscriptionId/resourceGroups/$ResourceGroupName/providers/Microsoft.Databricks/workspaces/$WorkspaceName"
        }
    }
    $global:DatabricksURI = "https://$AzureRegion.$AzureDatabricksURISuffix"
    Write-Verbose "Globals at end of Connect:" 
    Write-Globals
    if ($PSBoundParameters.ContainsKey('TestConnectDatabricks')) {
        Write-Verbose "Connecting to Workspace to verify connection details are correct:" 
        if ($PSCmdlet.ParameterSetName -eq "Bearer") {
            Test-ConnectDatabricks -Region $AzureRegion -BearerToken $BearerToken
        }
        else {
            Test-ConnectDatabricks
        }
    }
}