Executing SQL Queries on Databricks: A Guide to Using the SQL Execution API with Python to build Data Applications

Databricks is a powerful platform for big data processing and analytics. One of its key features is the ability to execute SQL queries on large datasets from Databricks SQL Warehouses. In this post, i’ll show you how to use the Databricks SQL Execution API and Python to execute SQL queries on your Databricks serverless Warehouse and retrieve the results.

What are key features of Databricks SQL Execution API?

The Databricks SQL Statement Execution API makes it easier for businesses to access their data and build custom data applications. It removes the need to manage connections and install drivers, making it more efficient and user-friendly.

How to use the Databricks SQL Execution API with Python?

Here’s an example of how you can use the Databricks SQL Execution API with Python:

#!/usr/bin/env python3

import json
import os
from urllib.parse import urljoin, urlencode

import pyarrow
import requests

# NOTE set debuglevel = 1 (or higher) for http debug logging
from http.client import HTTPConnection
HTTPConnection.debuglevel = 0


# Expects environment variables;

HOST = os.getenv("HOST")
URL = os.getenv("URL")
WAREHOUSE = os.getenv("WAREHOUSE")
AUTH_TOKEN = os.environ.get("AUTH_TOKEN")

if HOST and not URL:
    URL = f"https://{HOST}/api/2.0/sql/statements/"

assert URL, "Required: URL"
assert WAREHOUSE, "Required: WAREHOUSE"
assert AUTH_TOKEN, "Required: AUTH_TOKEN"


# example statement big enough to exceed 5MB data requirement
limit = 1000000
sql_statement = "SELECT l_orderkey, l_extendedprice, l_shipdate FROM samples.tpch.lineitem ".format(limit)
if limit:
    sql_statement = f"{sql_statement} LIMIT {limit}"

# check api reference for all options available https://docs.databricks.com/api-explorer/workspace/statementexecution
payload = json.dumps({
    "statement": sql_statement,
    "warehouse_id": WAREHOUSE,
    "catalog": "samples",
    "schema": "tpch",
    "wait_timeout": "10s",
    "disposition": "EXTERNAL_LINKS",
    "format": "ARROW_STREAM"
})

headers = {
    'Content-Type': 'application/json'
}
auth=('token', AUTH_TOKEN)

def process_success(response, limit=None):
    chunks = response.json()["manifest"]["chunks"]
    tables = []

    print("{} chunks(s) in result set".format(len(chunks)))
    for idx, chunkInfo in enumerate(chunks):

        if len(chunks)>1:
            stmt_url = urljoin(URL, statement_id) + "/"
            row_offset_param = urlencode({'row_offset': chunkInfo["row_offset"]})
            print(stmt_url)
            resolve_external_link_url = urljoin(stmt_url, "result/chunks/{}?{}".format(
                chunkInfo["chunk_index"], row_offset_param))
            
            print("resolve_external_link_url is "  , resolve_external_link_url)

            response = requests.get(resolve_external_link_url, auth=auth, headers=headers)
            print("response is " , response)
            assert response.status_code == 200

            external_url = urljoin(stmt_url,response.json()["external_links"][0]["external_link"])
            print("external_url is " , external_url)

        else:
            external_url = response.json()["result"]["external_links"][0]["external_link"]
        
        # NOTE: do _NOT_ send the authorization header to external urls
        raw_response = requests.get(external_url, auth=None, headers=None)
        print("raw_response " , raw_response)
        assert raw_response.status_code == 200

        arrow_table = pyarrow.ipc.open_stream(raw_response.content).read_all()
        
        tables.append(arrow_table)
        print("chunk {} received".format(idx))

    full_table = pyarrow.concat_tables(tables).to_pandas()

    # Export DataFrame to CSV file
    full_table.to_csv(f'{statement_id}_query_results.csv', sep=';', header='true' , encoding='utf-8', index=False)
    print(full_table)

print("Using URL:", URL)

response = requests.post(URL, auth=auth, headers=headers, data=payload)
print("Statement POST got HTTP status code:", response.status_code)
assert response.status_code == 200
state = response.json()["status"]["state"]
statement_id = response.json()["statement_id"]

while state in ["PENDING", "RUNNING"]:
    stmt_url = urljoin(URL, statement_id)
    response = requests.get(stmt_url, auth=auth, headers=headers)
    print("Statement GET got HTTP status code:", response.status_code)
    assert response.status_code == 200
    state = response.json()["status"]["state"]

assert state == "SUCCEEDED"

process_success(response, limit)

That’s it! I hope this post has helped you understand how to use the Databricks SQL Execution API with Python to execute SQL queries on your Databricks serverless Warehouse. If you have any questions, feel free to ask!

Some of the useful references on this topic :

https://www.databricks.com/blog/2023/03/07/databricks-sql-statement-execution-api-announcing-public-preview.html

https://docs.databricks.com/api-explorer/workspace/statementexecution