19
19
OperationalError ,
20
20
SessionAlreadyClosedError ,
21
21
CursorAlreadyClosedError ,
22
+ Error ,
23
+ NotSupportedError ,
22
24
)
23
25
from databricks .sql .thrift_api .TCLIService import ttypes
24
26
from databricks .sql .backend .thrift_backend import ThriftDatabricksClient
@@ -249,6 +251,20 @@ def read(self) -> Optional[OAuthToken]:
249
251
+ str (self .get_session_id_hex ())
250
252
)
251
253
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
+
252
268
self .use_inline_params = self ._set_use_inline_params_with_warning (
253
269
kwargs .get ("use_inline_params" , False )
254
270
)
@@ -294,7 +310,7 @@ def __exit__(self, exc_type, exc_value, traceback):
294
310
return False
295
311
296
312
def __del__ (self ):
297
- if self .open :
313
+ if self .session . open :
298
314
logger .debug (
299
315
"Closing unclosed connection for session "
300
316
"{}" .format (self .get_session_id_hex ())
@@ -328,13 +344,6 @@ def get_protocol_version(openSessionResp):
328
344
"""Delegate to Session class static method"""
329
345
return Session .get_protocol_version (openSessionResp )
330
346
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
-
338
347
def cursor (
339
348
self ,
340
349
arraysize : int = DEFAULT_ARRAY_SIZE ,
@@ -345,7 +354,7 @@ def cursor(
345
354
346
355
Will throw an Error if the connection has been closed.
347
356
"""
348
- if not self .open :
357
+ if not self .session . open :
349
358
raise Error ("Cannot create cursor from closed connection" )
350
359
351
360
cursor = Cursor (
@@ -1490,7 +1499,7 @@ def close(self) -> None:
1490
1499
if (
1491
1500
self .op_state != ttypes .TOperationState .CLOSED_STATE
1492
1501
and not self .has_been_closed_server_side
1493
- and self .connection .open
1502
+ and self .connection .session . open
1494
1503
):
1495
1504
self .backend .close_command (self .command_id )
1496
1505
except RequestError as e :
0 commit comments