warehouseClusterSchedule.py

# Databricks notebook source
import requests
import json

# COMMAND ----------

dbutils.widgets.text("workspaceUrl", "https://adb-123456789101112.13.azuredatabricks.net" ,"Workspace URL")
dbutils.widgets.text("clusterActivity", "start" ,"Cluster Activity")
dbutils.widgets.text("clusterId", "444abcdefgh666" ,"Cluster ID")

# COMMAND ----------

pat = dbutils.secrets.get("KeyVault", "DatabricksPAT")
dbxWorkspaceURL = dbutils.widgets.get("workspaceUrl")
clusterActivity = dbutils.widgets.get("clusterActivity")
clusterId = dbutils.widgets.get("clusterId")

requestHeader = {'Authorization': 'Bearer ' + pat}

# COMMAND ----------

def getClusterInfo(clusterId):
  endpoint = f'/api/2.0/sql/warehouses/{clusterId}'
  response = requests.get(dbxWorkspaceURL + endpoint,headers=requestHeader)
  return json.loads(response.content)

# COMMAND ----------

def startCluster(clusterId):
  endpoint = f'/api/2.0/sql/warehouses/{clusterId}/start'
  clusterInfo = getClusterInfo(clusterId)
  if clusterInfo['state'] == 'RUNNING':
    return f"{clusterId} is already running"
  else:
    response = requests.post(dbxWorkspaceURL + endpoint,headers=requestHeader)
    return json.loads(response.content)

# COMMAND ----------

def listClusters():
  endpoint = '/api/2.0/sql/warehouses/'
  response = requests.get(dbxWorkspaceURL + endpoint,headers=requestHeader)
  return json.loads(response.content)

# COMMAND ----------

def terminateCluster(clusterId):
  endpoint = f'/api/2.0/sql/warehouses/{clusterId}/stop'
  clusterInfo = getClusterInfo(clusterId)
  if clusterInfo['state'] == 'STOPPED':
    return f"{clusterId} is already stopped"
  else:
    response = requests.post(dbxWorkspaceURL + endpoint,headers=requestHeader)
    return json.loads(response.content)

# COMMAND ----------

if clusterId == '':
  clusters = listClusters()
  for cluster in clusters["warehouses"]:
    if clusterActivity == "start":
      response = startCluster(clusterId)
      print(response)
    if clusterActivity == "stop":
      response = terminateCluster(clusterId)
      print(response)
else:
  if clusterActivity == "start":
    response = startCluster(clusterId)
    print(response)
  if clusterActivity == "stop":
    response = terminateCluster(clusterId)
    print(response)
  if clusterActivity == "info":
    response = getClusterInfo(clusterId)
    print(response)