Skip to content

Commit e863ac3

Browse files
fix: separate session opening logic from instantiation
ensures correctness of self.session.open call in Connection Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 9b4f9b3 commit e863ac3

File tree

2 files changed

+23
-10
lines changed

2 files changed

+23
-10
lines changed

src/databricks/sql/client.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,7 @@ def read(self) -> Optional[OAuthToken]:
241241
_use_arrow_native_complex_types,
242242
**kwargs,
243243
)
244+
self.session.open()
244245

245246
logger.info(
246247
"Successfully opened connection with session "
@@ -329,9 +330,7 @@ def get_protocol_version(openSessionResp):
329330
@property
330331
def open(self) -> bool:
331332
"""Return whether the connection is open by checking if the session is open."""
332-
# NOTE: we have to check for the existence of session in case the __del__ is called
333-
# before the session is instantiated
334-
return hasattr(self, "session") and self.session.open
333+
return self.session.is_open
335334

336335
def cursor(
337336
self,

src/databricks/sql/session.py

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,14 @@ def __init__(
2929
3030
This class handles all session-related behavior and communication with the backend.
3131
"""
32-
self.open = False
32+
self.is_open = False
3333
self.host = server_hostname
3434
self.port = kwargs.get("_port", 443)
3535

36+
self.session_configuration = session_configuration
37+
self.catalog = catalog
38+
self.schema = schema
39+
3640
auth_provider = get_python_sql_connector_auth_provider(
3741
server_hostname, **kwargs
3842
)
@@ -78,12 +82,16 @@ def __init__(
7882
**kwargs,
7983
)
8084

85+
self._session_handle = None
86+
self.protocol_version = None
87+
88+
def open(self) -> None:
8189
self._open_session_resp = self.thrift_backend.open_session(
82-
session_configuration, catalog, schema
90+
self.session_configuration, self.catalog, self.schema
8391
)
8492
self._session_handle = self._open_session_resp.sessionHandle
8593
self.protocol_version = self.get_protocol_version(self._open_session_resp)
86-
self.open = True
94+
self.is_open = True
8795
logger.info("Successfully opened session " + str(self.get_session_id_hex()))
8896

8997
@staticmethod
@@ -114,10 +122,16 @@ def get_session_handle(self):
114122
return self._session_handle
115123

116124
def get_session_id(self):
117-
return self.thrift_backend.handle_to_id(self._session_handle)
125+
session_handle = self.get_session_handle()
126+
if session_handle is None:
127+
return None
128+
return self.thrift_backend.handle_to_id(session_handle)
118129

119130
def get_session_id_hex(self):
120-
return self.thrift_backend.handle_to_hex_id(self._session_handle)
131+
session_handle = self.get_session_handle()
132+
if session_handle is None:
133+
return None
134+
return self.thrift_backend.handle_to_hex_id(session_handle)
121135

122136
def close(self) -> None:
123137
"""Close the underlying session."""
@@ -127,7 +141,7 @@ def close(self) -> None:
127141
return
128142

129143
try:
130-
self.thrift_backend.close_session(self._session_handle)
144+
self.thrift_backend.close_session(self.get_session_handle())
131145
except RequestError as e:
132146
if isinstance(e.args[1], SessionAlreadyClosedError):
133147
logger.info("Session was closed by a prior request")
@@ -143,4 +157,4 @@ def close(self) -> None:
143157
except Exception as e:
144158
logger.error(f"Attempt to close session raised a local exception: {e}")
145159

146-
self.open = False
160+
self.is_open = False

0 commit comments

Comments
 (0)