Skip to content

Commit 156421a

Browse files
Merge branch 'metadata-sea' into fetch-json-inline
2 parents 242307a + a4d5bdb commit 156421a

File tree

14 files changed

+162
-325
lines changed

14 files changed

+162
-325
lines changed

CHANGELOG.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
# Release History
22

3+
# 4.0.5 (2025-06-24)
4+
- Fix: Reverted change in cursor close handling which led to errors impacting users (databricks/databricks-sql-python#613 by @madhav-db)
5+
36
# 4.0.4 (2025-06-16)
47

58
- Update thrift client library after cleaning up unused fields and structs (databricks/databricks-sql-python#553 by @vikrantpuppala)

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[tool.poetry]
22
name = "databricks-sql-connector"
3-
version = "4.0.4"
3+
version = "4.0.5"
44
description = "Databricks SQL Connector for Python"
55
authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
66
license = "Apache-2.0"

src/databricks/sql/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ def __repr__(self):
6868
DATE = DBAPITypeObject("date")
6969
ROWID = DBAPITypeObject()
7070

71-
__version__ = "4.0.4"
71+
__version__ = "4.0.5"
7272
USER_AGENT_NAME = "PyDatabricksSqlConnector"
7373

7474
# These two functions are pyhive legacy

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -719,8 +719,14 @@ def get_tables(
719719
)
720720
assert result is not None, "execute_command returned None in synchronous mode"
721721

722+
from databricks.sql.result_set import SeaResultSet
723+
724+
assert isinstance(
725+
result, SeaResultSet
726+
), "execute_command returned a non-SeaResultSet"
727+
722728
# Apply client-side filtering by table_types
723-
from databricks.sql.backend.filters import ResultSetFilter
729+
from databricks.sql.backend.sea.utils.filters import ResultSetFilter
724730

725731
result = ResultSetFilter.filter_tables_by_type(result, table_types)
726732

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,8 @@ class MetadataCommands(Enum):
6060
SHOW_TABLES_ALL_CATALOGS = "SHOW TABLES IN ALL CATALOGS"
6161
SHOW_COLUMNS = "SHOW COLUMNS IN CATALOG {}"
6262

63-
SCHEMA_LIKE_PATTERN = " SCHEMA LIKE '{}'"
64-
TABLE_LIKE_PATTERN = " TABLE LIKE '{}'"
6563
LIKE_PATTERN = " LIKE '{}'"
64+
SCHEMA_LIKE_PATTERN = " SCHEMA" + LIKE_PATTERN
65+
TABLE_LIKE_PATTERN = " TABLE" + LIKE_PATTERN
6666

6767
CATALOG_SPECIFIC = "CATALOG {}"

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

Lines changed: 23 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,22 @@
44
This module provides filtering capabilities for result sets returned by different backends.
55
"""
66

7+
from __future__ import annotations
8+
79
import logging
810
from typing import (
911
List,
1012
Optional,
1113
Any,
1214
Callable,
1315
cast,
16+
TYPE_CHECKING,
1417
)
1518

16-
from databricks.sql.backend.sea.backend import SeaDatabricksClient
17-
from databricks.sql.backend.types import ExecuteResponse
19+
if TYPE_CHECKING:
20+
from databricks.sql.result_set import SeaResultSet
1821

19-
from databricks.sql.result_set import ResultSet, SeaResultSet
22+
from databricks.sql.backend.types import ExecuteResponse
2023

2124
logger = logging.getLogger(__name__)
2225

@@ -62,11 +65,11 @@ def _filter_sea_result_set(
6265
)
6366

6467
# Create a new ResultData object with filtered data
65-
6668
from databricks.sql.backend.sea.models.base import ResultData
6769

6870
result_data = ResultData(data=filtered_rows, external_links=None)
6971

72+
from databricks.sql.backend.sea.backend import SeaDatabricksClient
7073
from databricks.sql.result_set import SeaResultSet
7174

7275
# Create a new SeaResultSet with the filtered data
@@ -83,11 +86,11 @@ def _filter_sea_result_set(
8386

8487
@staticmethod
8588
def filter_by_column_values(
86-
result_set: ResultSet,
89+
result_set: SeaResultSet,
8790
column_index: int,
8891
allowed_values: List[str],
8992
case_sensitive: bool = False,
90-
) -> ResultSet:
93+
) -> SeaResultSet:
9194
"""
9295
Filter a result set by values in a specific column.
9396
@@ -105,34 +108,24 @@ def filter_by_column_values(
105108
if not case_sensitive:
106109
allowed_values = [v.upper() for v in allowed_values]
107110

108-
# Determine the type of result set and apply appropriate filtering
109-
from databricks.sql.result_set import SeaResultSet
110-
111-
if isinstance(result_set, SeaResultSet):
112-
return ResultSetFilter._filter_sea_result_set(
113-
result_set,
114-
lambda row: (
115-
len(row) > column_index
116-
and isinstance(row[column_index], str)
117-
and (
118-
row[column_index].upper()
119-
if not case_sensitive
120-
else row[column_index]
121-
)
122-
in allowed_values
123-
),
124-
)
125-
126-
# For other result set types, return the original (should be handled by specific implementations)
127-
logger.warning(
128-
f"Filtering not implemented for result set type: {type(result_set).__name__}"
111+
return ResultSetFilter._filter_sea_result_set(
112+
result_set,
113+
lambda row: (
114+
len(row) > column_index
115+
and isinstance(row[column_index], str)
116+
and (
117+
row[column_index].upper()
118+
if not case_sensitive
119+
else row[column_index]
120+
)
121+
in allowed_values
122+
),
129123
)
130-
return result_set
131124

132125
@staticmethod
133126
def filter_tables_by_type(
134-
result_set: ResultSet, table_types: Optional[List[str]] = None
135-
) -> ResultSet:
127+
result_set: SeaResultSet, table_types: Optional[List[str]] = None
128+
) -> SeaResultSet:
136129
"""
137130
Filter a result set of tables by the specified table types.
138131

src/databricks/sql/client.py

Lines changed: 3 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -281,13 +281,7 @@ def __enter__(self) -> "Connection":
281281
return self
282282

283283
def __exit__(self, exc_type, exc_value, traceback):
284-
try:
285-
self.close()
286-
except BaseException as e:
287-
logger.warning(f"Exception during connection close in __exit__: {e}")
288-
if exc_type is None:
289-
raise
290-
return False
284+
self.close()
291285

292286
def __del__(self):
293287
if self.open:
@@ -417,14 +411,7 @@ def __enter__(self) -> "Cursor":
417411
return self
418412

419413
def __exit__(self, exc_type, exc_value, traceback):
420-
try:
421-
logger.debug("Cursor context manager exiting, calling close()")
422-
self.close()
423-
except BaseException as e:
424-
logger.warning(f"Exception during cursor close in __exit__: {e}")
425-
if exc_type is None:
426-
raise
427-
return False
414+
self.close()
428415

429416
def __iter__(self):
430417
if self.active_result_set:
@@ -1091,21 +1078,7 @@ def cancel(self) -> None:
10911078
def close(self) -> None:
10921079
"""Close cursor"""
10931080
self.open = False
1094-
1095-
# Close active operation handle if it exists
1096-
if self.active_command_id:
1097-
try:
1098-
self.backend.close_command(self.active_command_id)
1099-
except RequestError as e:
1100-
if isinstance(e.args[1], CursorAlreadyClosedError):
1101-
logger.info("Operation was canceled by a prior request")
1102-
else:
1103-
logging.warning(f"Error closing operation handle: {e}")
1104-
except Exception as e:
1105-
logging.warning(f"Error closing operation handle: {e}")
1106-
finally:
1107-
self.active_command_id = None
1108-
1081+
self.active_command_id = None
11091082
if self.active_result_set:
11101083
self._close_and_clear_active_result_set()
11111084

tests/e2e/common/large_queries_mixin.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ def test_query_with_large_narrow_result_set(self):
8383
assert row[0] == row_id
8484

8585
def test_long_running_query(self):
86-
"""Incrementally increase query size until it takes at least 5 minutes,
86+
"""Incrementally increase query size until it takes at least 3 minutes,
8787
and asserts that the query completes successfully.
8888
"""
8989
minutes = 60
@@ -113,5 +113,5 @@ def test_long_running_query(self):
113113
duration = time.time() - start
114114
current_fraction = duration / min_duration
115115
print("Took {} s with scale factor={}".format(duration, scale_factor))
116-
# Extrapolate linearly to reach 5 min and add 50% padding to push over the limit
116+
# Extrapolate linearly to reach 3 min and add 50% padding to push over the limit
117117
scale_factor = math.ceil(1.5 * scale_factor / current_fraction)

tests/e2e/common/staging_ingestion_tests.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
4646
) as conn:
4747

4848
cursor = conn.cursor()
49-
query = f"PUT '{temp_path}' INTO 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' OVERWRITE"
49+
query = f"PUT '{temp_path}' INTO 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' OVERWRITE"
5050
cursor.execute(query)
5151

5252
# GET should succeed
@@ -57,7 +57,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
5757
extra_params={"staging_allowed_local_path": new_temp_path}
5858
) as conn:
5959
cursor = conn.cursor()
60-
query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
60+
query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"
6161
cursor.execute(query)
6262

6363
with open(new_fh, "rb") as fp:
@@ -67,7 +67,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
6767

6868
# REMOVE should succeed
6969

70-
remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv'"
70+
remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv'"
7171

7272
with self.connection(extra_params={"staging_allowed_local_path": "/"}) as conn:
7373
cursor = conn.cursor()
@@ -79,7 +79,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
7979
Error, match="Staging operation over HTTP was unsuccessful: 404"
8080
):
8181
cursor = conn.cursor()
82-
query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
82+
query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"
8383
cursor.execute(query)
8484

8585
os.remove(temp_path)

0 commit comments

Comments
 (0)