Public/CommandExecution.ps1

Function Start-DatabricksCommand {
  <#
      .SYNOPSIS
      Run a command or file using an existing Execution Context.
      .DESCRIPTION
      Run a command or file using an existing Execution Context.
      Official API Documentation: https://docs.databricks.com/dev-tools/api/1.2/index.html#command-execution-1
      .PARAMETER ClusterID
      The id of an existing cluster you want to use.
      .PARAMETER ContextID
      The id of the context on which you want to run the command/file.
      .PARAMETER Language
      The language for which you want to create an execution context.
      .PARAMETER Command
      The actual command you want to run.
      .EXAMPLE
      Start-DatabricksCommand -ClusterID "1202-211320-brick1" -Language Scala -ContextID 6317282514101885389 -Command "sc.parallelize(1 to 10).collect"
  #>

  [CmdletBinding()]
  param
  (
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("cluster_id")] [string] $ClusterID, 
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [ValidateSet('scala', 'python', 'sql')] [string] $Language,
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("context_id")] [string] $ContextID,
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("cmd")] [string] $Command
  )
  begin {
    $requestMethod = "POST"
    $apiEndpoint = "/1.2/commands/execute"
  }
    
  process {    
    Write-Verbose "Building Body/Parameters for final API call ..."
    #Set parameters
    $parameters = @{
      clusterId = $ClusterID
      language = $Language
      contextId = $ContextID
      command = $command
    }

    $result = Invoke-DatabricksApiRequest -Method $requestMethod -EndPoint $apiEndpoint -Body $parameters
        
    # add original parameters so the returned value can be used properly for pipelining
    # need to use JSON conversion to get a PSCustomObject instead of HashTable
    $ret = [PSCustomObject]@{
      command_id = $result.id
      context_id = $ContextID
      cluster_id = $ClusterID
      language = $Language
    }

    return $ret
  }
}


Function Get-DatabricksCommand {
  <#
      .SYNOPSIS
      Show the status of an existing command.
      .DESCRIPTION
      Show the status of an existing command.
      Official API Documentation: https://docs.databricks.com/dev-tools/api/1.2/index.html#command-execution-1
      .PARAMETER ClusterID
      The id of an existing cluster you want to use.
      .PARAMETER ContextID
      The id of the context for which to retrieve the status.
      .PARAMETER CommandID
      The id of the command for which to retrieve the status.
      .EXAMPLE
      Get-DatabricksCommand -ClusterID "1202-211320-brick1" -ContextId 221312313123 -CommandID 5220029674192230006
  #>

  [CmdletBinding()]
  param
  (
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("cluster_id")] [string] $ClusterID, 
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("context_id")] [string] $ContextID,
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("command_id")] [string] $CommandID
  )
  begin {
    $requestMethod = "GET"
    $apiEndpoint = "/1.2/commands/status"
  }
    
  process {    
    Write-Verbose "Building Body/Parameters for final API call ..."
    #Set parameters
    $parameters = @{
      clusterId = $ClusterID
      contextId = $ContextID
      commandId = $CommandID
    }

    $result = Invoke-DatabricksApiRequest -Method $requestMethod -EndPoint $apiEndpoint -Body $parameters
        
    return $result
  }
}

Function Stop-DatabricksCommand {
  <#
      .SYNOPSIS
      Stops/Cancels a Databricks Command
      .DESCRIPTION
      Stops/Cancels a Databricks Command
      Official API Documentation: https://docs.databricks.com/dev-tools/api/1.2/index.html#cancel-a-command
      .PARAMETER ClusterID
      The id of of the cluster on which you want to cancel the command.
      .PARAMETER ContextID
      The id of the context on which you want to cancel the command.
      .PARAMETER CommandID
      The id of the command you wan tto cancel
      .EXAMPLE
      Stop-DatabricksCommand -ClusterID "1202-211320-brick1" -ContextId 221312313123 -CommandID 5220029674192230006
  #>

  [CmdletBinding()]
  param
  (
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("cluster_id")] [string] $ClusterID, 
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("context_id")] [string] $ContextID,
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("command_id")] [string] $CommandID
  )
  begin {
    $requestMethod = "POST"
    $apiEndpoint = "/1.2/commands/cancel"
  }
    
  process {    
    Write-Verbose "Building Body/Parameters for final API call ..."
    #Set parameters
    $parameters = @{
      clusterId = $ClusterID
      contextId = $ContextID
      commandId = $CommandID
    }

    $result = Invoke-DatabricksApiRequest -Method $requestMethod -EndPoint $apiEndpoint -Body $parameters
        
    return $result
  }
}


Function Get-DatabricksCommandResult {
  <#
      .SYNOPSIS
      Show the status of an existing command.
      .DESCRIPTION
      Show the status of an existing command.
      Official API Documentation: https://docs.databricks.com/dev-tools/api/1.2/index.html#command-execution-1
      .PARAMETER ClusterID
      The id of an existing cluster you want to use.
      .PARAMETER CommandID
      The id of the command for which to retrieve the status.
      .PARAMETER ContextID
      The id of the context for which to retrieve the status.
      .EXAMPLE
      Get-DatabricksExecutionContextStatus -ClusterID "1202-211320-brick1" -CommandID 5220029674192230006
  #>

  [CmdletBinding()]
  param
  (
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("cluster_id")] [string] $ClusterID, 
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("context_id")] [string] $ContextID,
    [Parameter(Mandatory = $true, ValueFromPipelineByPropertyName = $true)] [Alias("command_id")] [string] $CommandID,
    [Parameter(Mandatory = $false)] [switch] $AwaitCompletion
  )
  process {    
    $wait = $false
    do {
      if($wait) { Start-Sleep -Seconds 3 }
      else {$wait = $true }
      $apiResults = Get-DatabricksCommandStatus -ClusterID $ClusterID -ContextID $ContextID -CommandID $CommandID     
    } while ($AwaitCompletion -and $apiResults.status -notin @("Finished", "Cancelled", "Error"))
    

    if($apiResults.status -eq "Finished")
    {
      if($apiResults.results.resultType -eq "table")
      {
        $data = @()
        $schema = $apiResults.results.schema

        foreach($row in $apiResults.results.data)
        {
          $newRow = [PSCustomObject]@{}
          for ($i = 0; $i -lt $schema.Count; $i++) {
            $value = $row[$i] -as $schema[$i].type.trim('"')
            $newRow | Add-Member -NotePropertyName $schema[$i].name -NotePropertyValue $value
          }

          $data += $newRow
        }

        return $data
      }
      elseif($apiResults.results.resultType -eq "text") {
        return [string]($apiResults.results.data)
      }
      elseif($apiResults.results.resultType -eq "error") {
        Write-Error ($apiResults.results | ConvertTo-Json -Depth 10)
      }
    }
    elseif($apiResults.status -eq "Error")
    {
      Write-Error ($apiResults | ConvertTo-Json -Depth 10)
    }
    return $apiResults
  }
}