11from __future__ import annotations
22
33from 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
86try :
97 import pyarrow
1210
1311import 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+ )
2120from databricks .sql .backend .sea .utils .constants import ResultFormat
2221from databricks .sql .exc import ProgrammingError , ServerOperationError
2322from 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