Generalise Backend Layer #604

varun-edachali-dbx · 2025-06-18T03:50:36Z

What type of PR is this?

Refactor

Description

This PR refactors the codebase to aid the eventual introduction of a new backend client. It is composed of the following three PRs:

Separate the Session functionality from the Connection class to further abstract backend implementation details from the connection class: Separate Session related functionality from Connection class #571
Introduce a general backend interface (DatabricksClient): Introduce Backend Interface (DatabricksClient) #573
- Introduce the DatabricksClient interface and make the existing ThriftBackend implement it as ThriftDatabricksClient.
- Introduce SessionId and CommandId as consistent adapters for backend layers to represent sessions and commands instead of relying on Thrift specific types.
Introduce a general ResultSet interface: Implement ResultSet Abstraction (backend interfaces for fetch phase) #574
- Create a ResultSet interface and make the existing ResultSet interface implement it as ThriftResultSet.
- Make the concrete Thrift backend return a ThriftResultSet instead of ExecuteResponse from execution relevant commands.
- Generalise the return types of methods in the backend interface to prevent alignment with Thrift (introduce CommandState to prevent using Thrift's status types).

How is this tested?

Unit tests
E2E Tests
Manually
N/A

Related Tickets & Documents

Design Doc
PECOBLR-569

* 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> * add open property to Connection to ensure maintenance of existing API Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * update unit tests to address ThriftBackend through session instead of through Connection Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * chore: move session specific tests from test_client to test_session Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) as in CONTRIBUTING.md Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * use connection open property instead of long chain through session Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * trigger integration workflow Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix: ensure open attribute of Connection never fails 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 <varun.edachali@databricks.com> * fix: de-complicate earlier connection open logic 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 <varun.edachali@databricks.com> * Revert "fix: de-complicate earlier connection open logic" This reverts commit d6b1b19. Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * [empty commit] attempt to trigger ci e2e workflow Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * Update CODEOWNERS (#562) new codeowners Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * Enhance Cursor close handling and context manager exception management 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 <varun.edachali@databricks.com> * PECOBLR-86 improve logging on python driver (#556) * PECOBLR-86 Improve logging for debug level Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com> * PECOBLR-86 Improve logging for debug level Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com> * fixed format Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com> * used lazy logging Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com> * changed debug to error logs Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com> * used lazy logging Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com> --------- Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com> Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * Revert "Merge remote-tracking branch 'upstream/sea-migration' into decouple-session" This reverts commit dbb2ec5, reversing changes made to 7192f11. Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * Reapply "Merge remote-tracking branch 'upstream/sea-migration' into decouple-session" This reverts commit bdb8381. Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * 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> * fix: use is_open attribute to denote session availability Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix: access thrift backend through session Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * chore: use get_handle() instead of private session attribute in client Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix: remove accidentally removed assertions Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> --------- Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com> Co-authored-by: Jothi Prakash <jothi.prakash@databricks.com> Co-authored-by: Madhav Sainanee <madhav.sainanee@databricks.com> Co-authored-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

NOTE: the `test_complex_types` e2e test was not working at the time of this merge. The test must be triggered when the test is back up and running as intended. * remove excess logs, assertions, instantiations large merge artifacts Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) + remove excess log (merge artifact) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix typing Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove un-necessary check Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove un-necessary replace call Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * introduce __str__ methods for CommandId and SessionId Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * docstrings for DatabricksClient interface Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * stronger typing of Cursor and ExecuteResponse Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove utility functions from backend interface, fix circular import Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * rename info to properties Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * newline for cleanliness Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix circular import Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * to_hex_id -> get_hex_id Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * better comment on protocol version getter Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * move guid to hex id to new utils module Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * move staging allowed local path to connection props Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add strong return type for execute_command Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * skip auth, error handling in databricksclient interface Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * chore: docstring + line width Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * get_id -> get_guid Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * chore: docstring Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix: to_hex_id -> to_hex_guid Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> --------- Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

…574) * ensure backend client returns a ResultSet type in backend tests Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * newline for cleanliness Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix circular import Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * to_hex_id -> get_hex_id Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * better comment on protocol version getter Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * stricter typing for cursor Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * correct typing Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * correct tests and merge artifacts Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove accidentally modified workflow files remnants of old merge Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * chore: remove accidentally modified workflow files Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add back accidentally removed docstrings Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * clean up docstrings Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * log hex Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove unnecessary _replace call Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add __str__ for CommandId Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * take TOpenSessionResp in get_protocol_version to maintain existing interface Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * active_op_handle -> active_mmand_id Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * ensure None returned for close_command Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * account for ResultSet return in new pydocs Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * pydoc for types Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * move common state to ResultSet aprent Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * stronger typing in resultSet behaviour Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove redundant patch in test Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add has_been_closed_server_side assertion Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove redundancies in tests Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * more robust close check Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * use normalised state in e2e test Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * simplify corrected test Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add line gaps after multi-line pydocs for consistency Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * use normalised CommandState type in ExecuteResponse Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> --------- Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

