Added example for async execute query (#537)

Added examples and fixed the async execute not working without pyarrow

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: jprakash-db

src/databricks/sql/client.py

@@ -877,6 +877,19 @@ def is_query_pending(self):
             ttypes.TOperationState.PENDING_STATE,
         ]
 
+    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):
         """
 

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")