From beffa2fb9cc3a86238105e40d6e2567f4aa07053 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Mon, 4 Nov 2024 12:02:29 +0530 Subject: [PATCH] Added integration tests for execute_async --- tests/e2e/test_driver.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index cfd1e969..0d6d3d29 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -36,6 +36,7 @@ compare_dbr_versions, is_thrift_v5_plus, ) +from databricks.sql.thrift_api.TCLIService import ttypes from tests.e2e.common.core_tests import CoreTestMixin, SmokeTestMixin from tests.e2e.common.large_queries_mixin import LargeQueriesMixin from tests.e2e.common.timestamp_tests import TimestampTestsMixin @@ -175,6 +176,26 @@ def test_cloud_fetch(self): for i in range(len(cf_result)): assert cf_result[i] == noop_result[i] + def test_execute_async(self): + def isExecuting(operation_state): + return not operation_state or operation_state in [ + ttypes.TOperationState.RUNNING_STATE, + ttypes.TOperationState.PENDING_STATE, + ] + + 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 + while isExecuting(cursor.get_query_state()): + log.info("Polling the status in test_execute_async") + + cursor.get_execution_result() + result = cursor.fetchall() + + assert result[0].asDict() == {"count(1)": 0} + # Exclude Retry tests because they require specific setups, and LargeQueries too slow for core # tests