jayantsing-db

Great work, i think the implementation of sql connector makes more sense now and much more open to extension. Some comments inline.

Thanks very much @jprakash-db for the review.

src/databricks/sql/backend/databricks_client.py

src/databricks/sql/client.py

src/databricks/sql/result_set.py

src/databricks/sql/utils.py

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

vikrantpuppala

telemetry changes lgtm

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx and others added 4 commits June 18, 2025 03:47

remove un-necessary initialisation assertions

450b80d

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx had a problem deploying to azure-prod June 18, 2025 03:50 — with GitHub Actions Failure

remove un-necessary line break s

a926f02

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx had a problem deploying to azure-prod June 18, 2025 03:56 — with GitHub Actions Failure

varun-edachali-dbx added 2 commits June 18, 2025 03:57

more un-necessary line breaks

55ad001

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

constrain diff of test_closing_connection_closes_commands

fa15730

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod June 18, 2025 04:13 — with GitHub Actions Inactive

reduce diff of test_closing_connection_closes_commands

019c7fb

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx had a problem deploying to azure-prod June 18, 2025 04:15 — with GitHub Actions Failure

use pytest-like assertions for test_closing_connection_closes_commands

726abe7

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod June 18, 2025 04:18 — with GitHub Actions Inactive

ensure command_id is not None

bf6d41c

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod June 18, 2025 05:46 — with GitHub Actions Inactive

line breaks after multi-line pyfocs

5afa733

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod June 18, 2025 05:57 — with GitHub Actions Inactive

ensure non null operationHandle for commandId creation

e3dfd36

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx had a problem deploying to azure-prod June 18, 2025 06:05 — with GitHub Actions Failure

use command_id methods instead of explicit guid_to_hex_id conversion

63360b3

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod June 18, 2025 06:11 — with GitHub Actions Inactive

remove un-necessary artifacts in test_session, add back assertion

13ffb8d

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx had a problem deploying to azure-prod June 18, 2025 06:22 — with GitHub Actions Failure

varun-edachali-dbx marked this pull request as ready for review June 18, 2025 06:28

varun-edachali-dbx requested review from deeksha-db, samikshya-db, jprakash-db, jackyhu-db and madhav-db as code owners June 18, 2025 06:28

varun-edachali-dbx temporarily deployed to azure-prod July 4, 2025 04:49 — with GitHub Actions Inactive

vikrantpuppala removed their request for review July 8, 2025 03:37

varun-edachali-dbx removed request for deeksha-db, samikshya-db, jackyhu-db, madhav-db, gopalldb and shivam2680 July 10, 2025 05:22

Merge branch 'main' into backend-refactors

45f5c26

varun-edachali-dbx had a problem deploying to azure-prod July 10, 2025 05:46 — with GitHub Actions Failure

formatting (black)

2c9368a

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx had a problem deploying to azure-prod July 10, 2025 05:47 — with GitHub Actions Failure

remove repetition from Session.__init__

9b1b1f5

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod July 10, 2025 05:58 — with GitHub Actions Inactive

jayantsing-db approved these changes Jul 13, 2025

View reviewed changes

varun-edachali-dbx added 3 commits July 14, 2025 07:14

mention that if catalog / schema name is None, we fetch across all

d72be7f

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

mention fetching across all tables if null table name

0e1f410

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

remove lazy import of ThriftResultSet

bbef428

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx requested review from vikrantpuppala and jprakash-db July 14, 2025 02:04

Merge branch 'main' into backend-refactors

deb93a0

varun-edachali-dbx temporarily deployed to azure-prod July 14, 2025 02:28 — with GitHub Actions Inactive

remove unused import

e1a8c43

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod July 14, 2025 02:37 — with GitHub Actions Inactive

vikrantpuppala approved these changes Jul 14, 2025

View reviewed changes

better docstrings

8d06f5e

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod July 15, 2025 04:37 — with GitHub Actions Inactive

clarified role of cursor in docstring

044b3e5

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

varun-edachali-dbx temporarily deployed to azure-prod July 15, 2025 06:00 — with GitHub Actions Inactive

varun-edachali-dbx merged commit ba1eab3 into main Jul 15, 2025
23 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Generalise Backend Layer #604

Generalise Backend Layer #604

Uh oh!

varun-edachali-dbx commented Jun 18, 2025 •

edited

Loading

Uh oh!

jayantsing-db left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vikrantpuppala left a comment

Uh oh!

Uh oh!

Uh oh!

Generalise Backend Layer #604

Generalise Backend Layer #604

Uh oh!

Conversation

varun-edachali-dbx commented Jun 18, 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vikrantpuppala left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

varun-edachali-dbx commented Jun 18, 2025 •

edited

Loading