Skip to content

Commit 329c9af

Browse files
Merge branch 'ext-links-sea' into merge-http-client
2 parents 92a9260 + 82f9d6b commit 329c9af

16 files changed

+1080
-255
lines changed

src/databricks/sql/backend/databricks_client.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
from abc import ABC, abstractmethod
1212
from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING
1313

14+
from databricks.sql.types import SSLOptions
15+
1416
if TYPE_CHECKING:
1517
from databricks.sql.client import Cursor
1618

@@ -25,6 +27,13 @@
2527

2628

2729
class DatabricksClient(ABC):
30+
def __init__(self, ssl_options: SSLOptions, **kwargs):
31+
self._use_arrow_native_complex_types = kwargs.get(
32+
"_use_arrow_native_complex_types", True
33+
)
34+
self._max_download_threads = kwargs.get("max_download_threads", 10)
35+
self._ssl_options = ssl_options
36+
2837
# == Connection and Session Management ==
2938
@abstractmethod
3039
def open_session(

src/databricks/sql/backend/filters.py

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,8 @@
99
List,
1010
Optional,
1111
Any,
12-
Dict,
1312
Callable,
14-
TypeVar,
15-
Generic,
1613
cast,
17-
TYPE_CHECKING,
1814
)
1915

2016
from databricks.sql.backend.types import ExecuteResponse, CommandId
@@ -29,16 +25,13 @@
2925

3026
class ResultSetFilter:
3127
"""
32-
A general-purpose filter for result sets that can be applied to any backend.
33-
34-
This class provides methods to filter result sets based on various criteria,
35-
similar to the client-side filtering in the JDBC connector.
28+
A general-purpose filter for result sets.
3629
"""
3730

3831
@staticmethod
3932
def _filter_sea_result_set(
40-
result_set: "SeaResultSet", filter_func: Callable[[List[Any]], bool]
41-
) -> "SeaResultSet":
33+
result_set: SeaResultSet, filter_func: Callable[[List[Any]], bool]
34+
) -> SeaResultSet:
4235
"""
4336
Filter a SEA result set using the provided filter function.
4437
@@ -49,15 +42,13 @@ def _filter_sea_result_set(
4942
Returns:
5043
A filtered SEA result set
5144
"""
45+
5246
# Get all remaining rows
5347
all_rows = result_set.results.remaining_rows()
5448

5549
# Filter rows
5650
filtered_rows = [row for row in all_rows if filter_func(row)]
5751

58-
# Import SeaResultSet here to avoid circular imports
59-
from databricks.sql.result_set import SeaResultSet
60-
6152
# Reuse the command_id from the original result set
6253
command_id = result_set.command_id
6354

@@ -73,10 +64,13 @@ def _filter_sea_result_set(
7364
)
7465

7566
# Create a new ResultData object with filtered data
67+
7668
from databricks.sql.backend.sea.models.base import ResultData
7769

7870
result_data = ResultData(data=filtered_rows, external_links=None)
7971

72+
from databricks.sql.result_set import SeaResultSet
73+
8074
# Create a new SeaResultSet with the filtered data
8175
from databricks.sql.backend.sea.backend import SeaDatabricksClient
8276

@@ -93,11 +87,11 @@ def _filter_sea_result_set(
9387

9488
@staticmethod
9589
def filter_by_column_values(
96-
result_set: "ResultSet",
90+
result_set: ResultSet,
9791
column_index: int,
9892
allowed_values: List[str],
9993
case_sensitive: bool = False,
100-
) -> "ResultSet":
94+
) -> ResultSet:
10195
"""
10296
Filter a result set by values in a specific column.
10397
@@ -110,6 +104,7 @@ def filter_by_column_values(
110104
Returns:
111105
A filtered result set
112106
"""
107+
113108
# Convert to uppercase for case-insensitive comparison if needed
114109
if not case_sensitive:
115110
allowed_values = [v.upper() for v in allowed_values]
@@ -140,8 +135,8 @@ def filter_by_column_values(
140135

141136
@staticmethod
142137
def filter_tables_by_type(
143-
result_set: "ResultSet", table_types: Optional[List[str]] = None
144-
) -> "ResultSet":
138+
result_set: ResultSet, table_types: Optional[List[str]] = None
139+
) -> ResultSet:
145140
"""
146141
Filter a result set of tables by the specified table types.
147142
@@ -156,6 +151,7 @@ def filter_tables_by_type(
156151
Returns:
157152
A filtered result set containing only tables of the specified types
158153
"""
154+
159155
# Default table types if none specified
160156
DEFAULT_TABLE_TYPES = ["TABLE", "VIEW", "SYSTEM TABLE"]
161157
valid_types = (

0 commit comments

Comments
 (0)