Skip to content

Commit 574df21

Browse files
create concrete ThriftResultSet instead of ResultSet type
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent d4a49ca commit 574df21

File tree

1 file changed

+5
-4
lines changed

1 file changed

+5
-4
lines changed

tests/unit/test_fetches.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import databricks.sql.client as client
1111
from databricks.sql.utils import ExecuteResponse, ArrowQueue
1212
from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
13+
from databricks.sql.result_set import ThriftResultSet
1314

1415

1516
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
@@ -38,9 +39,8 @@ def make_dummy_result_set_from_initial_results(initial_results):
3839
# If the initial results have been set, then we should never try and fetch more
3940
schema, arrow_table = FetchTests.make_arrow_table(initial_results)
4041
arrow_queue = ArrowQueue(arrow_table, len(initial_results), 0)
41-
rs = client.ResultSet(
42+
rs = ThriftResultSet(
4243
connection=Mock(),
43-
backend=None,
4444
execute_response=ExecuteResponse(
4545
status=None,
4646
has_been_closed_server_side=True,
@@ -52,6 +52,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
5252
arrow_schema_bytes=schema.serialize().to_pybytes(),
5353
is_staging_operation=False,
5454
),
55+
thrift_client=None,
5556
)
5657
num_cols = len(initial_results[0]) if initial_results else 0
5758
rs.description = [
@@ -84,9 +85,8 @@ def fetch_results(
8485
mock_thrift_backend.fetch_results = fetch_results
8586
num_cols = len(batch_list[0][0]) if batch_list and batch_list[0] else 0
8687

87-
rs = client.ResultSet(
88+
rs = ThriftResultSet(
8889
connection=Mock(),
89-
backend=mock_thrift_backend,
9090
execute_response=ExecuteResponse(
9191
status=None,
9292
has_been_closed_server_side=False,
@@ -101,6 +101,7 @@ def fetch_results(
101101
arrow_schema_bytes=None,
102102
is_staging_operation=False,
103103
),
104+
thrift_client=mock_thrift_backend,
104105
)
105106
return rs
106107

0 commit comments

Comments
 (0)