Courtesy of my colleague Sri, here's some sample library code to execute on a databricks cluster with a short SLA
import logging
import textwrap
import time
from typing import Text
from databricks_cli.sdk import ApiClient, ClusterService
# Create a custom logging
logging.basicConfig(format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S')
logging.getLogger().setLevel(logging.INFO)
class ClusterExecutionContext(object):
def __init__(self, api_client_1_2: ApiClient, cluster_service: ClusterService, cluster_id: Text, language: Text, context_id: Text = None):
self.context_id = context_id
self.language = language
self.cluster_id = cluster_id
self.api_client_1_2 = api_client_1_2
self.cluster_service = cluster_service
def __enter__(self):
logging.info("Creating Cluster Execution Context")
self._create_context()
self.wait_for_context_ready()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
logging.info("Exiting Cluster Execution Context Deleting Context")
self._delete_context()
def _verify_cluster_running_state(self):
cluster_resp = self.cluster_service.get_cluster(self.cluster_id)
if cluster_resp["state"] != "RUNNING":
raise ValueError("State is not in Running cannot generate context!")
def _verify_context_exists(self):
if self.context_id is None or len(self.context_id) <= 0:
raise ValueError("Context ID does not exist.")
def _create_context(self):
self._verify_cluster_running_state()
data = {"language": self.language, "clusterId": self.cluster_id}
resp = self.api_client_1_2.perform_query("POST", "/contexts/create", data)
self.context_id = resp["id"]
logging.debug(resp)
return resp
def _read_context(self):
self._verify_cluster_running_state()
self._verify_context_exists()
data = {"contextId": self.context_id, "clusterId": self.cluster_id}
resp = self.api_client_1_2.perform_query("GET", "/contexts/status", data)
logging.debug(resp)
return resp
def wait_for_context_ready(self):
self._verify_cluster_running_state()
self._verify_context_exists()
# Will attempt 5 tries with a 10 second delay
for i in range(100):
context = self._read_context()
status = context["status"]
logging.debug("Context is in a {} state...".format(status))
if status == "Running":
logging.info("Context is in running state ready to run commands!")
return
time.sleep(0.5)
raise Exception("Context failed to be created after 50 seconds")
def run_command(self, command_string, timeout_seconds=10):
with ClusterCommandExecutionContext(cluster_execution_context=self, command_string=command_string,
timeout_seconds=timeout_seconds) as ccec:
return ccec
def _delete_context(self):
self._verify_cluster_running_state()
self._verify_context_exists()
data = {"contextId": self.context_id, "clusterId": self.cluster_id}
resp = self.api_client_1_2.perform_query("POST", "/contexts/destroy", data)
return resp
class ClusterCommandExecutionContext(object):
def __init__(self, cluster_execution_context: ClusterExecutionContext, command_string: Text, timeout_seconds: int, command_id: Text=None):
# assumes to be a running context
self.command_id = command_id
self.timeout_seconds = timeout_seconds
self.command_string = command_string
self.cluster_execution_context = cluster_execution_context
def __enter__(self):
self._create_command()
try:
self._wait_for_command_finish()
return self._read_command()
except TimeoutError as te:
logging.info("Command Timed out: {}".format(te))
self._delete_command()
self._wait_for_command_finish()
logging.info("Command Canceled Successfully!")
return None
def __exit__(self, exc_type, exc_val, exc_tb):
logging.debug("Exiting Command Execution Context")
def _verify_cluster_running_state(self):
cluster_resp = self.cluster_execution_context.cluster_service.get_cluster(self.cluster_execution_context.cluster_id)
if cluster_resp["state"] != "RUNNING":
raise ValueError("State is not in Running cannot generate context!")
def _verify_command_exists(self):
if self.command_id is None or len(self.command_id) <= 0:
raise ValueError("Command ID does not exist.")
def _create_command(self):
self._verify_cluster_running_state()
data = {"language": self.cluster_execution_context.language,
"clusterId": self.cluster_execution_context.cluster_id,
"contextId": self.cluster_execution_context.context_id,
"command": self.command_string}
resp = self.cluster_execution_context.api_client_1_2.perform_query("POST", "/commands/execute", data)
self.command_id = resp["id"]
logging.debug(resp)
return resp
def _read_command(self):
self._verify_cluster_running_state()
self._verify_command_exists()
data = {
"clusterId": self.cluster_execution_context.cluster_id,
"contextId": self.cluster_execution_context.context_id,
"commandId": self.command_id}
resp = self.cluster_execution_context.api_client_1_2.perform_query("GET", "/commands/status", data)
logging.debug(resp)
return resp
def _wait_for_command_finish(self):
self._verify_cluster_running_state()
self._verify_command_exists()
# Will attempt 5 tries with a 10 second delay
end_states = ["Finished", "Cancelled", "Error"]
for i in range(self.timeout_seconds):
command = self._read_command()
status = command["status"]
if i % 10 == 0:
logging.debug("Command is in a {} state...".format(status))
if status in end_states:
logging.debug("Command is in a terminal state: {}".format(status))
return
time.sleep(0.5)
raise TimeoutError("Command timed out after: {} seconds".format(self.timeout_seconds))
def _delete_command(self):
self._verify_cluster_running_state()
data = {"clusterId": self.cluster_execution_context.cluster_id,
"contextId": self.cluster_execution_context.context_id,
"commandId": self.command_id}
resp = self.cluster_execution_context.api_client_1_2.perform_query("POST", "/commands/cancel", data)
return resp
To use code above use the following
host = ""
token = ""
cluster = ""
language = "python"
command = """
import json
import time
import requests
# raise Exception("hello world")
display(spark.sql("SHOW TABLES IN default"))
"""
command2 = """
import json
import time
display(spark.sql("SELECT * FROM airlines_gold"))
"""
db_api_client = ApiClient(host=host, token=token, apiVersion="1.2")
cluster_service = ClusterService(ApiClient(host=host, token=token, apiVersion="2.0"))
with ClusterExecutionContext(api_client_1_2=db_api_client, cluster_service=cluster_service, cluster_id=cluster, language=language) as cec:
for i in range(10):
data = cec.run_command(textwrap.dedent(command), timeout_seconds=12)
error_is_true = True if data.get("results", {}).get("resultType", None) == "error" else False
if error_is_true:
raise Exception(data.get("results", {}).get("summary", None))
logging.info(data)
