10
10
import databricks .sql .client as client
11
11
from databricks .sql .utils import ExecuteResponse , ArrowQueue
12
12
from databricks .sql .backend .thrift_backend import ThriftDatabricksClient
13
+ from databricks .sql .result_set import ThriftResultSet
13
14
14
15
15
16
@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):
38
39
# If the initial results have been set, then we should never try and fetch more
39
40
schema , arrow_table = FetchTests .make_arrow_table (initial_results )
40
41
arrow_queue = ArrowQueue (arrow_table , len (initial_results ), 0 )
41
- rs = client . ResultSet (
42
+ rs = ThriftResultSet (
42
43
connection = Mock (),
43
- backend = None ,
44
44
execute_response = ExecuteResponse (
45
45
status = None ,
46
46
has_been_closed_server_side = True ,
@@ -52,6 +52,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
52
52
arrow_schema_bytes = schema .serialize ().to_pybytes (),
53
53
is_staging_operation = False ,
54
54
),
55
+ thrift_client = None ,
55
56
)
56
57
num_cols = len (initial_results [0 ]) if initial_results else 0
57
58
rs .description = [
@@ -84,9 +85,8 @@ def fetch_results(
84
85
mock_thrift_backend .fetch_results = fetch_results
85
86
num_cols = len (batch_list [0 ][0 ]) if batch_list and batch_list [0 ] else 0
86
87
87
- rs = client . ResultSet (
88
+ rs = ThriftResultSet (
88
89
connection = Mock (),
89
- backend = mock_thrift_backend ,
90
90
execute_response = ExecuteResponse (
91
91
status = None ,
92
92
has_been_closed_server_side = False ,
@@ -101,6 +101,7 @@ def fetch_results(
101
101
arrow_schema_bytes = None ,
102
102
is_staging_operation = False ,
103
103
),
104
+ thrift_client = mock_thrift_backend ,
104
105
)
105
106
return rs
106
107
0 commit comments