Separate Session related functionality from Connection class #571

varun-edachali-dbx · 2025-05-26T05:22:40Z

What type of PR is this?

Refactor

Description

Separate the Session related functionality from the Connection class in order to further abstract the backend implementation details from the connection, to eventually give way to the introduction of SEA.

How is this tested?

Unit tests
- some of the existing unit tests were slightly altered to account for the introduction of the Session class. No unit tests were removed or introduced.
E2E Tests
Manually
N/A

Related Tickets & Documents

https://docs.google.com/document/d/1Y-eXLhNqqhrMVGnOlG8sdFrCxBTN1GdQvuKG4IfHmo0/edit?usp=sharing
https://databricks.atlassian.net/browse/PECOBLR-438?atlOrigin=eyJpIjoiYjBiMjlhYjBjM2RhNDdiOGE4OGU4NjgwYzU1ZmIzNzciLCJwIjoiaiJ9

ensure maintenance of current APIs of Connection while delegating responsibility Signed-off-by: varun-edachali-dbx <[email protected]>

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

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

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

as in CONTRIBUTING.md Signed-off-by: varun-edachali-dbx <[email protected]>

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

in case the openSession takes long, the initialisation of the session will not complete immediately. This could make the session attribute inaccessible. If the Connection is deleted in this time, the open() check will throw because the session attribute does not exist. Thus, we default to the Connection being closed in this case. This was not an issue before because open was a direct attribute of the Connection class. Caught in the integration tests. Signed-off-by: varun-edachali-dbx <[email protected]>

earlier, one of the integration tests was failing because 'session was not an attribute of Connection'. This is likely tied to a local configuration issue related to unittest that was causing an error in the test suite itself. The tests are now passing without checking for the session attribute. c676f9b Signed-off-by: varun-edachali-dbx <[email protected]>

This reverts commit d6b1b19. Signed-off-by: varun-edachali-dbx <[email protected]>

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

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

…t to prevent server side resource leaks (#554) * Enhance Cursor close handling and context manager exception management * tests * fmt * Fix Cursor.close() to properly handle CursorAlreadyClosedError * Remove specific test message from Cursor.close() error handling * Improve error handling in connection and cursor context managers to ensure proper closure during exceptions, including KeyboardInterrupt. Add tests for nested cursor management and verify operation closure on server-side errors. * add * add Signed-off-by: varun-edachali-dbx <[email protected]>

* PECOBLR-86 Improve logging for debug level Signed-off-by: Sai Shree Pradhan <[email protected]> * PECOBLR-86 Improve logging for debug level Signed-off-by: Sai Shree Pradhan <[email protected]> * fixed format Signed-off-by: Sai Shree Pradhan <[email protected]> * used lazy logging Signed-off-by: Sai Shree Pradhan <[email protected]> * changed debug to error logs Signed-off-by: Sai Shree Pradhan <[email protected]> * used lazy logging Signed-off-by: Sai Shree Pradhan <[email protected]> --------- Signed-off-by: Sai Shree Pradhan <[email protected]> Signed-off-by: varun-edachali-dbx <[email protected]>

…couple-session" This reverts commit dbb2ec5, reversing changes made to 7192f11. Signed-off-by: varun-edachali-dbx <[email protected]>

…ecouple-session" This reverts commit bdb8381. Signed-off-by: varun-edachali-dbx <[email protected]>

ensures correctness of self.session.open call in Connection Signed-off-by: varun-edachali-dbx <[email protected]>

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

…ession

jayantsing-db

Hi Varun, thanks for these changes. I have added some minor comments. Please take a look.

src/databricks/sql/session.py

jayantsing-db · 2025-05-26T06:28:40Z

src/databricks/sql/session.py

+        self,
+        server_hostname: str,
+        http_path: str,
+        http_headers: Optional[List[Tuple[str, str]]] = None,
+        session_configuration: Optional[Dict[str, Any]] = None,
+        catalog: Optional[str] = None,
+        schema: Optional[str] = None,
+        _use_arrow_native_complex_types: Optional[bool] = True,
+        **kwargs,


nit

these are identical to connection params. are all of them relevant to a DBSQL session?

should be introduce a type/namedtuple like ConnectionParams or a better name. The reason is that now for each added connection param, we would have to modify at two places

I agree, that makes sense.

Shouldn't we name it SessionParams though? These parameters will be passed through to the session constructor.

I believe Session is an internal abstraction and these params originate from Connection so better to name ConnectionParams. Additionally, i think this change might break a lot of things. Let's do it completely separately (let's log a JIRA ticket for now and take it up later)

