Forward porting all changes into 4.x.x. uptil v3.7.3 (#529)

* Base changes

* Black formatter

* Cache version fix

* Added the changed test_retry.py file

* retry_test_mixins changes

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

Author: jprakash-db

.github/workflows/code-quality-checks.yml

@@ -93,7 +93,7 @@ jobs:
       #----------------------------------------------
       -   name: Load cached venv
           id: cached-poetry-dependencies
-          uses: actions/cache@v2
+          uses: actions/cache@v4
           with:
             path: .venv-pyarrow
             key: venv-pyarrow-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ github.event.repository.name }}-${{ hashFiles('**/poetry.lock') }}

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

@@ -408,8 +408,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/backend/thrift_backend.py

@@ -68,7 +68,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),

src/databricks/sql/client.py

@@ -789,6 +789,14 @@ def execute(
             async_op=False,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
         )
+        self.active_result_set = ResultSet(
+            self.connection,
+            execute_response,
+            self.thrift_backend,
+            self.buffer_size_bytes,
+            self.arraysize,
+            self.connection.use_cloud_fetch,
+        )
 
         if self.active_result_set.is_staging_operation:
             self._handle_staging_operation(

tests/e2e/test_driver.py

@@ -179,14 +179,20 @@ 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 cursor.is_query_pending():
+            while self.isExecuting(cursor.get_query_state()):
                 time.sleep(self.POLLING_INTERVAL)
                 log.info("Polling the status in test_execute_async")
 
@@ -205,7 +211,7 @@ def test_execute_async__small_result(self):
             time.sleep(5)
 
             ## Polling after every POLLING_INTERVAL seconds
-            while cursor.is_query_pending():
+            while self.isExecuting(cursor.get_query_state()):
                 time.sleep(self.POLLING_INTERVAL)
                 log.info("Polling the status in test_execute_async")
 
@@ -235,7 +241,7 @@ def test_execute_async__large_result(self):
             time.sleep(5)
 
             ## Polling after every POLLING_INTERVAL seconds
-            while cursor.is_query_pending():
+            while self.isExecuting(cursor.get_query_state()):
                 time.sleep(self.POLLING_INTERVAL)
                 log.info("Polling the status in test_execute_async")