17
17
18
18
if TYPE_CHECKING :
19
19
from databricks .sql .client import Cursor
20
- from databricks .sql .result_set import SeaResultSet
20
+ from databricks .sql .backend . sea . result_set import SeaResultSet
21
21
22
22
from databricks .sql .backend .databricks_client import DatabricksClient
23
23
from databricks .sql .backend .types import (
@@ -253,7 +253,7 @@ def close_session(self, session_id: SessionId) -> None:
253
253
logger .debug ("SeaDatabricksClient.close_session(session_id=%s)" , session_id )
254
254
255
255
if session_id .backend_type != BackendType .SEA :
256
- raise ProgrammingError ("Not a valid SEA session ID" )
256
+ raise ValueError ("Not a valid SEA session ID" )
257
257
sea_session_id = session_id .to_sea_session_id ()
258
258
259
259
request_data = DeleteSessionRequest (
@@ -292,7 +292,7 @@ def get_allowed_session_configurations() -> List[str]:
292
292
293
293
def _extract_description_from_manifest (
294
294
self , manifest : ResultManifest
295
- ) -> Optional [ List ]:
295
+ ) -> List [ Tuple ]:
296
296
"""
297
297
Extract column description from a manifest object, in the format defined by
298
298
the spec: https://peps.python.org/pep-0249/#description
@@ -301,15 +301,12 @@ def _extract_description_from_manifest(
301
301
manifest: The ResultManifest object containing schema information
302
302
303
303
Returns:
304
- Optional[ List]: A list of column tuples or None if no columns are found
304
+ List[Tuple ]: A list of column tuples
305
305
"""
306
306
307
307
schema_data = manifest .schema
308
308
columns_data = schema_data .get ("columns" , [])
309
309
310
- if not columns_data :
311
- return None
312
-
313
310
columns = []
314
311
for col_data in columns_data :
315
312
# Format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
@@ -325,7 +322,7 @@ def _extract_description_from_manifest(
325
322
)
326
323
)
327
324
328
- return columns if columns else None
325
+ return columns
329
326
330
327
def _results_message_to_execute_response (
331
328
self , response : GetStatementResponse
@@ -431,7 +428,7 @@ def execute_command(
431
428
"""
432
429
433
430
if session_id .backend_type != BackendType .SEA :
434
- raise ProgrammingError ("Not a valid SEA session ID" )
431
+ raise ValueError ("Not a valid SEA session ID" )
435
432
436
433
sea_session_id = session_id .to_sea_session_id ()
437
434
@@ -510,9 +507,11 @@ def cancel_command(self, command_id: CommandId) -> None:
510
507
"""
511
508
512
509
if command_id .backend_type != BackendType .SEA :
513
- raise ProgrammingError ("Not a valid SEA command ID" )
510
+ raise ValueError ("Not a valid SEA command ID" )
514
511
515
512
sea_statement_id = command_id .to_sea_statement_id ()
513
+ if sea_statement_id is None :
514
+ raise ValueError ("Not a valid SEA command ID" )
516
515
517
516
request = CancelStatementRequest (statement_id = sea_statement_id )
518
517
self .http_client ._make_request (
@@ -533,9 +532,11 @@ def close_command(self, command_id: CommandId) -> None:
533
532
"""
534
533
535
534
if command_id .backend_type != BackendType .SEA :
536
- raise ProgrammingError ("Not a valid SEA command ID" )
535
+ raise ValueError ("Not a valid SEA command ID" )
537
536
538
537
sea_statement_id = command_id .to_sea_statement_id ()
538
+ if sea_statement_id is None :
539
+ raise ValueError ("Not a valid SEA command ID" )
539
540
540
541
request = CloseStatementRequest (statement_id = sea_statement_id )
541
542
self .http_client ._make_request (
@@ -562,6 +563,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
562
563
raise ProgrammingError ("Not a valid SEA command ID" )
563
564
564
565
sea_statement_id = command_id .to_sea_statement_id ()
566
+ if sea_statement_id is None :
567
+ raise ValueError ("Not a valid SEA command ID" )
565
568
566
569
request = GetStatementRequest (statement_id = sea_statement_id )
567
570
response_data = self .http_client ._make_request (
@@ -594,9 +597,11 @@ def get_execution_result(
594
597
"""
595
598
596
599
if command_id .backend_type != BackendType .SEA :
597
- raise ProgrammingError ("Not a valid SEA command ID" )
600
+ raise ValueError ("Not a valid SEA command ID" )
598
601
599
602
sea_statement_id = command_id .to_sea_statement_id ()
603
+ if sea_statement_id is None :
604
+ raise ValueError ("Not a valid SEA command ID" )
600
605
601
606
# Create the request model
602
607
request = GetStatementRequest (statement_id = sea_statement_id )
@@ -610,7 +615,7 @@ def get_execution_result(
610
615
response = GetStatementResponse .from_dict (response_data )
611
616
612
617
# Create and return a SeaResultSet
613
- from databricks .sql .result_set import SeaResultSet
618
+ from databricks .sql .backend . sea . result_set import SeaResultSet
614
619
615
620
execute_response = self ._results_message_to_execute_response (response )
616
621
0 commit comments