Merge branch 'metadata-sea' into fetch-json-inline

Author: varun-edachali-dbx

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Release History
 
+# 4.0.5 (2025-06-24)
+- Fix: Reverted change in cursor close handling which led to errors impacting users (databricks/databricks-sql-python#613 by @madhav-db)
+
 # 4.0.4 (2025-06-16)
 
 - Update thrift client library after cleaning up unused fields and structs (databricks/databricks-sql-python#553 by @vikrantpuppala)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.0.4"
+version = "4.0.5"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"

src/databricks/sql/__init__.py

@@ -68,7 +68,7 @@ def __repr__(self):
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "4.0.4"
+__version__ = "4.0.5"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

src/databricks/sql/backend/sea/backend.py

@@ -719,8 +719,14 @@ def get_tables(
         )
         assert result is not None, "execute_command returned None in synchronous mode"
 
+        from databricks.sql.result_set import SeaResultSet
+
+        assert isinstance(
+            result, SeaResultSet
+        ), "execute_command returned a non-SeaResultSet"
+
         # Apply client-side filtering by table_types
-        from databricks.sql.backend.filters import ResultSetFilter
+        from databricks.sql.backend.sea.utils.filters import ResultSetFilter
 
         result = ResultSetFilter.filter_tables_by_type(result, table_types)
 

src/databricks/sql/backend/sea/utils/constants.py

@@ -60,8 +60,8 @@ class MetadataCommands(Enum):
     SHOW_TABLES_ALL_CATALOGS = "SHOW TABLES IN ALL CATALOGS"
     SHOW_COLUMNS = "SHOW COLUMNS IN CATALOG {}"
 
-    SCHEMA_LIKE_PATTERN = " SCHEMA LIKE '{}'"
-    TABLE_LIKE_PATTERN = " TABLE LIKE '{}'"
     LIKE_PATTERN = " LIKE '{}'"
+    SCHEMA_LIKE_PATTERN = " SCHEMA" + LIKE_PATTERN
+    TABLE_LIKE_PATTERN = " TABLE" + LIKE_PATTERN
 
     CATALOG_SPECIFIC = "CATALOG {}"

src/databricks/sql/backend/filters.py renamed to src/databricks/sql/backend/sea/utils/filters.py

@@ -4,19 +4,22 @@
 This module provides filtering capabilities for result sets returned by different backends.
 """
 
+from __future__ import annotations
+
 import logging
 from typing import (
     List,
     Optional,
     Any,
     Callable,
     cast,
+    TYPE_CHECKING,
 )
 
-from databricks.sql.backend.sea.backend import SeaDatabricksClient
-from databricks.sql.backend.types import ExecuteResponse
+if TYPE_CHECKING:
+    from databricks.sql.result_set import SeaResultSet
 
-from databricks.sql.result_set import ResultSet, SeaResultSet
+from databricks.sql.backend.types import ExecuteResponse
 
 logger = logging.getLogger(__name__)
 
@@ -62,11 +65,11 @@ def _filter_sea_result_set(
         )
 
         # Create a new ResultData object with filtered data
-
         from databricks.sql.backend.sea.models.base import ResultData
 
         result_data = ResultData(data=filtered_rows, external_links=None)
 
+        from databricks.sql.backend.sea.backend import SeaDatabricksClient
         from databricks.sql.result_set import SeaResultSet
 
         # Create a new SeaResultSet with the filtered data
@@ -83,11 +86,11 @@ def _filter_sea_result_set(
 
     @staticmethod
     def filter_by_column_values(
-        result_set: ResultSet,
+        result_set: SeaResultSet,
         column_index: int,
         allowed_values: List[str],
         case_sensitive: bool = False,
-    ) -> ResultSet:
+    ) -> SeaResultSet:
         """
         Filter a result set by values in a specific column.
 
@@ -105,34 +108,24 @@ def filter_by_column_values(
         if not case_sensitive:
             allowed_values = [v.upper() for v in allowed_values]
 
-        # Determine the type of result set and apply appropriate filtering
-        from databricks.sql.result_set import SeaResultSet
-
-        if isinstance(result_set, SeaResultSet):
-            return ResultSetFilter._filter_sea_result_set(
-                result_set,
-                lambda row: (
-                    len(row) > column_index
-                    and isinstance(row[column_index], str)
-                    and (
-                        row[column_index].upper()
-                        if not case_sensitive
-                        else row[column_index]
-                    )
-                    in allowed_values
-                ),
-            )
-
-        # For other result set types, return the original (should be handled by specific implementations)
-        logger.warning(
-            f"Filtering not implemented for result set type: {type(result_set).__name__}"
+        return ResultSetFilter._filter_sea_result_set(
+            result_set,
+            lambda row: (
+                len(row) > column_index
+                and isinstance(row[column_index], str)
+                and (
+                    row[column_index].upper()
+                    if not case_sensitive
+                    else row[column_index]
+                )
+                in allowed_values
+            ),
         )
-        return result_set
 
     @staticmethod
     def filter_tables_by_type(
-        result_set: ResultSet, table_types: Optional[List[str]] = None
-    ) -> ResultSet:
+        result_set: SeaResultSet, table_types: Optional[List[str]] = None
+    ) -> SeaResultSet:
         """
         Filter a result set of tables by the specified table types.
 

src/databricks/sql/client.py

@@ -281,13 +281,7 @@ def __enter__(self) -> "Connection":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        try:
-            self.close()
-        except BaseException as e:
-            logger.warning(f"Exception during connection close in __exit__: {e}")
-            if exc_type is None:
-                raise
-        return False
+        self.close()
 
     def __del__(self):
         if self.open:
@@ -417,14 +411,7 @@ def __enter__(self) -> "Cursor":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        try:
-            logger.debug("Cursor context manager exiting, calling close()")
-            self.close()
-        except BaseException as e:
-            logger.warning(f"Exception during cursor close in __exit__: {e}")
-            if exc_type is None:
-                raise
-        return False
+        self.close()
 
     def __iter__(self):
         if self.active_result_set:
@@ -1091,21 +1078,7 @@ def cancel(self) -> None:
     def close(self) -> None:
         """Close cursor"""
         self.open = False
-
-        # Close active operation handle if it exists
-        if self.active_command_id:
-            try:
-                self.backend.close_command(self.active_command_id)
-            except RequestError as e:
-                if isinstance(e.args[1], CursorAlreadyClosedError):
-                    logger.info("Operation was canceled by a prior request")
-                else:
-                    logging.warning(f"Error closing operation handle: {e}")
-            except Exception as e:
-                logging.warning(f"Error closing operation handle: {e}")
-            finally:
-                self.active_command_id = None
-
+        self.active_command_id = None
         if self.active_result_set:
             self._close_and_clear_active_result_set()
 

tests/e2e/common/large_queries_mixin.py

@@ -83,7 +83,7 @@ def test_query_with_large_narrow_result_set(self):
                 assert row[0] == row_id
 
     def test_long_running_query(self):
-        """Incrementally increase query size until it takes at least 5 minutes,
+        """Incrementally increase query size until it takes at least 3 minutes,
         and asserts that the query completes successfully.
         """
         minutes = 60
@@ -113,5 +113,5 @@ def test_long_running_query(self):
                 duration = time.time() - start
                 current_fraction = duration / min_duration
                 print("Took {} s with scale factor={}".format(duration, scale_factor))
-                # Extrapolate linearly to reach 5 min and add 50% padding to push over the limit
+                # Extrapolate linearly to reach 3 min and add 50% padding to push over the limit
                 scale_factor = math.ceil(1.5 * scale_factor / current_fraction)

tests/e2e/common/staging_ingestion_tests.py

@@ -46,7 +46,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
         ) as conn:
 
             cursor = conn.cursor()
-            query = f"PUT '{temp_path}' INTO 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' OVERWRITE"
+            query = f"PUT '{temp_path}' INTO 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' OVERWRITE"
             cursor.execute(query)
 
         # GET should succeed
@@ -57,7 +57,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
             extra_params={"staging_allowed_local_path": new_temp_path}
         ) as conn:
             cursor = conn.cursor()
-            query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
+            query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"
             cursor.execute(query)
 
         with open(new_fh, "rb") as fp:
@@ -67,7 +67,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
 
         # REMOVE should succeed
 
-        remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv'"
+        remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv'"
 
         with self.connection(extra_params={"staging_allowed_local_path": "/"}) as conn:
             cursor = conn.cursor()
@@ -79,7 +79,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
                 Error, match="Staging operation over HTTP was unsuccessful: 404"
             ):
                 cursor = conn.cursor()
-                query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
+                query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"
                 cursor.execute(query)
 
         os.remove(temp_path)