testingDBX.R

library(httr)
library(infuser)


## See https://docs.databricks.com/api/1.2/index.html#execution-context
ctxMake <- function(instance,clusterId,user,password){
  url <- infuse("https://{{instance}}.cloud.databricks.com/api/1.2/contexts/create",instance=instance)
  pyctx<-POST(url,body=list(language='python', clusterId=clusterId)
    ,encode='form'
    ,authenticate(user,password))
  pyctxId <- content(pyctx)$id
}
  
## Context Status
ctxStatus <- function(instance, ctx, clusterId,user,password){
  url <- infuse("https://{{instance}}.cloud.databricks.com/api/1.2/contexts/status",instance=instance)
  ctxStatusCall<-GET(url,query=list(clusterId=clusterId,contextId=ctx)
            ,authenticate(user,password))
  content(ctxStatusCall)
}

## Context Destroy
ctxDestroy <- function( ctx, clusterId,instance,user, password){
  url <- infuse("https://{{instance}}.cloud.databricks.com/api/1.2/contexts/destroy",instance=instance)
  ctxDestroyCall<-POST(url,body=list(clusterId=clusterId,contextId=ctx)
                     ,encode='form'
                     ,authenticate(user,password))
  ctxDestroy <- content(ctxDestroyCall)
}


## Command execute
## https://docs.databricks.com/api/1.2/index.html#command-execution
runCommand <- function(command,ctx,clusterId,instance,user,password){
  url <- infuse("https://{{instance}}.cloud.databricks.com/api/1.2/commands/execute"
              ,instance=instance)
  commandUrl<-POST(url
            ,body=list(language='python'
                       ,clusterId=clusterId
                       ,contextId=ctx
                       ,command=command)
            ,encode='form'
            ,authenticate(user,password))
  commandCtx <- content(commandUrl)$id
  commandCtx
}


## Command  Status
commandStatus <- function(cmdId,ctx,clusterId, instance,user,password){
  url <- infuse("https://{{instance}}.cloud.databricks.com/api/1.2/commands/status",instance=instance)
  cmdStatusCall<-GET(url,query=list(clusterId=clusterId
                                  ,contextId=ctx
                                  ,commandId=cmdId)
                   ,authenticate(user,password))
  (cmdStatus <- content(cmdStatusCall))
  return(cmdStatus)
}

## Command Cancel
commandCancel <- function(cmdId,ctx,clusterId, instance,user,password){
  url <- infuse("https://{{instance}}.cloud.databricks.com/api/1.2/commands/cancel"
              ,instance=instance)
  cancelUrl<-POST(url
                 ,body=list(clusterId=clusterId
                            ,contextId=ctx
                            ,commandId=cmdId)
                 ,encode='form'
                 ,authenticate(user,password))
  (cancelCtx <- content(cancelUrl))
  return(cancelCtx)
}


## Examples
instance <- "dbc-caf9527b-e073"
clusterId='1013-020424-beset44'
user <- "[email protected]"
password <- "!!"



ctx <- ctxMake(instance=instance,clusterId=clusterId,user=user,password=password)
ctxStatus(instance=instance,ctx=ctx,clusterId=clusterId, user=user,password=password)
cmdId <- runCommand(command="ms",ctx=ctx,clusterId=clusterId,instance=instance, user=user
                    ,password=password)
commandStatus(cmdId,ctx=ctx, clusterId=clusterId, instance=instance, user=user,password=password)
commandCancel(cmdId,ctx=ctx,clusterId=clusterId,instance=instance, user=user,password=password)



cmd='
import sys
import datetime
import random
import subprocess

from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
ms = spark.read.option("mergeSchema", "true").parquet("s3://telemetry-parquet/main_summary/v4/")
ms.createOrReplaceTempView("ms")
'

id1 <- runCommand(command=cmd,ctx=ctx,clusterId=clusterId,instance=instance, user=user
                    ,password=password)
commandStatus(id1,ctx=ctx, clusterId=clusterId, instance=instance, user=user,password=password)
commandCancel(id1,ctx=ctx,clusterId=clusterId,instance=instance, user=user,password=password)



cmd2 = 'r=spark.sql("select active_ticks*5 as x,subsession_length as y from ms where sample_id=\'42\' and submission_date_s3=\'20170801\' limit 1000").toPandas()'
id2 <- runCommand(command=cmd2,ctx=ctx,clusterId=clusterId,instance=instance, user=user
                  ,password=password)
commandStatus(id2,ctx=ctx, clusterId=clusterId, instance=instance, user=user,password=password)
commandCancel(id2,ctx=ctx,clusterId=clusterId,instance=instance, user=user,password=password)

cmd3 = 'print(r)'
id3 <- runCommand(command=cmd3,ctx=ctx,clusterId=clusterId,instance=instance, user=user
                  ,password=password)
commandStatus(id3,ctx=ctx, clusterId=clusterId, instance=instance, user=user,password=password)
commandCancel(id3,ctx=ctx,clusterId=clusterId,instance=instance, user=user,password=password)


cmd4='
from IPython.display import display
import matplotlib
import matplotlib.pyplot as plt
r.sort_values(by=["x"])
fig=r.plot(x="x",y="y",marker=".")
'
id4 <- runCommand(command=cmd4,ctx=ctx,clusterId=clusterId,instance=instance, user=user
                  ,password=password)
commandStatus(id4,ctx=ctx, clusterId=clusterId, instance=instance, user=user,password=password)
commandCancel(id3,ctx=ctx,clusterId=clusterId,instance=instance, user=user,password=password)



cmd5='display(fig)'
id5 <- runCommand(command=cmd5,ctx=ctx,clusterId=clusterId,instance=instance, user=user
                  ,password=password)
commandStatus(id5,ctx=ctx, clusterId=clusterId, instance=instance, user=user,password=password)
commandCancel(id5,ctx=ctx,clusterId=clusterId,instance=instance, user=user,password=password)