Skip to content

Commit 0bdf8f9

Browse files
make Manifest mandatory
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 3f1fd93 commit 0bdf8f9

File tree

2 files changed

+22
-3
lines changed

2 files changed

+22
-3
lines changed

src/databricks/sql/result_set.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -456,7 +456,7 @@ def __init__(
456456
execute_response: ExecuteResponse,
457457
sea_client: SeaDatabricksClient,
458458
result_data: ResultData,
459-
manifest: Optional[ResultManifest] = None,
459+
manifest: ResultManifest,
460460
buffer_size_bytes: int = 104857600,
461461
arraysize: int = 10000,
462462
):

tests/unit/test_sea_result_set.py

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@
66
"""
77

88
import pytest
9-
from unittest.mock import patch, MagicMock, Mock
9+
from unittest.mock import Mock
1010

1111
from databricks.sql.result_set import SeaResultSet, Row
1212
from databricks.sql.utils import JsonQueue
13-
from databricks.sql.backend.types import CommandId, CommandState, BackendType
13+
from databricks.sql.backend.types import CommandId, CommandState
1414
from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
1515

1616

@@ -88,6 +88,18 @@ def json_queue(self, sample_data):
8888
"""Create a JsonQueue with sample data."""
8989
return JsonQueue(sample_data)
9090

91+
def empty_manifest(self):
92+
"""Create an empty manifest."""
93+
return ResultManifest(
94+
format="JSON_ARRAY",
95+
schema={},
96+
total_row_count=0,
97+
total_byte_count=0,
98+
total_chunk_count=0,
99+
truncated=False,
100+
is_volume_operation=False,
101+
)
102+
91103
def test_init_with_execute_response(
92104
self, mock_connection, mock_sea_client, execute_response
93105
):
@@ -97,6 +109,7 @@ def test_init_with_execute_response(
97109
execute_response=execute_response,
98110
sea_client=mock_sea_client,
99111
result_data=ResultData(data=[]),
112+
manifest=self.empty_manifest(),
100113
buffer_size_bytes=1000,
101114
arraysize=100,
102115
)
@@ -117,6 +130,7 @@ def test_close(self, mock_connection, mock_sea_client, execute_response):
117130
execute_response=execute_response,
118131
sea_client=mock_sea_client,
119132
result_data=ResultData(data=[]),
133+
manifest=self.empty_manifest(),
120134
buffer_size_bytes=1000,
121135
arraysize=100,
122136
)
@@ -138,6 +152,7 @@ def test_close_when_already_closed_server_side(
138152
execute_response=execute_response,
139153
sea_client=mock_sea_client,
140154
result_data=ResultData(data=[]),
155+
manifest=self.empty_manifest(),
141156
buffer_size_bytes=1000,
142157
arraysize=100,
143158
)
@@ -161,6 +176,7 @@ def test_close_when_connection_closed(
161176
execute_response=execute_response,
162177
sea_client=mock_sea_client,
163178
result_data=ResultData(data=[]),
179+
manifest=self.empty_manifest(),
164180
buffer_size_bytes=1000,
165181
arraysize=100,
166182
)
@@ -317,6 +333,7 @@ def test_fetchmany_arrow_not_implemented(
317333
execute_response=execute_response,
318334
sea_client=mock_sea_client,
319335
result_data=ResultData(data=None, external_links=[]),
336+
manifest=self.empty_manifest(),
320337
buffer_size_bytes=1000,
321338
arraysize=100,
322339
)
@@ -336,6 +353,7 @@ def test_fetchall_arrow_not_implemented(
336353
execute_response=execute_response,
337354
sea_client=mock_sea_client,
338355
result_data=ResultData(data=None, external_links=[]),
356+
manifest=self.empty_manifest(),
339357
buffer_size_bytes=1000,
340358
arraysize=100,
341359
)
@@ -353,6 +371,7 @@ def test_is_staging_operation(
353371
execute_response=execute_response,
354372
sea_client=mock_sea_client,
355373
result_data=ResultData(data=[]),
374+
manifest=self.empty_manifest(),
356375
buffer_size_bytes=1000,
357376
arraysize=100,
358377
)

0 commit comments

Comments
 (0)