Skip to content

Commit 54e0570

Browse files
remove string literals around type defs
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent c1a9e90 commit 54e0570

File tree

1 file changed

+12
-13
lines changed

1 file changed

+12
-13
lines changed

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

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
from __future__ import annotations
22

33
from abc import ABC
4-
from typing import List, Optional, Tuple, Union
5-
6-
from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
4+
from typing import List, Optional, Tuple, Union, TYPE_CHECKING
75

86
try:
97
import pyarrow
@@ -12,12 +10,13 @@
1210

1311
import dateutil
1412

15-
from databricks.sql.backend.sea.backend import SeaDatabricksClient
16-
from databricks.sql.backend.sea.models.base import (
17-
ExternalLink,
18-
ResultData,
19-
ResultManifest,
20-
)
13+
if TYPE_CHECKING:
14+
from databricks.sql.backend.sea.backend import SeaDatabricksClient
15+
from databricks.sql.backend.sea.models.base import (
16+
ExternalLink,
17+
ResultData,
18+
ResultManifest,
19+
)
2120
from databricks.sql.backend.sea.utils.constants import ResultFormat
2221
from databricks.sql.exc import ProgrammingError, ServerOperationError
2322
from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
@@ -109,7 +108,7 @@ def __init__(
109108
result_data: ResultData,
110109
max_download_threads: int,
111110
ssl_options: SSLOptions,
112-
sea_client: "SeaDatabricksClient",
111+
sea_client: SeaDatabricksClient,
113112
statement_id: str,
114113
total_chunk_count: int,
115114
lz4_compressed: bool = False,
@@ -159,7 +158,7 @@ def __init__(
159158
# Initialize table and position
160159
self.table = self._create_table_from_link(first_link)
161160

162-
def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink:
161+
def _convert_to_thrift_link(self, link: ExternalLink) -> TSparkArrowResultLink:
163162
"""Convert SEA external links to Thrift format for compatibility with existing download manager."""
164163
# Parse the ISO format expiration time
165164
expiry_time = int(dateutil.parser.parse(link.expiration).timestamp())
@@ -172,7 +171,7 @@ def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink
172171
httpHeaders=link.http_headers or {},
173172
)
174173

175-
def _get_chunk_link(self, chunk_index: int) -> "ExternalLink":
174+
def _get_chunk_link(self, chunk_index: int) -> ExternalLink:
176175
"""Progress to the next chunk link."""
177176
if chunk_index >= self._total_chunk_count:
178177
raise ValueError(
@@ -191,7 +190,7 @@ def _get_chunk_link(self, chunk_index: int) -> "ExternalLink":
191190
)
192191

193192
def _create_table_from_link(
194-
self, link: "ExternalLink"
193+
self, link: ExternalLink
195194
) -> Union["pyarrow.Table", None]:
196195
"""Create a table from a link."""
197196

0 commit comments

Comments
 (0)