Skip to content

Updated the DatabricksRM class to use Databricks service principals with REST API path #8327

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 18 commits into
base: main
Choose a base branch
from

Conversation

willsmithDB
Copy link
Contributor

@willsmithDB willsmithDB commented Jun 4, 2025

This PR is to add service principal support with the Databricks REST API path for the databricks retriever in DSPy. This is mainly the path for work outside of Databricks environments (i.e VS Code with Databricks Connect) where the Databricks SDK is not automatically installed or if the REST API is preferred.

  • Updated logic for the _databricks_sdk_installed flag to not fail catastrophically when working locally without the package installed.
  • Updated the base class to include the optional parameters.
  • Updated print statements for which auth method is invoked.
  • Modified parameter checks to not incorrectly trigger an exception when service principal parameters exist, but SDK is not installed and the personal access token / endpoint are missing.
  • Updated docstrings.
  • Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id.
  • Updated logic on the _query_via_databricks_sdk to use the SP credentials if they exist otherwise will fallback to PAT.

Note: I synced the fork with the main branch which added a number of commits.

Screenshots of REST API with SP tests:

image Screenshot 2025-06-04 at 10 06 33 AM Screenshot 2025-06-04 at 10 07 49 AM image

willsmithDB and others added 18 commits May 27, 2025 14:31
- Updated the base class to include the optional parameters
- Updated logic on the _query_via_databricks_sdk to use the SP credentials if they exist otherwise will fallback to PAT.
…ity' into dbx_service_principal_functionality

# Conflicts:
#	dspy/retrieve/databricks_rm.py
Updated print statements for which auth method is invoked.
Updated docstrings.
Created helper function: _get_oauth_token
Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
…ity' into dbx_service_principal_functionality

# Conflicts:
#	dspy/retrieve/databricks_rm.py
Updated print statements for which auth method is invoked.
Updated docstrings.
Created helper function: _get_oauth_token
Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
Updated print statements for which auth method is invoked.
Updated docstrings.
Created helper function: _get_oauth_token
Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
Updated print statements for which auth method is invoked.
Updated docstrings.
Created helper function: _get_oauth_token
Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
…ity' into dbx_service_principal_functionality
Updated print statements for which auth method is invoked.
Updated docstrings.
Created helper function: _get_oauth_token
Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
…ity' into dbx_service_principal_functionality
…ality

Updated the DatabricksRM class to use Databricks service principals.
@willsmithDB willsmithDB changed the title Updated the DatabricksRM class to use Databricks service principals. Updated the DatabricksRM class to use Databricks service principals with REST API path. Jun 4, 2025
@willsmithDB willsmithDB changed the title Updated the DatabricksRM class to use Databricks service principals with REST API path. Updated the DatabricksRM class to use Databricks service principals with REST API path Jun 4, 2025
Copy link
Collaborator

@chenmoneygithub chenmoneygithub left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR!

I feel this is going too complex with supporting all Databricks auth methods. We should move this implementation into dastabricks-ai-bridge IMO. I will discuss with the team and make a decision.

@willsmithDB
Copy link
Contributor Author

willsmithDB commented Jun 4, 2025

@chenmoneygithub ,

Pardon me for the extra PR. This was supposed to be included with my other PR to add the same SP support for both methods that were already present (SDK and REST API). It's the same mechanism, but inadvertently developed for the SDK before the API when they should've been added together if that makes sense.

Thank you for your thoughts on this!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants