Added example for async execute query (#537)

Added examples and fixed the async execute not working without pyarrow

Author: jprakash-db

examples/query_async_execute.py

@@ -0,0 +1,32 @@
+from databricks import sql
+import os
+import time
+
+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:
+        long_running_query = """
+            SELECT COUNT(*) FROM RANGE(10000 * 16) x 
+            JOIN RANGE(10000) y 
+            ON FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') LIKE '%not%a%date%'
+        """
+
+        # Non-blocking call
+        cursor.execute_async(long_running_query)
+
+        # Polling every 5 seconds until the query is no longer pending
+        while cursor.is_query_pending():
+            print("POLLING")
+            time.sleep(5)
+
+        # Blocking call: fetch results when execution completes
+        cursor.get_async_execution_result()
+
+        result = cursor.fetchall()
+
+        for res in result:
+            print(res)

src/databricks/sql/client.py

@@ -896,6 +896,19 @@ def get_query_state(self) -> "TOperationState":
         self._check_not_closed()
         return self.thrift_backend.get_query_state(self.active_op_handle)
 
+    def is_query_pending(self):
+        """
+        Checks whether the async executing query is in pending state or not
+
+        :return:
+        """
+        operation_state = self.get_query_state()
+
+        return not operation_state or operation_state in [
+            ttypes.TOperationState.RUNNING_STATE,
+            ttypes.TOperationState.PENDING_STATE,
+        ]
+
     def get_async_execution_result(self):
         """
 
@@ -905,13 +918,7 @@ def get_async_execution_result(self):
         """
         self._check_not_closed()
 
-        def is_executing(operation_state) -> "bool":
-            return not operation_state or operation_state in [
-                ttypes.TOperationState.RUNNING_STATE,
-                ttypes.TOperationState.PENDING_STATE,
-            ]
-
-        while is_executing(self.get_query_state()):
+        while self.is_query_pending():
             # Poll after some default time
             time.sleep(self.ASYNC_DEFAULT_POLLING_INTERVAL)
 

src/databricks/sql/thrift_backend.py

@@ -797,12 +797,15 @@ def get_execution_result(self, op_handle, cursor):
             t_result_set_metadata_resp.schema
         )
 
-        schema_bytes = (
-            t_result_set_metadata_resp.arrowSchema
-            or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
-            .serialize()
-            .to_pybytes()
-        )
+        if pyarrow:
+            schema_bytes = (
+                t_result_set_metadata_resp.arrowSchema
+                or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
+                .serialize()
+                .to_pybytes()
+            )
+        else:
+            schema_bytes = None
 
         queue = ResultSetQueueFactory.build_queue(
             row_set_type=resp.resultSetMetadata.resultFormat,

tests/e2e/test_driver.py

@@ -179,20 +179,14 @@ def test_cloud_fetch(self):
 
 
 class TestPySQLAsyncQueriesSuite(PySQLPytestTestCase):
-    def isExecuting(self, operation_state):
-        return not operation_state or operation_state in [
-            ttypes.TOperationState.RUNNING_STATE,
-            ttypes.TOperationState.PENDING_STATE,
-        ]
-
     def test_execute_async__long_running(self):
 
         long_running_query = "SELECT COUNT(*) FROM RANGE(10000 * 16) x JOIN RANGE(10000) y ON FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') LIKE '%not%a%date%'"
         with self.cursor() as cursor:
             cursor.execute_async(long_running_query)
 
             ## Polling after every POLLING_INTERVAL seconds
-            while self.isExecuting(cursor.get_query_state()):
+            while cursor.is_query_pending():
                 time.sleep(self.POLLING_INTERVAL)
                 log.info("Polling the status in test_execute_async")
 
@@ -211,7 +205,7 @@ def test_execute_async__small_result(self):
             time.sleep(5)
 
             ## Polling after every POLLING_INTERVAL seconds
-            while self.isExecuting(cursor.get_query_state()):
+            while cursor.is_query_pending():
                 time.sleep(self.POLLING_INTERVAL)
                 log.info("Polling the status in test_execute_async")
 
@@ -241,7 +235,7 @@ def test_execute_async__large_result(self):
             time.sleep(5)
 
             ## Polling after every POLLING_INTERVAL seconds
-            while self.isExecuting(cursor.get_query_state()):
+            while cursor.is_query_pending():
                 time.sleep(self.POLLING_INTERVAL)
                 log.info("Polling the status in test_execute_async")