handler.py

import boto3
import logging

from http.client import HTTPSConnection
from json import dumps, loads
from os import environ

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

kms_client = boto3.client('kms')
    
DATABRICKS_HOST=environ["DATABRICKS_HOST"]
DATABRICKS_TOKEN=environ["DATABRICKS_TOKEN"]
WAREHOUSE_IDS = environ["WAREHOUSE_IDS"].split(",")

def getRunningQueries(warehouseId: str) -> dict:
    conn = HTTPSConnection(DATABRICKS_HOST)
    url = f"https://{DATABRICKS_HOST}/api/2.0/sql/history/queries"
    headers = {
        "Authorization": f"Bearer {DATABRICKS_TOKEN}"
    }

    conn.request(
        method="GET",
        url=url,
        headers=headers,
        body=dumps(
            {
                "filter_by": {
                    "warehouse_ids": [warehouseId],
                    "statuses": ["RUNNING", "QUEUED"]
                },
                "max_results": 1000,
                "include_metrics": True
            }
        )
    )
    response = loads(conn.getresponse().read())
    logger.info(response)
    conn.close()
    return response

def stopWarehouse(warehouseId: str) -> int:
    conn = HTTPSConnection(DATABRICKS_HOST)
    url = f"https://{DATABRICKS_HOST}//api/2.0/sql/warehouses/{warehouseId}/stop"
    headers = {
        "Authorization": f"Bearer {DATABRICKS_TOKEN}"
    }
    conn.request(
        method="POST",
        url=url,
        headers=headers
    )

    response = loads(conn.getresponse().read())
    logger.info(response)
    conn.close()
    return response

def run(event, context):
    for warehouseId in WAREHOUSE_IDS:
        runningQueries = getRunningQueries(warehouseId)
        if len(runningQueries.get("res", [])) == 0:
            stopWarehouse(warehouseId)