https://databricks.atlassian.net/browse/PECOBLR-447?atlOrigin=eyJpIjoiNDM5OWZkZjlhOGRjNGU5MWI0NGU3NTNiMjIwMTU5ZjIiLCJwIjoiaiJ9

jayantsing-db · 2025-05-26T06:34:32Z

src/databricks/sql/session.py

+        precedence over the serverProtocolVersion defined in the OpenSessionResponse.
+        """


Suggested change

precedence over the serverProtocolVersion defined in the OpenSessionResponse.

"""

precedence over the serverProtocolVersion defined in the OpenSessionResponse.

"""

i think there is a line gap after a multi-line pydoc. @jprakash-db do we follow any python coding guidelines like for docstring: https://peps.python.org/pep-0257/

in databricks we follow this https://databricks.atlassian.net/wiki/spaces/UN/pages/3334538555/Python+Guidelines+go+py but this could be different for OSS repo.

We use the black formatter which follows the PEP-257 style.

got it. is there a linter? in the CI or do we have to run the linter manually?

Currently we have pylint but it needs to be run manually.

jayantsing-db · 2025-05-26T06:40:10Z

src/databricks/sql/client.py

+    @property
+    def protocol_version(self):
+        """Get the protocol version from the Session object"""
+        return self.session.protocol_version
+
+    @staticmethod
+    def get_protocol_version(openSessionResp):
+        """Delegate to Session class static method"""
+        return Session.get_protocol_version(openSessionResp)
+
+    @property
+    def open(self) -> bool:
+        """Return whether the connection is open by checking if the session is open."""
+        return self.session.is_open


nit: should we group together property and staticmethod? @jprakash-db any coding/lint guidelines OSS python driver follows?

There are no specific standard in python, because saying something as private etc has no meaning, we can access anything anytime. There are some general standards but nothing concrete

jayantsing-db · 2025-05-26T06:43:36Z

src/databricks/sql/session.py

+        except Exception as e:
+            logger.error(f"Attempt to close session raised a local exception: {e}")
+
+        self.is_open = False


when opening session the flag is set at the end which makes sense. for closing session call, should we be eager to unset the flag in the very beginning?

What if the session close fails? Shouldn't the session remain open?

if the client code has called on the Session class to close the session, then the client assumes that method will close the session. I think unsetting the flag right away makes more sense then. However, an interesting question is do we use this flag internally in Session class to make unsetting meaningful (i.e., when flag is false, do we give null or throw exception when getting session handle?)

Currently there does not seem to exist such a dependency, but I'm still not clear on this.

If the close() call raises an exception isn't the client expected to retry?

https://databricks.atlassian.net/browse/PECOBLR-448?atlOrigin=eyJpIjoiMzJmZTZmMmU4ZDcyNGU5NWJiMmZlZjEwZDVjMDU3NTIiLCJwIjoiaiJ9

tests/e2e/test_driver.py

tests/unit/test_client.py

tests/unit/test_session.py

github-actions · 2025-05-26T08:24:22Z

Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase (git rebase -i main).

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

…sion

jayantsing-db

thanks!

varun-edachali-dbx and others added 23 commits May 24, 2025 02:25

decouple session class from existing Connection

f97c81d

ensure maintenance of current APIs of Connection while delegating responsibility Signed-off-by: varun-edachali-dbx <[email protected]>

add open property to Connection to ensure maintenance of existing API

fe0af87

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

update unit tests to address ThriftBackend through session instead of…

18f8f67

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

chore: move session specific tests from test_client to test_session

fd8decb

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

formatting (black)

1a92b77

as in CONTRIBUTING.md Signed-off-by: varun-edachali-dbx <[email protected]>

use connection open property instead of long chain through session

1b9a50a

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

trigger integration workflow

0bf2794

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

Revert "fix: de-complicate earlier connection open logic"

63b10c3

This reverts commit d6b1b19. Signed-off-by: varun-edachali-dbx <[email protected]>

[empty commit] attempt to trigger ci e2e workflow

f2b3fd5

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

Update CODEOWNERS (#562)

53f16ab

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

Revert "Merge remote-tracking branch 'upstream/sea-migration' into de…

923bbb6

…couple-session" This reverts commit dbb2ec5, reversing changes made to 7192f11. Signed-off-by: varun-edachali-dbx <[email protected]>

Reapply "Merge remote-tracking branch 'upstream/sea-migration' into d…

8df8c33

…ecouple-session" This reverts commit bdb8381. Signed-off-by: varun-edachali-dbx <[email protected]>

fix: separate session opening logic from instantiation

bcf5994

ensures correctness of self.session.open call in Connection Signed-off-by: varun-edachali-dbx <[email protected]>

fix: use is_open attribute to denote session availability

500dd0b

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

fix: access thrift backend through session

510b454

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

chore: use get_handle() instead of private session attribute in client

634faa9

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

formatting (black)

a32862b

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

Merge remote-tracking branch 'upstream/sea-migration' into decouple-s…

88b728d

…ession

Merge remote-tracking branch 'upstream/sea-migration' into decouple-s…

ed04584

…ession

varun-edachali-dbx requested review from deeksha-db, samikshya-db, jprakash-db, jackyhu-db, madhav-db, gopalldb and jayantsing-db as code owners May 26, 2025 05:22

varun-edachali-dbx requested review from vikrantpuppala and shivam2680 as code owners May 26, 2025 05:22

varun-edachali-dbx temporarily deployed to azure-prod May 26, 2025 05:22 — with GitHub Actions Inactive

varun-edachali-dbx mentioned this pull request May 26, 2025

Separate Session related functionality from Connection class #567

Closed

5 tasks

jayantsing-db requested changes May 26, 2025

View reviewed changes

jayantsing-db reviewed May 26, 2025

View reviewed changes

tests/e2e/test_driver.py Show resolved Hide resolved

tests/unit/test_client.py Outdated Show resolved Hide resolved

tests/unit/test_client.py Show resolved Hide resolved

jayantsing-db reviewed May 26, 2025

View reviewed changes

tests/unit/test_session.py Show resolved Hide resolved

varun-edachali-dbx temporarily deployed to azure-prod May 26, 2025 08:24 — with GitHub Actions Inactive

fix: remove accidentally removed assertions

ff842d7

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

varun-edachali-dbx force-pushed the decouple-session branch from 99ee979 to ff842d7 Compare May 26, 2025 08:31

varun-edachali-dbx temporarily deployed to azure-prod May 26, 2025 08:31 — with GitHub Actions Inactive

varun-edachali-dbx requested a review from jayantsing-db May 26, 2025 12:45

Merge remote-tracking branch 'origin/sea-migration' into decouple-ses…

9190a33

…sion

varun-edachali-dbx temporarily deployed to azure-prod May 28, 2025 10:25 — with GitHub Actions Inactive

jayantsing-db approved these changes May 28, 2025

View reviewed changes

varun-edachali-dbx merged commit 5bf5d4c into sea-migration May 28, 2025
23 checks passed

		precedence over the serverProtocolVersion defined in the OpenSessionResponse.
		"""

Separate Session related functionality from Connection class #571

Separate Session related functionality from Connection class #571

Uh oh!

Conversation

varun-edachali-dbx commented May 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What type of PR is this?

Description

How is this tested?

Related Tickets & Documents

Uh oh!

jayantsing-db left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

varun-edachali-dbx May 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented May 26, 2025

Uh oh!

jayantsing-db left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

varun-edachali-dbx commented May 26, 2025 •

edited

Loading

varun-edachali-dbx May 26, 2025 •

edited

Loading