Base changes

Author: jprakash-db

CHANGELOG.md

@@ -5,6 +5,16 @@
 - Split the connector into two separate packages: `databricks-sql-connector` and `databricks-sqlalchemy`. The `databricks-sql-connector` package contains the core functionality of the connector, while the `databricks-sqlalchemy` package contains the SQLAlchemy dialect for the connector. 
 - Pyarrow dependency is now optional in `databricks-sql-connector`. Users needing arrow are supposed to explicitly install pyarrow
 
+# 3.7.3 (2025-03-28)
+
+- Fix: Unable to poll small results in execute_async function (databricks/databricks-sql-python#515 by @jprakash-db) 
+- Updated log messages to show the status code and error messages of requests (databricks/databricks-sql-python#511 by @jprakash-db) 
+- Fix: Incorrect metadata was fetched in case of queries with the same alias (databricks/databricks-sql-python#505 by @jprakash-db)
+
+# 3.7.2 (2025-01-31)
+
+- Updated the retry_dela_max and retry_timeout (databricks/databricks-sql-python#497 by @jprakash-db)
+
 # 3.7.1 (2025-01-07)
 
 - Relaxed the number of Http retry attempts (databricks/databricks-sql-python#486 by @jprakash-db)

src/databricks/sql/auth/retry.py

@@ -290,8 +290,9 @@ def sleep_for_retry(self, response: BaseHTTPResponse) -> bool:
         else:
             proposed_wait = self.get_backoff_time()
 
-        proposed_wait = min(proposed_wait, self.delay_max)
+        proposed_wait = max(proposed_wait, self.delay_max)
         self.check_proposed_wait(proposed_wait)
+        logger.debug(f"Retrying after {proposed_wait} seconds")
         time.sleep(proposed_wait)
         return True
 
@@ -344,23 +345,24 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if a retry would violate the configured policy.
         """
 
+        logger.info(f"Received status code {status_code} for {method} request")
+
         # Request succeeded. Don't retry.
         if status_code == 200:
             return False, "200 codes are not retried"
 
         if status_code == 401:
-            raise NonRecoverableNetworkError(
-                "Received 401 - UNAUTHORIZED. Confirm your authentication credentials."
+            return (
+                False,
+                "Received 401 - UNAUTHORIZED. Confirm your authentication credentials.",
             )
 
         if status_code == 403:
-            raise NonRecoverableNetworkError(
-                "Received 403 - FORBIDDEN. Confirm your authentication credentials."
-            )
+            return False, "403 codes are not retried"
 
         # Request failed and server said NotImplemented. This isn't recoverable. Don't retry.
         if status_code == 501:
-            raise NonRecoverableNetworkError("Received code 501 from server.")
+            return False, "Received code 501 from server."
 
         # Request failed and this method is not retryable. We only retry POST requests.
         if not self._is_method_retryable(method):
@@ -399,8 +401,9 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
             and status_code not in self.status_forcelist
             and status_code not in self.force_dangerous_codes
         ):
-            raise UnsafeToRetryError(
-                "ExecuteStatement command can only be retried for codes 429 and 503. Received code: {status_code}"
+            return (
+                False,
+                "ExecuteStatement command can only be retried for codes 429 and 503",
             )
 
         # Request failed with a dangerous code, was an ExecuteStatement, but user forced retries for this

src/databricks/sql/auth/thrift_http_client.py

@@ -198,6 +198,12 @@ def flush(self):
         self.message = self.__resp.reason
         self.headers = self.__resp.headers
 
+        logger.info(
+            "HTTP Response with status code {}, message: {}".format(
+                self.code, self.message
+            )
+        )
+
     @staticmethod
     def basic_proxy_auth_headers(proxy):
         if proxy is None or not proxy.username:

src/databricks/sql/client.py

@@ -744,6 +744,7 @@ def execute(
         self,
         operation: str,
         parameters: Optional[TParameterCollection] = None,
+        enforce_embedded_schema_correctness=False
     ) -> "Cursor":
         """
         Execute a query and wait for execution to complete.
@@ -808,6 +809,7 @@ def execute(
             use_cloud_fetch=self.connection.use_cloud_fetch,
             parameters=prepared_params,
             async_op=False,
+            enforce_embedded_schema_correctness=enforce_embedded_schema_correctness
         )
         self.active_result_set = ResultSet(
             self.connection,
@@ -829,6 +831,7 @@ def execute_async(
         self,
         operation: str,
         parameters: Optional[TParameterCollection] = None,
+        enforce_embedded_schema_correctness=False
     ) -> "Cursor":
         """
 
@@ -869,6 +872,7 @@ def execute_async(
             use_cloud_fetch=self.connection.use_cloud_fetch,
             parameters=prepared_params,
             async_op=True,
+            enforce_embedded_schema_correctness=enforce_embedded_schema_correctness
         )
 
         return self

src/databricks/sql/thrift_backend.py

@@ -66,7 +66,7 @@
 # - 900s attempts-duration lines up w ODBC/JDBC drivers (for cluster startup > 10 mins)
 _retry_policy = {  # (type, default, min, max)
     "_retry_delay_min": (float, 1, 0.1, 60),
-    "_retry_delay_max": (float, 30, 5, 3600),
+    "_retry_delay_max": (float, 60, 5, 3600),
     "_retry_stop_after_attempts_count": (int, 30, 1, 60),
     "_retry_stop_after_attempts_duration": (float, 900, 1, 86400),
     "_retry_delay_default": (float, 5, 1, 60),
@@ -883,6 +883,7 @@ def execute_command(
         use_cloud_fetch=True,
         parameters=[],
         async_op=False,
+        enforce_embedded_schema_correctness=False
     ):
         assert session_handle is not None
 
@@ -898,8 +899,12 @@ def execute_command(
             sessionHandle=session_handle,
             statement=operation,
             runAsync=True,
-            getDirectResults=ttypes.TSparkGetDirectResults(
-                maxRows=max_rows, maxBytes=max_bytes
+            # For async operation we don't want the direct results
+            getDirectResults=None
+            if async_op
+            else ttypes.TSparkGetDirectResults(
+                maxRows=max_rows,
+                maxBytes=max_bytes,
             ),
             canReadArrowResult=True if pyarrow else False,
             canDecompressLZ4Result=lz4_compression,
@@ -910,6 +915,7 @@ def execute_command(
             },
             useArrowNativeTypes=spark_arrow_types,
             parameters=parameters,
+            enforceEmbeddedSchemaCorrectness=enforce_embedded_schema_correctness
         )
         resp = self.make_request(self._client.ExecuteStatement, req)
 

tests/e2e/common/large_queries_mixin.py

@@ -94,7 +94,7 @@ def test_long_running_query(self):
         scale_factor = 1
         with self.cursor() as cursor:
             while duration < min_duration:
-                assert scale_factor < 512, "Detected infinite loop"
+                assert scale_factor < 1024, "Detected infinite loop"
                 start = time.time()
 
                 cursor.execute(

tests/e2e/test_driver.py

@@ -177,19 +177,22 @@ 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,
-            ]
+
+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 isExecuting(cursor.get_query_state()):
+            while self.isExecuting(cursor.get_query_state()):
                 time.sleep(self.POLLING_INTERVAL)
                 log.info("Polling the status in test_execute_async")
 
@@ -198,6 +201,55 @@ def isExecuting(operation_state):
 
             assert result[0].asDict() == {"count(1)": 0}
 
+    def test_execute_async__small_result(self):
+        small_result_query = "SELECT 1"
+
+        with self.cursor() as cursor:
+            cursor.execute_async(small_result_query)
+
+            ## Fake sleep for 5 secs
+            time.sleep(5)
+
+            ## Polling after every POLLING_INTERVAL seconds
+            while self.isExecuting(cursor.get_query_state()):
+                time.sleep(self.POLLING_INTERVAL)
+                log.info("Polling the status in test_execute_async")
+
+            cursor.get_async_execution_result()
+            result = cursor.fetchall()
+
+            assert result[0].asDict() == {"1": 1}
+
+    def test_execute_async__large_result(self):
+        x_dimension = 1000
+        y_dimension = 1000
+        large_result_query = f"""
+                SELECT
+                    x.id AS x_id,
+                    y.id AS y_id,
+                    FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') AS date
+                FROM
+                    RANGE({x_dimension}) x
+                JOIN
+                    RANGE({y_dimension}) y
+            """
+
+        with self.cursor() as cursor:
+            cursor.execute_async(large_result_query)
+
+            ## Fake sleep for 5 secs
+            time.sleep(5)
+
+            ## Polling after every POLLING_INTERVAL seconds
+            while self.isExecuting(cursor.get_query_state()):
+                time.sleep(self.POLLING_INTERVAL)
+                log.info("Polling the status in test_execute_async")
+
+            cursor.get_async_execution_result()
+            result = cursor.fetchall()
+
+            assert len(result) == x_dimension * y_dimension
+
 
 # Exclude Retry tests because they require specific setups, and LargeQueries too slow for core
 # tests
@@ -828,4 +880,4 @@ def test_initial_namespace(self):
             cursor.execute("select current_catalog()")
             assert cursor.fetchone()[0] == self.arguments["catalog"]
             cursor.execute("select current_database()")
-            assert cursor.fetchone()[0] == table_name
+            assert cursor.fetchone()[0] == table_name