Skip to content

Commit cd119e9

Browse files
fix type issues
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 10f55f0 commit cd119e9

File tree

3 files changed

+11
-9
lines changed

3 files changed

+11
-9
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import logging
44
import time
55
import re
6-
from typing import Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
6+
from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
77

88
from databricks.sql.backend.sea.models.base import ExternalLink, ResultManifest
99
from databricks.sql.backend.sea.utils.constants import (

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
class SeaResultSetQueueFactory(ABC):
3333
@staticmethod
3434
def build_queue(
35-
sea_result_data: ResultData,
35+
result_data: ResultData,
3636
manifest: ResultManifest,
3737
statement_id: str,
3838
ssl_options: Optional[SSLOptions] = None,
@@ -45,7 +45,7 @@ def build_queue(
4545
Factory method to build a result set queue for SEA backend.
4646
4747
Args:
48-
sea_result_data (ResultData): Result data from SEA response
48+
result_data (ResultData): Result data from SEA response
4949
manifest (ResultManifest): Manifest from SEA response
5050
statement_id (str): Statement ID for the query
5151
description (List[List[Any]]): Column descriptions
@@ -59,7 +59,7 @@ def build_queue(
5959

6060
if manifest.format == ResultFormat.JSON_ARRAY.value:
6161
# INLINE disposition with JSON_ARRAY format
62-
return JsonQueue(sea_result_data.data)
62+
return JsonQueue(result_data.data)
6363
elif manifest.format == ResultFormat.ARROW_STREAM.value:
6464
# EXTERNAL_LINKS disposition
6565
if not max_download_threads:
@@ -74,11 +74,13 @@ def build_queue(
7474
raise ValueError(
7575
"SEA client is required for EXTERNAL_LINKS disposition"
7676
)
77-
if not manifest:
78-
raise ValueError("Manifest is required for EXTERNAL_LINKS disposition")
77+
if not result_data.external_links:
78+
raise ValueError(
79+
"External links are required for EXTERNAL_LINKS disposition"
80+
)
7981

8082
return SeaCloudFetchQueue(
81-
initial_links=sea_result_data.external_links,
83+
initial_links=result_data.external_links,
8284
max_download_threads=max_download_threads,
8385
ssl_options=ssl_options,
8486
sea_client=sea_client,
@@ -125,7 +127,7 @@ def __init__(
125127
statement_id: str,
126128
total_chunk_count: int,
127129
lz4_compressed: bool = False,
128-
description: Optional[List[Tuple]] = None,
130+
description: List[Tuple] = [],
129131
):
130132
"""
131133
Initialize the SEA CloudFetchQueue.

src/databricks/sql/utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -369,7 +369,7 @@ def __init__(
369369
start_row_offset: int = 0,
370370
result_links: Optional[List[TSparkArrowResultLink]] = None,
371371
lz4_compressed: bool = True,
372-
description: Optional[List[Tuple]] = None,
372+
description: List[Tuple] = [],
373373
):
374374
"""
375375
Initialize the Thrift CloudFetchQueue.

0 commit comments

Comments
 (0)