[Data] Add credential provider abstraction for Databricks UC datasource

[ankur-anyscale]

Introduce a credential provider pattern for Databricks authentication, enabling custom credential sources while maintaining backward compatibility.

Changes:

• Add DatabricksCredentialProvider base class with StaticCredentialProvider and EnvironmentCredentialProvider implementations
• Add credential_provider parameter to DatabricksUCDatasource and read_databricks_tables()
• Add UnityCatalogConnector class for reading Unity Catalog tables directly (supports Delta/Parquet formats with AWS, Azure, and GCP credential handoff)
• Add retry on 401 with credential invalidation via shared request_with_401_retry() helper
• Centralize common code (build_headers, request_with_401_retry) in databricks_credentials.py
• Move Databricks tests to dedicated test files with shared test utilities

The credential provider abstraction allows users to implement custom credential sources that support token refresh and other authentication patterns beyond static tokens.

Backward compatibility: Existing code using the DATABRICKS_TOKEN and DATABRICKS_HOST environment variables continues to work unchanged.

[gemini-code-assist]

Code Review

This pull request introduces a new, extensible Databricks credential provider system for Ray Data. It defines an abstract DatabricksCredentialProvider base class and provides concrete implementations for static tokens (StaticCredentialProvider) and environment variables (EnvironmentCredentialProvider), the latter including logic to detect the Databricks host from the runtime environment. The DatabricksUCDatasource was refactored to utilize this new provider, dynamically fetching authentication tokens for each request and implementing a single-retry mechanism for 401 Unauthorized responses during both statement polling and chunk resolution. The read_databricks_tables API was updated to accept an optional credential_provider parameter, replacing its previous direct environment variable parsing. Corresponding unit tests were added for the new credential providers and their integration, including tests for the 401 retry logic. Review comments highlighted areas for improvement, such as making the EnvironmentCredentialProvider's host resolution error message more precise, evaluating the portability of the _detect_databricks_host method's reliance on IPython, and enhancing the 401 retry logic in both the polling and chunk resolution loops with a maximum retry limit or backoff strategy to prevent potential infinite loops. Additionally, a comment noted the importance of ensuring the DatabricksCredentialProvider instance is serializable if it holds state for remote workers.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 005cdf5265

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[goutamvenkat-anyscale]

left a few comments

[goutamvenkat-anyscale]

Lets make these consts

[ankur-anyscale]

Done.

[goutamvenkat-anyscale]

Why are we bypassing all exceptions?

[ankur-anyscale]

Logged the exception here. If the data runtime is not available and host is also not set, then it will be handled by the current flow.

[goutamvenkat-anyscale]

Make this into a reusable function

[ankur-anyscale]

Done.

[goutamvenkat-anyscale]

Nit: dict[str, str]

[ankur-anyscale]

Done.

[goutamvenkat-anyscale]

Why is this being nuked?

[ankur-anyscale]

This was just moved to other file. I have updated the tests as well as per current guidelines as well as your comments.

[goutamvenkat-anyscale]

use rows_same

[ankur-anyscale]

Done

[goutamvenkat-anyscale]

Use rows_same

[ankur-anyscale]

Done.

[goutamvenkat-anyscale]

/gemini summary

[gemini-code-assist]

Summary of Changes

This pull request significantly enhances the Databricks Unity Catalog datasource by introducing a robust and extensible credential provider abstraction. This change improves authentication flexibility, allowing users to define custom credential sources and enabling automatic token refresh mechanisms. It also adds resilience to transient authentication failures by implementing a 401 retry mechanism with credential invalidation, ensuring smoother data operations.

Highlights

