add new feature (pining, ip access list) (#174)

* 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 <[email protected]>
DataThirstLtd · Jun 4, 2021 · 7b6b870 · 7b6b870
1 parent 38e5e74
commit 7b6b870
Show file tree

Hide file tree

Showing 14 changed files with 613 additions and 2 deletions.
diff --git a/Public/Add-DatabricksIPAccessList.ps1 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": "<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
+}
diff --git a/Public/Get-DatabricksClusterPinStatus.ps1 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
@@ -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
+}
diff --git a/Public/Get-DatabricksIPAccessListStatus.ps1 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)
+}
diff --git a/Public/Remove-DatabricksIPAccessList.ps1 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
+}
diff --git a/Public/Set-DatabricksClusterPinStatus.ps1 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
@@ -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"
+}
diff --git 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
     }
 }