From 7b6b8703ba32f549b996d81053d938a16253b075 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?De=20L=C3=A8ne=20Mirouze?= Date: Fri, 4 Jun 2021 18:31:13 +0000 Subject: [PATCH] add new feature (pining, ip access list) (#174) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Setting Access List * IP access list Tests * Set-DatabricksClusterPinStatus + tests * tests * deboging some tests * cleaning code Co-authored-by: Frédéric De Lène Mirouze --- Public/Add-DatabricksIPAccessList.ps1 | 80 +++++++++++++++++++ Public/Get-DatabricksClusterPinStatus.ps1 | 31 +++++++ Public/Get-DatabricksIPAccessList.ps1 | 45 +++++++++++ Public/Get-DatabricksIPAccessListStatus.ps1 | 42 ++++++++++ Public/Remove-DatabricksIPAccessList.ps1 | 50 ++++++++++++ Public/Set-DatabricksClusterPinStatus.ps1 | 46 +++++++++++ Public/Set-DatabricksIPAccessListStatus.ps1 | 49 ++++++++++++ Tests/Add-DatabricksClusterPolicy.tests.ps1 | 2 +- Tests/Add-DatabricksIPAccessList.tests.ps1 | 63 +++++++++++++++ Tests/Get-DatabricksIPAccessList.tests.ps1 | 48 +++++++++++ ...Get-DatabricksIPAccessListStatus.tests.ps1 | 39 +++++++++ Tests/Remove-DatabricksIPAccessList.tests.ps1 | 52 ++++++++++++ .../Set-DatabricksClusterPinStatus.tests.ps1 | 64 +++++++++++++++ azure.databricks.cicd.tools.psd1 | 4 +- 14 files changed, 613 insertions(+), 2 deletions(-) create mode 100644 Public/Add-DatabricksIPAccessList.ps1 create mode 100644 Public/Get-DatabricksClusterPinStatus.ps1 create mode 100644 Public/Get-DatabricksIPAccessList.ps1 create mode 100644 Public/Get-DatabricksIPAccessListStatus.ps1 create mode 100644 Public/Remove-DatabricksIPAccessList.ps1 create mode 100644 Public/Set-DatabricksClusterPinStatus.ps1 create mode 100644 Public/Set-DatabricksIPAccessListStatus.ps1 create mode 100644 Tests/Add-DatabricksIPAccessList.tests.ps1 create mode 100644 Tests/Get-DatabricksIPAccessList.tests.ps1 create mode 100644 Tests/Get-DatabricksIPAccessListStatus.tests.ps1 create mode 100644 Tests/Remove-DatabricksIPAccessList.tests.ps1 create mode 100644 Tests/Set-DatabricksClusterPinStatus.tests.ps1 diff --git a/Public/Add-DatabricksIPAccessList.ps1 b/Public/Add-DatabricksIPAccessList.ps1 new file mode 100644 index 0000000..741cd58 --- /dev/null +++ b/Public/Add-DatabricksIPAccessList.ps1 @@ -0,0 +1,80 @@ +<# + +.SYNOPSIS + Add an IP access list. + +.DESCRIPTION +The IP Access List API enables Azure Databricks admins to configure IP allow lists and block lists for a workspace. +If the feature is disabled for a workspace, all access is allowed. +There is support for allow lists (inclusion) and block lists (exclusion). + +Be sure to check the doc before using this feature: +https://docs.microsoft.com/en-us/azure/databricks/security/network/ip-access-list + +.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 ListName + Label for this list + +.PARAMETER ListType + Either ALLOW (allow list) or BLOCK (a block list, which means exclude even if in allow list). + +.PARAMETER ListIPs + A string array of IP addresses and CIDR ranges, as String values. + + .OUTPUTS + A structure describing the new Access List IP. Looks like: + + { + "list_id": "", + "label": "office", + "ip_addresses": [ + "1.1.1.1", + "2.2.2.2/21" + ], + "address_count": 2, + "list_type": "ALLOW", + "created_at": 1578423494457, + "created_by": 6476783916686816, + "updated_at": 1578423494457, + "updated_by": 6476783916686816, + "enabled": true + } +#> + +Function Add-DatabricksIPAccessList { + [cmdletbinding()] + param ( + [parameter(Mandatory = $true, ParameterSetName = 'Bearer')] + [string]$BearerToken, + + [parameter(Mandatory = $false, ParameterSetName = 'Bearer')] + [parameter(Mandatory = $false, ParameterSetName = 'AAD')] + [string]$Region, + + [parameter(Mandatory = $true)][string]$ListName, + [parameter(Mandatory = $true, HelpMessage = "Enter an operation type: ALLOW or BLOCK")][string] + [ValidateSet("ALLOW", "BLOCK")] + $ListType, + [parameter(Mandatory = $true)][string[]]$ListIPs + ) + + [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12 + $Headers = GetHeaders $PSBoundParameters + + $URI = "$global:DatabricksURI/api/2.0/ip-access-lists" + + $Body = @{ + label = $ListName + list_type = $ListType + ip_addresses = $ListIPs + } + $BodyText = $Body | ConvertTo-Json -Depth 10 + + $response = Invoke-RestMethod -Method Post -Uri $URI -Headers $Headers -Body $BodyText + return $response.ip_access_list +} \ No newline at end of file diff --git a/Public/Get-DatabricksClusterPinStatus.ps1 b/Public/Get-DatabricksClusterPinStatus.ps1 new file mode 100644 index 0000000..1b06669 --- /dev/null +++ b/Public/Get-DatabricksClusterPinStatus.ps1 @@ -0,0 +1,31 @@ +<# + +.SYNOPSIS +Return information about all pinned clusters, active clusters, up to 100 of the most recently terminated all-purpose clusters in the past 30 days, and up to 30 of the most recently terminated job clusters in the past 30 days. + +.DESCRIPTION +Return information about all pinned clusters, active clusters, up to 100 of the most recently terminated all-purpose clusters in the past 30 days, and up to 30 of the most recently terminated job clusters in the past 30 days. + +.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 + + +#> + +Function Get-DatabricksClusterPinStatus { + [cmdletbinding()] + param ( + [parameter(Mandatory = $false)][string]$BearerToken, + [parameter(Mandatory = $false)][string]$Region + ) + + $Headers = GetHeaders $PSBoundParameters + $response = Invoke-RestMethod -Method Get -Body $body -Uri "$global:DatabricksURI/api/2.0/clusters/list" -Headers $Headers + + return $response.clusters +} + + diff --git a/Public/Get-DatabricksIPAccessList.ps1 b/Public/Get-DatabricksIPAccessList.ps1 new file mode 100644 index 0000000..062bdb2 --- /dev/null +++ b/Public/Get-DatabricksIPAccessList.ps1 @@ -0,0 +1,45 @@ +<# + +.SYNOPSIS + Add an IP access list. + +.DESCRIPTION +The IP Access List API enables Azure Databricks admins to configure IP allow lists and block lists for a workspace. +If the feature is disabled for a workspace, all access is allowed. +There is support for allow lists (inclusion) and block lists (exclusion). + +Be sure to check the doc before using this feature: +https://docs.microsoft.com/en-us/azure/databricks/security/network/ip-access-list + +.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 + +.OUTPUTS + List of defined IP Access list + See documentation +#> + +Function Get-DatabricksIPAccessList { + [cmdletbinding()] + param ( + [parameter(Mandatory = $true, ParameterSetName = 'Bearer')] + [string]$BearerToken, + + [parameter(Mandatory = $false, ParameterSetName = 'Bearer')] + [parameter(Mandatory = $false, ParameterSetName = 'AAD')] + [string]$Region + ) + + [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12 + $Headers = GetHeaders $PSBoundParameters + + $accessLists = $null + + $response = Invoke-RestMethod -Method Get -Body $body -Uri "$global:DatabricksURI/api/2.0/ip-access-lists" -Headers $Headers + $accessLists = $response.ip_access_lists + + return $accessLists +} \ No newline at end of file diff --git a/Public/Get-DatabricksIPAccessListStatus.ps1 b/Public/Get-DatabricksIPAccessListStatus.ps1 new file mode 100644 index 0000000..5234b71 --- /dev/null +++ b/Public/Get-DatabricksIPAccessListStatus.ps1 @@ -0,0 +1,42 @@ +<# + +.SYNOPSIS + Get if IP access list is activated for the workspace + +.DESCRIPTION +The IP Access List API enables Azure Databricks admins to configure IP allow lists and block lists for a workspace. +If the feature is disabled for a workspace, all access is allowed. +There is support for allow lists (inclusion) and block lists (exclusion). + +Be sure to check the doc before using this feature: +https://docs.microsoft.com/en-us/azure/databricks/security/network/ip-access-list + +.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 + + +#> + +Function Get-DatabricksIPAccessListStatus { + [cmdletbinding()] + param ( + [parameter(Mandatory = $true, ParameterSetName = 'Bearer')] + [string]$BearerToken, + + [parameter(Mandatory = $false, ParameterSetName = 'Bearer')] + [parameter(Mandatory = $false, ParameterSetName = 'AAD')] + [string]$Region + ) + + [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12 + $Headers = GetHeaders $PSBoundParameters + + $response = Invoke-RestMethod -Method Get ` + -Uri "$global:DatabricksURI/api/2.0/workspace-conf?keys=enableIpAccessLists" ` + -Headers $Headers + + return [boolean]::Parse($response.enableIpAccessLists) +} \ No newline at end of file diff --git a/Public/Remove-DatabricksIPAccessList.ps1 b/Public/Remove-DatabricksIPAccessList.ps1 new file mode 100644 index 0000000..c587106 --- /dev/null +++ b/Public/Remove-DatabricksIPAccessList.ps1 @@ -0,0 +1,50 @@ +<# + +.SYNOPSIS + Remove an IP access list. + +.DESCRIPTION +The IP Access List API enables Azure Databricks admins to configure IP allow lists and block lists for a workspace. +If the feature is disabled for a workspace, all access is allowed. +There is support for allow lists (inclusion) and block lists (exclusion). + +Be sure to check the doc before using this feature: +https://docs.microsoft.com/en-us/azure/databricks/security/network/ip-access-list + +.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 ListName + Label for this list + +.PARAMETER ListType + Either ALLOW (allow list) or BLOCK (a block list, which means exclude even if in allow list). + +.PARAMETER ListID + Id of the access list to delete. + +#> + +Function Remove-DatabricksIPAccessList { + [cmdletbinding()] + param ( + [parameter(Mandatory = $true, ParameterSetName = 'Bearer')] + [string]$BearerToken, + + [parameter(Mandatory = $false, ParameterSetName = 'Bearer')] + [parameter(Mandatory = $false, ParameterSetName = 'AAD')] + [string]$Region, + + [parameter(Mandatory = $true)][string]$ListID + ) + + [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12 + $Headers = GetHeaders $PSBoundParameters + + $URI = "$global:DatabricksURI/api/2.0/ip-access-lists/" + $ListID + + Invoke-RestMethod -Method Delete -Uri $URI -Headers $Headers +} \ No newline at end of file diff --git a/Public/Set-DatabricksClusterPinStatus.ps1 b/Public/Set-DatabricksClusterPinStatus.ps1 new file mode 100644 index 0000000..2937537 --- /dev/null +++ b/Public/Set-DatabricksClusterPinStatus.ps1 @@ -0,0 +1,46 @@ +<# + +.SYNOPSIS + Pin or unpin a DB cluster + +.DESCRIPTION + Pin or unpin a DB cluster + +.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 $enablePin + $true to pin, $false to unpin + +.PARAMETER $clusterId + Id of the cluster to be processed + + + +#> + +Function Set-DatabricksClusterPinStatus { + [cmdletbinding()] + param ( + [parameter(Mandatory = $false)][string]$BearerToken, + [parameter(Mandatory = $false)][string]$Region, + [parameter(Mandatory = $true)][boolean]$enablePin, + [parameter(Mandatory = $true)][string]$clusterId + ) + + $Headers = GetHeaders $PSBoundParameters + + $body = '{ "cluster_id": "' + $clusterId + '"}' + + if ($enablePin) { + Invoke-RestMethod -Method Post -Body $body -Uri "$global:DatabricksURI/api/2.0/clusters/pin" -Headers $Headers + } + else { + Invoke-RestMethod -Method Post -Body $body -Uri "$global:DatabricksURI/api/2.0/clusters/unpin" -Headers $Headers + } +} + + diff --git a/Public/Set-DatabricksIPAccessListStatus.ps1 b/Public/Set-DatabricksIPAccessListStatus.ps1 new file mode 100644 index 0000000..8862404 --- /dev/null +++ b/Public/Set-DatabricksIPAccessListStatus.ps1 @@ -0,0 +1,49 @@ +<# + +.SYNOPSIS + Enable/disable IP access list feature for the workspace + +.DESCRIPTION +The IP Access List API enables Azure Databricks admins to configure IP allow lists and block lists for a workspace. +If the feature is disabled for a workspace, all access is allowed. +There is support for allow lists (inclusion) and block lists (exclusion). + +Be sure to check the doc before using this feature: +https://docs.microsoft.com/en-us/azure/databricks/security/network/ip-access-list + +.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 enabled + $true enables IP access list feature for the workspace. + $false disables it. + +#> + +Function Set-DatabricksIPAccessListStatus { + [cmdletbinding()] + param ( + [parameter(Mandatory = $true, ParameterSetName = 'Bearer')] + [string]$BearerToken, + + [parameter(Mandatory = $false, ParameterSetName = 'Bearer')] + [parameter(Mandatory = $false, ParameterSetName = 'AAD')] + [string]$Region, + + [parameter(Mandatory=$true)][boolean]$enabled + ) + + [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12 + $Headers = GetHeaders $PSBoundParameters + + $Body = '{"enableIpAccessLists": "' + $enabled.ToString().ToLower() + '"}' + + Invoke-RestMethod -Method Patch ` + -Body $Body ` + -Uri "$global:DatabricksURI/api/2.0/workspace-conf" ` + -Headers $Headers ` + -ContentType "application/json" +} \ No newline at end of file diff --git a/Tests/Add-DatabricksClusterPolicy.tests.ps1 b/Tests/Add-DatabricksClusterPolicy.tests.ps1 index 4e85475..9676c5c 100644 --- a/Tests/Add-DatabricksClusterPolicy.tests.ps1 +++ b/Tests/Add-DatabricksClusterPolicy.tests.ps1 @@ -43,7 +43,7 @@ Describe "Add-DatabricksClusterPolicy" { catch { $errorThrown = $true } - $errorThrown | Should Be $true + $errorThrown | Should -Be $true } } diff --git a/Tests/Add-DatabricksIPAccessList.tests.ps1 b/Tests/Add-DatabricksIPAccessList.tests.ps1 new file mode 100644 index 0000000..2aee5d3 --- /dev/null +++ b/Tests/Add-DatabricksIPAccessList.tests.ps1 @@ -0,0 +1,63 @@ +param( + [ValidateSet('Bearer', 'ServicePrincipal')][string]$Mode = "Bearer" +) + +Set-Location $PSScriptRoot +Import-Module "..\azure.databricks.cicd.tools.psd1" -Force +$Config = (Get-Content '.\config.json' | ConvertFrom-Json) + +switch ($mode) { + ("Bearer") { + Connect-Databricks -Region $Config.Region -BearerToken $Config.BearerToken + } + ("ServicePrincipal") { + Connect-Databricks -Region $Config.Region -DatabricksOrgId $Config.DatabricksOrgId -ApplicationId $Config.ApplicationId -Secret $Config.Secret -TenantId $Config.TenantId + } +} + + +Describe "Add-DatabricksIPAccessList" { + + BeforeAll { + $sites = @('https://api.ipify.org', 'https://ifconfig.me/ip', 'https://ipinfo.io') + $myIP = foreach ($site in $sites) { + $return = Invoke-RestMethod -Uri $site + $ip = ([IPAddress] $return).IPAddressToString + if ($ip) { + $ip + break + } + } + } + + AfterAll { + Set-DatabricksIPAccessListStatus -enabled $false + } + + It "Can't provide unexpected ListType value" { + try { + Add-DatabricksIPAccessList -ListName "test" -ListType 'FOO' -ListIPs $myIP + } + catch { + $errorThrown = $true + } + + $errorThrown | Should -Be $true + } + + It "Can add a single IP access list" { + $name = "testList" + (Get-Random) + $response = Add-DatabricksIPAccessList -ListName $name -ListType 'ALLOW' -ListIPs $myIP + + $response.label | Should -Be $name + } + + It "Can add an array of IP access list" { + $name = "testList" + (Get-Random) + $ips = $myIP, "127.0.0.1" + $response = Add-DatabricksIPAccessList -ListName $name -ListType 'ALLOW' -ListIPs $ips + + $response.address_count | Should -Be 2 + } +} + \ No newline at end of file diff --git a/Tests/Get-DatabricksIPAccessList.tests.ps1 b/Tests/Get-DatabricksIPAccessList.tests.ps1 new file mode 100644 index 0000000..5e88ccb --- /dev/null +++ b/Tests/Get-DatabricksIPAccessList.tests.ps1 @@ -0,0 +1,48 @@ +param( + [ValidateSet('Bearer', 'ServicePrincipal')][string]$Mode = "Bearer" +) + +Set-Location $PSScriptRoot +Import-Module "..\azure.databricks.cicd.tools.psd1" -Force +$Config = (Get-Content '.\config.json' | ConvertFrom-Json) + +switch ($mode) { + ("Bearer") { + Connect-Databricks -Region $Config.Region -BearerToken $Config.BearerToken + } + ("ServicePrincipal") { + Connect-Databricks -Region $Config.Region -DatabricksOrgId $Config.DatabricksOrgId -ApplicationId $Config.ApplicationId -Secret $Config.Secret -TenantId $Config.TenantId + } +} + + +Describe "Get-DatabricksIPAccessList" { + + BeforeAll { + $name = "testList" + (Get-Random) + + $sites = @('https://api.ipify.org', 'https://ifconfig.me/ip', 'https://ipinfo.io') + $myIP = foreach ($site in $sites) { + $return = Invoke-RestMethod -Uri $site + $ip = ([IPAddress] $return).IPAddressToString + if ($ip) { + $ip + break + } + } + + Set-DatabricksIPAccessListStatus -enabled $true + } + + AfterAll { + Set-DatabricksIPAccessListStatus -enabled $false + } + + It "Add access IP list" { + Add-DatabricksIPAccessList -ListName $name -ListType 'ALLOW' -ListIPs $myIP + $accessList = Get-DatabricksIPAccessList + + $accessList.label | Should -Contain $name + } +} + \ No newline at end of file diff --git a/Tests/Get-DatabricksIPAccessListStatus.tests.ps1 b/Tests/Get-DatabricksIPAccessListStatus.tests.ps1 new file mode 100644 index 0000000..674b77e --- /dev/null +++ b/Tests/Get-DatabricksIPAccessListStatus.tests.ps1 @@ -0,0 +1,39 @@ +param( + [ValidateSet('Bearer', 'ServicePrincipal')][string]$Mode = "Bearer" +) + +Set-Location $PSScriptRoot +Import-Module "..\azure.databricks.cicd.tools.psd1" -Force +$Config = (Get-Content '.\config.json' | ConvertFrom-Json) + +switch ($mode) { + ("Bearer") { + Connect-Databricks -Region $Config.Region -BearerToken $Config.BearerToken + } + ("ServicePrincipal") { + Connect-Databricks -Region $Config.Region -DatabricksOrgId $Config.DatabricksOrgId -ApplicationId $Config.ApplicationId -Secret $Config.Secret -TenantId $Config.TenantId + } +} + + +Describe "Get-DatabricksIPAccessListStatus" { + + AfterAll { + Set-DatabricksIPAccessListStatus -enabled $false + } + + It "Get Current Status, set false" { + Set-DatabricksIPAccessListStatus -enabled $false + + $status = Get-DatabricksIPAccessListStatus + $status | Should -BeFalse + } + + It "Switch status" { + Set-DatabricksIPAccessListStatus -enabled $false + Set-DatabricksIPAccessListStatus -enabled $true + $status = Get-DatabricksIPAccessListStatus + $status | Should -BeTrue + } +} + diff --git a/Tests/Remove-DatabricksIPAccessList.tests.ps1 b/Tests/Remove-DatabricksIPAccessList.tests.ps1 new file mode 100644 index 0000000..d1fdbe5 --- /dev/null +++ b/Tests/Remove-DatabricksIPAccessList.tests.ps1 @@ -0,0 +1,52 @@ +param( + [ValidateSet('Bearer', 'ServicePrincipal')][string]$Mode = "Bearer" +) + +Set-Location $PSScriptRoot +Import-Module "..\azure.databricks.cicd.tools.psd1" -Force +$Config = (Get-Content '.\config.json' | ConvertFrom-Json) + +switch ($mode) { + ("Bearer") { + Connect-Databricks -Region $Config.Region -BearerToken $Config.BearerToken + } + ("ServicePrincipal") { + Connect-Databricks -Region $Config.Region -DatabricksOrgId $Config.DatabricksOrgId -ApplicationId $Config.ApplicationId -Secret $Config.Secret -TenantId $Config.TenantId + } +} + + +Describe "Remove-DatabricksIPAccessList" { + + BeforeAll { + $name = "testList" + (Get-Random) + + $sites = @('https://api.ipify.org', 'https://ifconfig.me/ip', 'https://ipinfo.io') + $myIP = foreach ($site in $sites) { + $return = Invoke-RestMethod -Uri $site + $ip = ([IPAddress] $return).IPAddressToString + if ($ip) { + $ip + break + } + } + + Set-DatabricksIPAccessListStatus -enabled $true + Add-DatabricksIPAccessList -ListName $name -ListType 'ALLOW' -ListIPs $myIP + } + + AfterAll { + Set-DatabricksIPAccessListStatus -enabled $false + } + + It "Remove all access IP" { + $accessList = Get-DatabricksIPAccessList + foreach ($item in $accessList) { + Remove-DatabricksIPAccessList -ListID $item.list_id + } + + $accessList = Get-DatabricksIPAccessList + $accessList.Count | Should -Be 0 + } +} + \ No newline at end of file diff --git a/Tests/Set-DatabricksClusterPinStatus.tests.ps1 b/Tests/Set-DatabricksClusterPinStatus.tests.ps1 new file mode 100644 index 0000000..ce09ad5 --- /dev/null +++ b/Tests/Set-DatabricksClusterPinStatus.tests.ps1 @@ -0,0 +1,64 @@ +param( + [ValidateSet('Bearer', 'ServicePrincipal')][string]$Mode = "Bearer" +) + +Set-Location $PSScriptRoot +Import-Module "..\azure.databricks.cicd.tools.psd1" -Force +$Config = (Get-Content '.\config.json' | ConvertFrom-Json) + +switch ($mode) { + ("Bearer") { + Connect-Databricks -Region $Config.Region -BearerToken $Config.BearerToken + } + ("ServicePrincipal") { + Connect-Databricks -Region $Config.Region -DatabricksOrgId $Config.DatabricksOrgId -ApplicationId $Config.ApplicationId -Secret $Config.Secret -TenantId $Config.TenantId + } +} + +$ClusterName = "TestCluster7" +$SparkVersion = "7.3.x-scala2.12" +$NodeType = "Standard_D3_v2" +$MinNumberOfWorkers = 1 +$MaxNumberOfWorkers = 1 +$Spark_conf = @{"spark.speculation" = $true; "spark.streaming.ui.retainedBatches" = 5 } +$CustomTags = @{CreatedBy = "SimonDM" } #;NumOfNodes=2;CanDelete=$true +$InitScripts = "dbfs:/script/script1" , "dbfs:/script/script2" +$SparkEnvVars = @{SPARK_WORKER_MEMORY = "29000m" } #;SPARK_LOCAL_DIRS="/local_disk0" +$AutoTerminationMinutes = 15 +$PythonVersion = 3 +$ClusterLogPath = "dbfs:/logs/mycluster" +$AzureAttributes = @{first_on_demand = 1; availability = "SPOT_WITH_FALLBACK_AZURE"; spot_bid_max_price = -1 } + +Describe "Set-DatabricksClusterPinStatus" { + + BeforeAll { + $ClusterId = New-DatabricksCluster -ClusterName $ClusterName -SparkVersion $SparkVersion -NodeType $NodeType ` + -MinNumberOfWorkers $MinNumberOfWorkers -MaxNumberOfWorkers $MaxNumberOfWorkers ` + -Spark_conf $Spark_conf -CustomTags $CustomTags -AutoTerminationMinutes $AutoTerminationMinutes -ClusterLogPath $ClusterLogPath ` + -Verbose -SparkEnvVars $SparkEnvVars -PythonVersion $PythonVersion -InitScripts $InitScripts + } + + AfterAll { + Start-Sleep -Seconds 5 + Remove-DatabricksCluster -ClusterName $ClusterName + } + + It "Pin a cluster" { + Set-DatabricksClusterPinStatus -clusterId $ClusterId -enablePin $true + + $status = Get-DatabricksClusterPinStatus | where-object { $_.cluster_id -eq $ClusterId } + + $status.pinned_by_user_name | should -not -be $null + } + + It "Unpin a cluster" { + Set-DatabricksClusterPinStatus -clusterId $ClusterId -enablePin $true + Set-DatabricksClusterPinStatus -clusterId $ClusterId -enablePin $false + + $listStatus = Get-DatabricksClusterPinStatus + $status = Get-DatabricksClusterPinStatus | where-object { $_.cluster_id -eq $ClusterId } + + $status.pinned_by_user_name | should -be $null + } +} + \ No newline at end of file diff --git a/azure.databricks.cicd.tools.psd1 b/azure.databricks.cicd.tools.psd1 index 278f4c8..28f026f 100644 --- a/azure.databricks.cicd.tools.psd1 +++ b/azure.databricks.cicd.tools.psd1 @@ -83,7 +83,9 @@ 'Get-DatabricksPermissionLevels', 'Invoke-DatabricksAPI', 'Remove-DatabricksSecret', 'Get-DatabricksWorkspaceFolder', 'Get-DatabricksInstancePool', 'Add-DatabricksInstancePool', 'Remove-DatabricksInstancePool', 'New-DatabricksBearerToken', 'Remove-DatabricksBearerToken', 'Get-DatabricksBearerToken', 'Get-DatabricksJob', 'Export-DatabricksJobs', - 'Get-DatabricksSecretByScope', 'Get-DatabricksSCIMURL', 'Get-DatabricksClusterPolicies', 'Add-DatabricksClusterPolicy','Remove-DatabricksClusterPolicy', 'Remove-DatabricksUser', 'Get-DatabricksUsers' + 'Get-DatabricksSecretByScope', 'Get-DatabricksSCIMURL', 'Get-DatabricksClusterPolicies', 'Add-DatabricksClusterPolicy','Remove-DatabricksClusterPolicy', 'Remove-DatabricksUser', 'Get-DatabricksUsers', + 'Get-DatabricksIPAccessList', 'Add-DatabricksIPAccessList', 'Remove-DatabricksIPAccessList', 'Set-DatabricksIPAccessListStatus', 'Get-DatabricksIPAccessListStatus', + 'Set-DatabricksClusterPinStatus', 'Get-DatabricksClusterPinStatus' # Cmdlets to export from this module CmdletsToExport = '*'