query_cancel.py

from databricks import sql
import os, threading, time

"""
The current operation of a cursor may be cancelled by calling its `.cancel()` method as shown in the example below.
"""

with sql.connect(
    server_hostname=os.getenv("DATABRICKS_SERVER_HOSTNAME"),
    http_path=os.getenv("DATABRICKS_HTTP_PATH"),
    access_token=os.getenv("DATABRICKS_TOKEN"),
) as connection:

    with connection.cursor() as cursor:

        def execute_really_long_query():
            try:
                cursor.execute(
                    "SELECT SUM(A.id - B.id) "
                    + "FROM range(1000000000) A CROSS JOIN range(100000000) B "
                    + "GROUP BY (A.id - B.id)"
                )
            except sql.exc.RequestError:
                print("It looks like this query was cancelled.")

        exec_thread = threading.Thread(target=execute_really_long_query)

        print("\n Beginning to execute long query")
        exec_thread.start()

        # Make sure the query has started before cancelling
        print("\n Waiting 15 seconds before canceling", end="", flush=True)

        seconds_waited = 0
        while seconds_waited < 15:
            seconds_waited += 1
            print(".", end="", flush=True)
            time.sleep(1)

        print("\n Cancelling the cursor's operation. This can take a few seconds.")
        cursor.cancel()

        print("\n Now checking the cursor status:")
        exec_thread.join(5)

        assert not exec_thread.is_alive()
        print("\n The previous command was successfully canceled")

        print("\n Now reusing the cursor to run a separate query.")

        # We can still execute a new command on the cursor
        cursor.execute("SELECT * FROM range(3)")

        print("\n Execution was successful. Results appear below:")

        print(cursor.fetchall())