|
| 1 | +from dataclasses import dataclass |
1 | 2 | from enum import Enum
|
2 |
| -from typing import Dict, Optional, Any |
| 3 | +from typing import Dict, List, Optional, Any, Tuple |
3 | 4 | import logging
|
4 | 5 |
|
5 | 6 | from databricks.sql.backend.utils import guid_to_hex_id
|
@@ -80,6 +81,26 @@ def from_thrift_state(
|
80 | 81 | else:
|
81 | 82 | return None
|
82 | 83 |
|
| 84 | + @classmethod |
| 85 | + def from_sea_state(cls, state: str) -> Optional["CommandState"]: |
| 86 | + """ |
| 87 | + Map SEA state string to CommandState enum. |
| 88 | + Args: |
| 89 | + state: SEA state string |
| 90 | + Returns: |
| 91 | + CommandState: The corresponding CommandState enum value |
| 92 | + """ |
| 93 | + state_mapping = { |
| 94 | + "PENDING": cls.PENDING, |
| 95 | + "RUNNING": cls.RUNNING, |
| 96 | + "SUCCEEDED": cls.SUCCEEDED, |
| 97 | + "FAILED": cls.FAILED, |
| 98 | + "CLOSED": cls.CLOSED, |
| 99 | + "CANCELED": cls.CANCELLED, |
| 100 | + } |
| 101 | + |
| 102 | + return state_mapping.get(state, None) |
| 103 | + |
83 | 104 |
|
84 | 105 | class BackendType(Enum):
|
85 | 106 | """
|
@@ -394,3 +415,19 @@ def to_hex_guid(self) -> str:
|
394 | 415 | return guid_to_hex_id(self.guid)
|
395 | 416 | else:
|
396 | 417 | return str(self.guid)
|
| 418 | + |
| 419 | + |
| 420 | +@dataclass |
| 421 | +class ExecuteResponse: |
| 422 | + """Response from executing a SQL command.""" |
| 423 | + |
| 424 | + command_id: CommandId |
| 425 | + status: CommandState |
| 426 | + description: Optional[ |
| 427 | + List[Tuple[str, str, None, None, Optional[int], Optional[int], bool]] |
| 428 | + ] = None |
| 429 | + has_more_rows: bool = False |
| 430 | + results_queue: Optional[Any] = None |
| 431 | + has_been_closed_server_side: bool = False |
| 432 | + lz4_compressed: bool = True |
| 433 | + is_staging_operation: bool = False |
0 commit comments