• Credential Provider Abstraction: Introduced a new DatabricksCredentialProvider abstract base class and concrete implementations (StaticCredentialProvider, EnvironmentCredentialProvider) to standardize and extend Databricks authentication methods.
• Enhanced Authentication Flexibility: The DatabricksUCDatasource and read_databricks_tables() now accept a credential_provider parameter, allowing users to implement custom authentication logic, including token refresh mechanisms.
• Robust 401 Retry Mechanism: Implemented automatic retry logic for 401 Unauthorized HTTP responses across all Databricks API calls (initial POST, polling GETs, and chunk fetching GETs), with credential invalidation to prompt token refresh.
• Improved Test Structure: Databricks-related tests have been refactored and moved into dedicated test files (test_databricks_credentials.py, test_databricks_uc_datasource.py) for better organization and maintainability.
• Backward Compatibility: Existing authentication methods using the token parameter or DATABRICKS_TOKEN environment variable remain functional and are handled by the default EnvironmentCredentialProvider.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/ray/data/_internal/datasource/databricks_credentials.py
  • Added new file to define Databricks credential provider interfaces and implementations.
  • Defined DatabricksCredentialProvider as an abstract base class with get_token, get_host, and invalidate methods.
  • Implemented StaticCredentialProvider for handling fixed authentication tokens and hosts.
  • Implemented EnvironmentCredentialProvider to retrieve credentials from environment variables, including Databricks runtime host detection and token re-reading on invalidate().
  • Added resolve_credential_provider utility function to determine the appropriate credential provider based on input.
• python/ray/data/_internal/datasource/databricks_uc_datasource.py
  • Imported DatabricksCredentialProvider for type hinting and usage.
  • Modified DatabricksUCDatasource.__init__ to accept a credential_provider instance instead of separate host and token arguments.
  • Introduced _build_headers and _request_with_401_retry helper functions to centralize header construction and implement 401 retry logic.
  • Updated all internal requests.post and requests.get calls to utilize the new _request_with_401_retry function, providing resilience against expired tokens.
  • Ensured the credential_provider is correctly captured for serialization when creating read tasks.
• python/ray/data/read_api.py
  • Removed direct os imports for environment variable access related to Databricks credentials.
  • Imported DatabricksCredentialProvider and resolve_credential_provider.
  • Added an optional credential_provider parameter to the read_databricks_tables function.
  • Updated the read_databricks_tables function to use resolve_credential_provider to obtain the authentication provider.
  • Removed previous direct environment variable parsing for Databricks host and token within read_databricks_tables.
  • Updated documentation for read_databricks_tables to include usage examples for custom credential providers.
• python/ray/data/tests/datasource/test_databricks_credentials.py
  • Added new file containing comprehensive unit tests for the DatabricksCredentialProvider abstract base class and its concrete implementations (StaticCredentialProvider, EnvironmentCredentialProvider).
  • Includes tests for proper initialization, credential retrieval, host detection, token refresh on invalidate(), and serialization capabilities of the providers.
• python/ray/data/tests/datasource/test_databricks_uc_datasource.py
  • Added new file containing integration tests for the DatabricksUCDatasource.
  • Tests reading data using both table names and SQL queries with various parallelism settings.
  • Includes specific tests for credential provider integration, ensuring fresh tokens are fetched for each request.
  • Verifies the 401 retry mechanism functions correctly during initial POST requests, polling, and chunk fetching.
  • Tests handling of empty query results.
• python/ray/data/tests/datasource/test_sql.py
  • Removed Databricks Unity Catalog datasource tests, which have been migrated to the new test_databricks_uc_datasource.py file.

Activity

• gemini-code-assist[bot] raised concerns regarding the precision of host resolution error messages, the dependency on IPython for host detection, and the lack of retry limits/backoff for 401 errors in polling and chunk fetching. It also questioned the serializability of DatabricksCredentialProvider.
• ankur-anyscale (author) responded by stating that existing logic was moved without change, added unit tests to confirm serializability, and implemented token refresh within EnvironmentCredentialProvider's invalidate method.
• chatgpt-codex-connector[bot] suggested refreshing environment-backed tokens on invalidate(), which was subsequently addressed by the author.
• cursor[bot] identified several issues, including MockResponse in tests missing the status_code attribute, the initial POST request lacking 401 retry logic, the presence of unused helper functions/fixtures, redundant credential provider classes in tests, and duplicated test setup code.
• ankur-anyscale addressed the MockResponse issue, implemented 401 retry for the initial POST, and refactored test code to remove unused elements and reduce duplication.
• goutamvenkat-anyscale provided feedback suggesting the use of constants, creating reusable functions, applying specific type hints (dict[str, str]), utilizing rows_same for DataFrame comparisons, and parameterizing tests. These suggestions were addressed by the author.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.