Skip to content

Commit 8ef6ed6

Browse files
decouple session class from existing Connection
ensure maintenance of current APIs of Connection while delegating responsibility Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 3658a91 commit 8ef6ed6

File tree

1 file changed

+19
-10
lines changed

1 file changed

+19
-10
lines changed

src/databricks/sql/client.py

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
OperationalError,
2020
SessionAlreadyClosedError,
2121
CursorAlreadyClosedError,
22+
Error,
23+
NotSupportedError,
2224
)
2325
from databricks.sql.thrift_api.TCLIService import ttypes
2426
from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
@@ -249,6 +251,20 @@ def read(self) -> Optional[OAuthToken]:
249251
+ str(self.get_session_id_hex())
250252
)
251253

254+
# Create the session
255+
self.session = Session(
256+
server_hostname,
257+
http_path,
258+
http_headers,
259+
session_configuration,
260+
catalog,
261+
schema,
262+
_use_arrow_native_complex_types,
263+
**kwargs
264+
)
265+
266+
logger.info("Successfully opened connection with session " + str(self.get_session_id_hex()))
267+
252268
self.use_inline_params = self._set_use_inline_params_with_warning(
253269
kwargs.get("use_inline_params", False)
254270
)
@@ -294,7 +310,7 @@ def __exit__(self, exc_type, exc_value, traceback):
294310
return False
295311

296312
def __del__(self):
297-
if self.open:
313+
if self.session.open:
298314
logger.debug(
299315
"Closing unclosed connection for session "
300316
"{}".format(self.get_session_id_hex())
@@ -328,13 +344,6 @@ def get_protocol_version(openSessionResp):
328344
"""Delegate to Session class static method"""
329345
return Session.get_protocol_version(openSessionResp)
330346

331-
@property
332-
def open(self) -> bool:
333-
"""Return whether the connection is open by checking if the session is open."""
334-
# NOTE: we have to check for the existence of session in case the __del__ is called
335-
# before the session is instantiated
336-
return hasattr(self, "session") and self.session.open
337-
338347
def cursor(
339348
self,
340349
arraysize: int = DEFAULT_ARRAY_SIZE,
@@ -345,7 +354,7 @@ def cursor(
345354
346355
Will throw an Error if the connection has been closed.
347356
"""
348-
if not self.open:
357+
if not self.session.open:
349358
raise Error("Cannot create cursor from closed connection")
350359

351360
cursor = Cursor(
@@ -1490,7 +1499,7 @@ def close(self) -> None:
14901499
if (
14911500
self.op_state != ttypes.TOperationState.CLOSED_STATE
14921501
and not self.has_been_closed_server_side
1493-
and self.connection.open
1502+
and self.connection.session.open
14941503
):
14951504
self.backend.close_command(self.command_id)
14961505
except RequestError as e:

0 commit comments

Comments
 (0)