Public/Connect-Databricks.ps1

function Connect-Databricks {
  [CmdletBinding()]
  param(
    [Parameter(Mandatory = $true, HelpMessage = 'Workspace or account URL')]
    [ValidateNotNullOrEmpty()]
    [string] $Url,

    [Parameter(Mandatory = $false)]
    [switch] $Workspace,

    [Parameter(Mandatory = $false)]
    [switch] $Account,

    [Parameter(Mandatory = $false, HelpMessage = 'Account ID, for example: 1234')]
    [string] $Id
  )

  $ParsedUrl = $null
  if (-not [System.Uri]::TryCreate($Url, [System.UriKind]::Absolute, [ref]$ParsedUrl)) {
    throw 'Url must be an absolute URL, for example: https://myworkspace.cloud.databricks.com'
  }

  if ([string]::IsNullOrWhiteSpace($ParsedUrl.Host)) {
    throw 'Url must include a host.'
  }

  $Url = $ParsedUrl.GetLeftPart([System.UriPartial]::Authority)

  $AccountIdFromUrl = $null
  if ($ParsedUrl.Query) {
    foreach ($Pair in $ParsedUrl.Query.TrimStart('?').Split('&')) {
      if (-not $Pair) {
        continue
      }

      $Parts = $Pair.Split('=', 2)
      $Name = [System.Uri]::UnescapeDataString($Parts[0])
      if ($Name -ieq 'account_id' -and $Parts.Count -gt 1) {
        $AccountIdFromUrl = [System.Uri]::UnescapeDataString($Parts[1])
        break
      }
    }
  }

  if (-not [string]::IsNullOrWhiteSpace($AccountIdFromUrl) -and [string]::IsNullOrWhiteSpace($Id)) {
    $Id = $AccountIdFromUrl
  }

  $IsAccountHost = $ParsedUrl.Host -ieq 'accounts.cloud.databricks.com'
  $UseAccountAuth = $Account -or $IsAccountHost -or -not [string]::IsNullOrWhiteSpace($Id)

  if ($UseAccountAuth) {
    if ([string]::IsNullOrWhiteSpace($Id)) {
      throw 'Account login requires account ID. Provide -Id or include account_id in URL query.'
    }

    & databricks auth login --host $Url --account-id $Id
  }
  else {
    & databricks auth login --host $Url
  }
}