fix: separate session opening logic from instantiation

ensures correctness of self.session.open call in Connection

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

Author: varun-edachali-dbx

src/databricks/sql/client.py

@@ -241,6 +241,7 @@ def read(self) -> Optional[OAuthToken]:
             _use_arrow_native_complex_types,
             **kwargs,
         )
+        self.session.open()
 
         logger.info(
             "Successfully opened connection with session "
@@ -329,9 +330,7 @@ def get_protocol_version(openSessionResp):
     @property
     def open(self) -> bool:
         """Return whether the connection is open by checking if the session is open."""
-        # NOTE: we have to check for the existence of session in case the __del__ is called
-        # before the session is instantiated
-        return hasattr(self, "session") and self.session.open
+        return self.session.is_open
 
     def cursor(
         self,

src/databricks/sql/session.py

@@ -29,10 +29,14 @@ def __init__(
 
         This class handles all session-related behavior and communication with the backend.
         """
-        self.open = False
+        self.is_open = False
         self.host = server_hostname
         self.port = kwargs.get("_port", 443)
 
+        self.session_configuration = session_configuration
+        self.catalog = catalog
+        self.schema = schema
+
         auth_provider = get_python_sql_connector_auth_provider(
             server_hostname, **kwargs
         )
@@ -78,12 +82,16 @@ def __init__(
             **kwargs,
         )
 
+        self._session_handle = None
+        self.protocol_version = None
+
+    def open(self) -> None:
         self._open_session_resp = self.thrift_backend.open_session(
-            session_configuration, catalog, schema
+            self.session_configuration, self.catalog, self.schema
         )
         self._session_handle = self._open_session_resp.sessionHandle
         self.protocol_version = self.get_protocol_version(self._open_session_resp)
-        self.open = True
+        self.is_open = True
         logger.info("Successfully opened session " + str(self.get_session_id_hex()))
 
     @staticmethod
@@ -114,10 +122,16 @@ def get_session_handle(self):
         return self._session_handle
 
     def get_session_id(self):
-        return self.thrift_backend.handle_to_id(self._session_handle)
+        session_handle = self.get_session_handle()
+        if session_handle is None:
+            return None
+        return self.thrift_backend.handle_to_id(session_handle)
 
     def get_session_id_hex(self):
-        return self.thrift_backend.handle_to_hex_id(self._session_handle)
+        session_handle = self.get_session_handle()
+        if session_handle is None:
+            return None
+        return self.thrift_backend.handle_to_hex_id(session_handle)
 
     def close(self) -> None:
         """Close the underlying session."""
@@ -127,7 +141,7 @@ def close(self) -> None:
             return
 
         try:
-            self.thrift_backend.close_session(self._session_handle)
+            self.thrift_backend.close_session(self.get_session_handle())
         except RequestError as e:
             if isinstance(e.args[1], SessionAlreadyClosedError):
                 logger.info("Session was closed by a prior request")
@@ -143,4 +157,4 @@ def close(self) -> None:
         except Exception as e:
             logger.error(f"Attempt to close session raised a local exception: {e}")
 
-        self.open = False
+        self.is_open = False