add docs for functions

Author: jprakash-db

src/databricks/sql/client.py

@@ -46,7 +46,7 @@
 from databricks.sql.experimental.oauth_persistence import OAuthPersistence
 
 from databricks.sql.thrift_api.TCLIService.ttypes import (
-    TSparkParameter,
+    TSparkParameter, TOperationState,
 )
 
 
@@ -764,6 +764,11 @@ def execute(
         Both will result in the query equivalent to "SELECT * FROM table WHERE field = 'foo'
         being sent to the server
 
+        async_op:
+        Denotes whether the execute command will execute the request asynchronously or not
+        By default it is set to False, if set True the execution request will be submitted and the code
+        will be non-blocking. User can later poll and request the result when ready
+
         :returns self
         """
 
@@ -819,14 +824,35 @@ def execute_async(
         self,
         operation: str,
         parameters: Optional[TParameterCollection] = None,
-    ):
-        return self.execute(operation, parameters, True)
+    ) -> "Cursor":
+        """
+
+        Execute a query and do not wait for it to complete and just move ahead
+
+        Internally it calls execute function with async_op=True
+        :param operation:
+        :param parameters:
+        :return:
+        """
+        self.execute(operation, parameters, True)
+        return self
 
-    def get_query_state(self):
+    def get_query_state(self) -> "TOperationState":
+        """
+        Get the state of the async executing query or basically poll the status of the query
+
+        :return:
+        """
         self._check_not_closed()
         return self.thrift_backend.get_query_state(self.active_op_handle)
 
     def get_execution_result(self):
+        """
+
+        Checks for the status of the async executing query and fetches the result if the query is finished
+        If executed sets the active_result_set to the obtained result
+        :return:
+        """
         self._check_not_closed()
 
         operation_state = self.get_query_state()

src/databricks/sql/thrift_backend.py

@@ -7,6 +7,8 @@
 import threading
 from typing import List, Union
 
+from databricks.sql.thrift_api.TCLIService.ttypes import TOperationState
+
 try:
     import pyarrow
 except ImportError:
@@ -844,7 +846,7 @@ def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
             self._check_command_not_in_error_or_closed_state(op_handle, poll_resp)
         return operation_state
 
-    def get_query_state(self, op_handle):
+    def get_query_state(self, op_handle) -> "TOperationState":
         poll_resp = self._poll_for_status(op_handle)
         operation_state = poll_resp.operationState
         self._check_command_not_in_error_or_closed_state(op_handle, poll_resp)

tests/e2e/test_driver.py

@@ -187,8 +187,9 @@ def isExecuting(operation_state):
         with self.cursor() as cursor:
             cursor.execute_async(long_running_query)
 
-            ## Polling
+            ## Polling after every 10 seconds
             while isExecuting(cursor.get_query_state()):
+                time.sleep(10)
                 log.info("Polling the status in test_execute_async")
 
             cursor.get_execution_result()