Skip to content

Commit 634faa9

Browse files
chore: use get_handle() instead of private session attribute in client
Signed-off-by: varun-edachali-dbx <[email protected]>
1 parent 510b454 commit 634faa9

File tree

2 files changed

+23
-23
lines changed

2 files changed

+23
-23
lines changed

src/databricks/sql/client.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -306,11 +306,11 @@ def __del__(self):
306306

307307
def get_session_id(self):
308308
"""Get the session ID from the Session object"""
309-
return self.session.get_session_id()
309+
return self.session.get_id()
310310

311311
def get_session_id_hex(self):
312312
"""Get the session ID in hex format from the Session object"""
313-
return self.session.get_session_id_hex()
313+
return self.session.get_id_hex()
314314

315315
@staticmethod
316316
def server_parameterized_queries_enabled(protocolVersion):
@@ -776,7 +776,7 @@ def execute(
776776
self._close_and_clear_active_result_set()
777777
execute_response = self.thrift_backend.execute_command(
778778
operation=prepared_operation,
779-
session_handle=self.connection.session._session_handle,
779+
session_handle=self.connection.session.get_handle(),
780780
max_rows=self.arraysize,
781781
max_bytes=self.buffer_size_bytes,
782782
lz4_compression=self.connection.lz4_compression,
@@ -839,7 +839,7 @@ def execute_async(
839839
self._close_and_clear_active_result_set()
840840
self.thrift_backend.execute_command(
841841
operation=prepared_operation,
842-
session_handle=self.connection.session._session_handle,
842+
session_handle=self.connection.session.get_handle(),
843843
max_rows=self.arraysize,
844844
max_bytes=self.buffer_size_bytes,
845845
lz4_compression=self.connection.lz4_compression,
@@ -935,7 +935,7 @@ def catalogs(self) -> "Cursor":
935935
self._check_not_closed()
936936
self._close_and_clear_active_result_set()
937937
execute_response = self.thrift_backend.get_catalogs(
938-
session_handle=self.connection.session._session_handle,
938+
session_handle=self.connection.session.get_handle(),
939939
max_rows=self.arraysize,
940940
max_bytes=self.buffer_size_bytes,
941941
cursor=self,
@@ -961,7 +961,7 @@ def schemas(
961961
self._check_not_closed()
962962
self._close_and_clear_active_result_set()
963963
execute_response = self.thrift_backend.get_schemas(
964-
session_handle=self.connection.session._session_handle,
964+
session_handle=self.connection.session.get_handle(),
965965
max_rows=self.arraysize,
966966
max_bytes=self.buffer_size_bytes,
967967
cursor=self,
@@ -994,7 +994,7 @@ def tables(
994994
self._close_and_clear_active_result_set()
995995

996996
execute_response = self.thrift_backend.get_tables(
997-
session_handle=self.connection.session._session_handle,
997+
session_handle=self.connection.session.get_handle(),
998998
max_rows=self.arraysize,
999999
max_bytes=self.buffer_size_bytes,
10001000
cursor=self,
@@ -1029,7 +1029,7 @@ def columns(
10291029
self._close_and_clear_active_result_set()
10301030

10311031
execute_response = self.thrift_backend.get_columns(
1032-
session_handle=self.connection.session._session_handle,
1032+
session_handle=self.connection.session.get_handle(),
10331033
max_rows=self.arraysize,
10341034
max_bytes=self.buffer_size_bytes,
10351035
cursor=self,

src/databricks/sql/session.py

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -82,17 +82,17 @@ def __init__(
8282
**kwargs,
8383
)
8484

85-
self._session_handle = None
85+
self._handle = None
8686
self.protocol_version = None
8787

8888
def open(self) -> None:
8989
self._open_session_resp = self.thrift_backend.open_session(
9090
self.session_configuration, self.catalog, self.schema
9191
)
92-
self._session_handle = self._open_session_resp.sessionHandle
92+
self._handle = self._open_session_resp.sessionHandle
9393
self.protocol_version = self.get_protocol_version(self._open_session_resp)
9494
self.is_open = True
95-
logger.info("Successfully opened session " + str(self.get_session_id_hex()))
95+
logger.info("Successfully opened session " + str(self.get_id_hex()))
9696

9797
@staticmethod
9898
def get_protocol_version(openSessionResp):
@@ -118,30 +118,30 @@ def server_parameterized_queries_enabled(protocolVersion):
118118
else:
119119
return False
120120

121-
def get_session_handle(self):
122-
return self._session_handle
121+
def get_handle(self):
122+
return self._handle
123123

124-
def get_session_id(self):
125-
session_handle = self.get_session_handle()
126-
if session_handle is None:
124+
def get_id(self):
125+
handle = self.get_handle()
126+
if handle is None:
127127
return None
128-
return self.thrift_backend.handle_to_id(session_handle)
128+
return self.thrift_backend.handle_to_id(handle)
129129

130-
def get_session_id_hex(self):
131-
session_handle = self.get_session_handle()
132-
if session_handle is None:
130+
def get_id_hex(self):
131+
handle = self.get_handle()
132+
if handle is None:
133133
return None
134-
return self.thrift_backend.handle_to_hex_id(session_handle)
134+
return self.thrift_backend.handle_to_hex_id(handle)
135135

136136
def close(self) -> None:
137137
"""Close the underlying session."""
138-
logger.info(f"Closing session {self.get_session_id_hex()}")
138+
logger.info(f"Closing session {self.get_id_hex()}")
139139
if not self.is_open:
140140
logger.debug("Session appears to have been closed already")
141141
return
142142

143143
try:
144-
self.thrift_backend.close_session(self.get_session_handle())
144+
self.thrift_backend.close_session(self.get_handle())
145145
except RequestError as e:
146146
if isinstance(e.args[1], SessionAlreadyClosedError):
147147
logger.info("Session was closed by a prior request")

0 commit comments

Comments
 (0)