chore: use get_handle() instead of private session attribute in client

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/client.py

@@ -306,11 +306,11 @@ def __del__(self):
 
     def get_session_id(self):
         """Get the session ID from the Session object"""
-        return self.session.get_session_id()
+        return self.session.get_id()
 
     def get_session_id_hex(self):
         """Get the session ID in hex format from the Session object"""
-        return self.session.get_session_id_hex()
+        return self.session.get_id_hex()
 
     @staticmethod
     def server_parameterized_queries_enabled(protocolVersion):
@@ -776,7 +776,7 @@ def execute(
         self._close_and_clear_active_result_set()
         execute_response = self.thrift_backend.execute_command(
             operation=prepared_operation,
-            session_handle=self.connection.session._session_handle,
+            session_handle=self.connection.session.get_handle(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             lz4_compression=self.connection.lz4_compression,
@@ -839,7 +839,7 @@ def execute_async(
         self._close_and_clear_active_result_set()
         self.thrift_backend.execute_command(
             operation=prepared_operation,
-            session_handle=self.connection.session._session_handle,
+            session_handle=self.connection.session.get_handle(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             lz4_compression=self.connection.lz4_compression,
@@ -935,7 +935,7 @@ def catalogs(self) -> "Cursor":
         self._check_not_closed()
         self._close_and_clear_active_result_set()
         execute_response = self.thrift_backend.get_catalogs(
-            session_handle=self.connection.session._session_handle,
+            session_handle=self.connection.session.get_handle(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,
@@ -961,7 +961,7 @@ def schemas(
         self._check_not_closed()
         self._close_and_clear_active_result_set()
         execute_response = self.thrift_backend.get_schemas(
-            session_handle=self.connection.session._session_handle,
+            session_handle=self.connection.session.get_handle(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,
@@ -994,7 +994,7 @@ def tables(
         self._close_and_clear_active_result_set()
 
         execute_response = self.thrift_backend.get_tables(
-            session_handle=self.connection.session._session_handle,
+            session_handle=self.connection.session.get_handle(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,
@@ -1029,7 +1029,7 @@ def columns(
         self._close_and_clear_active_result_set()
 
         execute_response = self.thrift_backend.get_columns(
-            session_handle=self.connection.session._session_handle,
+            session_handle=self.connection.session.get_handle(),
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,

src/databricks/sql/session.py

@@ -82,17 +82,17 @@ def __init__(
             **kwargs,
         )
 
-        self._session_handle = None
+        self._handle = None
         self.protocol_version = None
 
     def open(self) -> None:
         self._open_session_resp = self.thrift_backend.open_session(
             self.session_configuration, self.catalog, self.schema
         )
-        self._session_handle = self._open_session_resp.sessionHandle
+        self._handle = self._open_session_resp.sessionHandle
         self.protocol_version = self.get_protocol_version(self._open_session_resp)
         self.is_open = True
-        logger.info("Successfully opened session " + str(self.get_session_id_hex()))
+        logger.info("Successfully opened session " + str(self.get_id_hex()))
 
     @staticmethod
     def get_protocol_version(openSessionResp):
@@ -118,30 +118,30 @@ def server_parameterized_queries_enabled(protocolVersion):
         else:
             return False
 
-    def get_session_handle(self):
-        return self._session_handle
+    def get_handle(self):
+        return self._handle
 
-    def get_session_id(self):
-        session_handle = self.get_session_handle()
-        if session_handle is None:
+    def get_id(self):
+        handle = self.get_handle()
+        if handle is None:
             return None
-        return self.thrift_backend.handle_to_id(session_handle)
+        return self.thrift_backend.handle_to_id(handle)
 
-    def get_session_id_hex(self):
-        session_handle = self.get_session_handle()
-        if session_handle is None:
+    def get_id_hex(self):
+        handle = self.get_handle()
+        if handle is None:
             return None
-        return self.thrift_backend.handle_to_hex_id(session_handle)
+        return self.thrift_backend.handle_to_hex_id(handle)
 
     def close(self) -> None:
         """Close the underlying session."""
-        logger.info(f"Closing session {self.get_session_id_hex()}")
+        logger.info(f"Closing session {self.get_id_hex()}")
         if not self.is_open:
             logger.debug("Session appears to have been closed already")
             return
 
         try:
-            self.thrift_backend.close_session(self.get_session_handle())
+            self.thrift_backend.close_session(self.get_handle())
         except RequestError as e:
             if isinstance(e.args[1], SessionAlreadyClosedError):
                 logger.info("Session was closed by a prior request")