Skip to content

Commit 1b0fc9b

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 eac0433 commit 1b0fc9b

File tree

1 file changed

+18
-2
lines changed

1 file changed

+18
-2
lines changed

src/databricks/sql/client.py

Lines changed: 18 additions & 2 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
@@ -248,6 +250,20 @@ def read(self) -> Optional[OAuthToken]:
248250
+ str(self.get_session_id_hex())
249251
)
250252

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

295311
def __del__(self):
296-
if self.open:
312+
if self.session.open:
297313
logger.debug(
298314
"Closing unclosed connection for session "
299315
"{}".format(self.get_session_id_hex())
@@ -344,7 +360,7 @@ def cursor(
344360
345361
Will throw an Error if the connection has been closed.
346362
"""
347-
if not self.open:
363+
if not self.session.open:
348364
raise Error("Cannot create cursor from closed connection")
349365

350366
cursor = Cursor(

0 commit comments

Comments
 (0)