from prefect import task, Flow
from prefect.tasks.secrets import PrefectSecret
from prefect.tasks.databricks import DatabricksRunNow
import requests


conncfg = {"host":"<https://xxx.cloud.databricks.com/>", "token":"xxx"}


job_json = {
      "job_id": 24242,
      "notebook_params": {}
    }


@task
def run_databricks(job_config):
    notebook_run = DatabricksRunNow(json=job_config).run()
    run_id = notebook_run(databricks_conn_secret=conncfg)
    headers = { "Authorization" : "Bearer xxx" }
    r = requests.get('<https://xxx.cloud.databricks.com/api/2.0/jobs/runs/get-output?run_id=%s'%(runid)>, headers=headers)
    results = r.json()
    output =  results['metadata']['notebook_output']
    if output.has_key('message'):
        raise Exception(output['message'])



with Flow('testflow') as flow:
    run_databricks(job_json)

DatabricksRunNow().run(json=job_config)

.run()

DatabricksRunNow(json=job_config).run(databricks_conn_secret=conncfg)