lakefs_hooks_server.py

import os
import time
import requests
from flask import Flask, request
app = Flask(__name__)


def run_job(domain, token, job_id):
    api_response = requests.post(
        f'{domain}/api/2.0/jobs/run-now',
        headers={'Authorization': f'Bearer {token}'},
        json={'job_id': job_id },
    )
    return api_response.json().get('run_id')
  
def poll_run_id(domain, token, run_id):
    while True:
        poll_response = requests.get(
            f'{domain}/api/2.0/jobs/runs/get',
            headers={'Authorization': f'Bearer {token}'},
            params={'run_id': run_id},
        )
        state = poll_response.json().get('state', {}).get('result_state')
        if not state:
            time.sleep(1)
            continue  # Job hasn't completed yet
        return state
      
@app.route('/hooks/run-databricks-job', methods = ['POST'])
def webhook_databricks_job():
    dbs_domain = os.environ.get('DATABRICKS_DOMAIN')
    dbs_token = os.environ.get('DATABRICKS_API_TOKEN')
    job_id = request.args.get('job_id')  # Read query parameter
    run_id = run_job(dbs_domain, dbs_token, job_id)
    state = poll_run_id(dbs_domain, dbs_token, run_id)
    status_code = 200 if state == 'SUCCESS' else 400
    return state, status_code