Skip to content

Commit 3cc1a4f

Browse files
introduce normalised sessionId and CommandId for (near) complete backend abstraction
status is still aligned with Thrift
1 parent 1d026ad commit 3cc1a4f

File tree

6 files changed

+414
-145
lines changed

6 files changed

+414
-145
lines changed

src/databricks/sql/backend/databricks_client.py

Lines changed: 14 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from typing import Dict, Tuple, List, Optional, Any, Union
33

44
from databricks.sql.thrift_api.TCLIService import ttypes
5+
from databricks.sql.ids import SessionId, CommandId
56
from databricks.sql.utils import ExecuteResponse
67
from databricks.sql.types import SSLOptions
78

@@ -14,19 +15,19 @@ def open_session(
1415
session_configuration: Optional[Dict[str, Any]],
1516
catalog: Optional[str],
1617
schema: Optional[str],
17-
) -> ttypes.TOpenSessionResp:
18+
) -> SessionId:
1819
pass
1920

2021
@abstractmethod
21-
def close_session(self, session_handle: ttypes.TSessionHandle) -> None:
22+
def close_session(self, session_id: SessionId) -> None:
2223
pass
2324

2425
# == Query Execution, Command Management ==
2526
@abstractmethod
2627
def execute_command(
2728
self,
2829
operation: str,
29-
session_handle: ttypes.TSessionHandle,
30+
session_id: SessionId,
3031
max_rows: int,
3132
max_bytes: int,
3233
lz4_compression: bool,
@@ -39,25 +40,21 @@ def execute_command(
3940
pass
4041

4142
@abstractmethod
42-
def cancel_command(self, operation_handle: ttypes.TOperationHandle) -> None:
43+
def cancel_command(self, command_id: CommandId) -> None:
4344
pass
4445

4546
@abstractmethod
46-
def close_command(
47-
self, operation_handle: ttypes.TOperationHandle
48-
) -> ttypes.TStatus:
47+
def close_command(self, command_id: CommandId) -> ttypes.TStatus:
4948
pass
5049

5150
@abstractmethod
52-
def get_query_state(
53-
self, operation_handle: ttypes.TOperationHandle
54-
) -> ttypes.TOperationState:
51+
def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
5552
pass
5653

5754
@abstractmethod
5855
def get_execution_result(
5956
self,
60-
operation_handle: ttypes.TOperationHandle,
57+
command_id: CommandId,
6158
cursor: Any,
6259
) -> ExecuteResponse:
6360
pass
@@ -66,7 +63,7 @@ def get_execution_result(
6663
@abstractmethod
6764
def get_catalogs(
6865
self,
69-
session_handle: ttypes.TSessionHandle,
66+
session_id: SessionId,
7067
max_rows: int,
7168
max_bytes: int,
7269
cursor: Any,
@@ -76,7 +73,7 @@ def get_catalogs(
7673
@abstractmethod
7774
def get_schemas(
7875
self,
79-
session_handle: ttypes.TSessionHandle,
76+
session_id: SessionId,
8077
max_rows: int,
8178
max_bytes: int,
8279
cursor: Any,
@@ -88,7 +85,7 @@ def get_schemas(
8885
@abstractmethod
8986
def get_tables(
9087
self,
91-
session_handle: ttypes.TSessionHandle,
88+
session_id: SessionId,
9289
max_rows: int,
9390
max_bytes: int,
9491
cursor: Any,
@@ -102,7 +99,7 @@ def get_tables(
10299
@abstractmethod
103100
def get_columns(
104101
self,
105-
session_handle: ttypes.TSessionHandle,
102+
session_id: SessionId,
106103
max_rows: int,
107104
max_bytes: int,
108105
cursor: Any,
@@ -115,11 +112,11 @@ def get_columns(
115112

116113
# == Utility Methods ==
117114
@abstractmethod
118-
def handle_to_id(self, handle: ttypes.TSessionHandle) -> bytes:
115+
def handle_to_id(self, session_id: SessionId) -> Any:
119116
pass
120117

121118
@abstractmethod
122-
def handle_to_hex_id(self, handle: ttypes.TSessionHandle) -> str:
119+
def handle_to_hex_id(self, session_id: SessionId) -> str:
123120
pass
124121

125122
# Properties related to specific backend features

0 commit comments

Comments
 (0)