kestra-io / kestra-io/plugin-databricks
Add Lakebase Support Plugin
- Dominant language
- Java
- Stars
- 3
- Forks
- 10
- Avg merge
- 2d 59m
- Merged PRs (30d)
- 9
Description
## Summary
Lakebase is Databricks' managed Postgres (OLTP) offering — it speaks the standard Postgres wire protocol but authenticates with short-lived OAuth tokens rather than static passwords. This issue adds native Lakebase connectivity to `plugin-databricks` so flows can run transactional queries against a Lakebase database without hand-rolling token refresh logic in a script task.
## Motivation
Today, users who want to reach a Lakebase instance from Kestra have to either write a custom script task that calls the Databricks SDK to mint a token and then talks to Postgres manually, or misuse `plugin-jdbc-postgres` with a static token pasted into the `password` field — which silently breaks after ~60 minutes when the token expires. Neither is a good experience for an OLTP-style workload where a flow may run frequent, short queries against Lakebase over time. A dedicated Lakebase task/connection that mints a fresh credential per execution removes that footgun entirely and gives Databricks-centric flows a first-class way to read/write their OLTP layer alongside their existing cluster/job/SQL steps.
## Context
Part of the Databricks platform coverage EPIC: https://github.com/kestra-io/kestra-ee/issues/9384
We evaluated reusing `plugin-jdbc-postgres` as-is (see `PostgresConnectionInterface`/`PostgresService` in that repo): it already handles the JDBC URL, SSL modes, and client cert/key plumbing that Lakebase also needs, but it only exposes static `username`/`password` fields (`JdbcConnectionInterface`) with no mechanism to mint or refresh a short-lived OAuth token before each connection. That gap — token lifecycle, not SQL execution — is exactly what this issue needs to close, most naturally as a Lakebase-specific connection interface inside `plugin-databricks` that generates the credential via the Databricks SDK and then delegates to the same JDBC/SSL machinery `plugin-jdbc-postgres` already has.
## API Reference
- **Official docs**: https://docs.databricks.com/aws/en/oltp/projects/external-apps-connect?language=Java
- **Authentication**: OAuth machine-to-machine (M2M) — a Databricks service principal with "Workspace access" enabled calls `WorkspaceClient.postgres().generateDatabaseCredential(new GenerateDatabaseCredentialRequest().setEndpoint(endpointName))`, and the returned `DatabaseCredential.getToken()` is passed as the JDBC password. The username is the service principal's client ID (a UUID). Tokens expire after 60 minutes — connections/pools must recycle before then (Databricks recommends a 45-minute max lifetime if pooling).
- **Base URL pattern**: standard Postgres wire protocol — `jdbc:postgresql://:/` — Lakebase requires SSL.
- **SDK / client library**: Databricks SDK for Java, `com.databricks:databricks-sdk-java` (latest stable: `0.130.0`) — `WorkspaceClient.postgres()` exposes the `generateDatabaseCredential` call needed for token minting.
## Gradle Dependencies
Add to `build.gradle`:
```groovy
// Databricks SDK for Java — used to mint short-lived Lakebase database credentials
implementation "com.databricks:databricks-sdk-java:0.130.0"
// PostgreSQL JDBC driver — Lakebase speaks the standard Postgres wire protocol
implementation "org.postgresql:postgresql:42.7.11"
```
> Use the latest stable version available on Maven Central for both.
>
> **Kestra framework inclusions (do not add as dependencies):** Kestra's internal HTTP client (`io.kestra.core.http.client`) and Jackson serializers are provided by the framework. HikariCP-style connection pooling is not required here — each Kestra task run opens its own connection for the duration of the run, so a fresh token per execution (rather than a long-lived pool) is the correct pattern; only add pooling if a future task genuinely holds a connection open across a long-running batch.
## Plugin Structure
- **Repository**: `plugin-databricks` (existing)
- **Namespace**: `io.kestra.plugin.databricks.lakebase`
- **Sub-plugins**: `lakebase`
- **Categories**: `DATA`
> **Task class naming**: task class names must not repeat the plugin or package name as a prefix or suffix. Use concise action names: `Query`, `Batch`, `Trigger` — not `LakebaseQuery`, `QueryLakebase`, etc.
## Suggested Tasks
1. Implement a `LakebaseConnectionInterface` (extending `PostgresConnectionInterface`-style SSL/JDBC properties) that adds `endpoint`, `workspaceHost`, `clientId`, `clientSecret` properties and overrides connection-property construction to call `generateDatabaseCredential` and use the returned token as the password on every connection attempt
2. `Query` — execute a single SQL statement against Lakebase and return rows (`fetchType`: `FETCH_ONE` / `FETCH` / `STORE`)
3. `Batch` — execute a batch of parameterized statements (bulk insert/update)
4. `Trigger` — poll a table/query for new or changed rows since the last execution (`PollingTriggerInterface`), mirroring `plugin-jdbc-postgres`'s `Trigger`
5. Write unit + integration tests (Testcontainers if a Lakebase-compatible Postgres image can be used for local testing; otherwise mock `generateDatabaseCredential` and test the token-refresh path directly)
6. Add `package-info.java` with `@PluginSubGroup(category = PluginSubGroup.PluginCategory.DATA)`
7. Add `metadata/lakebase.yaml` and plugin icon SVG (reuse the existing `plugin-icon.svg` if no dedicated Lakebase icon exists)
8. Add YAML examples and update the plugin's how-to doc (`io.kestra.plugin.databricks.md`) with a `## Lakebase` authentication subsection
## YAML Examples
### Example 1 — Run a query against a Lakebase database
```yaml
id: lakebase_query
namespace: company.team
tasks:
- id: query_orders
type: io.kestra.plugin.databricks.lakebase.Query
workspaceHost: "{{ secret('DATABRICKS_HOST') }}"
clientId: "{{ secret('DATABRICKS_CLIENT_ID') }}"
clientSecret: "{{ secret('DATABRICKS_CLIENT_SECRET') }}"
endpoint: "{{ secret('LAKEBASE_ENDPOINT_NAME') }}"
database: "orders_db"
sql: "SELECT id, status, updated_at FROM orders WHERE status = 'pending'"
fetchType: FETCH
```
### Example 2 — Batch-insert rows produced by an upstream task
```yaml
id: lakebase_batch_insert
namespace: company.team
inputs:
- id: rows
type: ARRAY
itemType: STRING
tasks:
- id: insert_rows
type: io.kestra.plugin.databricks.lakebase.Batch
workspaceHost: "{{ secret('DATABRICKS_HOST') }}"
clientId: "{{ secret('DATABRICKS_CLIENT_ID') }}"
clientSecret: "{{ secret('DATABRICKS_CLIENT_SECRET') }}"
endpoint: "{{ secret('LAKEBASE_ENDPOINT_NAME') }}"
database: "orders_db"
sql: "INSERT INTO orders_audit (payload) VALUES (?)"
parameterGroups:
- parameters: "{{ inputs.rows }}"
```
### Example 3 — React to newly inserted rows
```yaml
id: lakebase_new_rows_trigger
namespace: company.team
triggers:
- id: on_new_order
type: io.kestra.plugin.databricks.lakebase.Trigger
workspaceHost: "{{ secret('DATABRICKS_HOST') }}"
clientId: "{{ secret('DATABRICKS_CLIENT_ID') }}"
clientSecret: "{{ secret('DATABRICKS_CLIENT_SECRET') }}"
endpoint: "{{ secret('LAKEBASE_ENDPOINT_NAME') }}"
database: "orders_db"
sql: "SELECT * FROM orders WHERE status = 'pending'"
interval: PT1M
tasks:
- id: handle_new_order
type: io.kestra.plugin.core.log.Log
message: "New pending order: {{ trigger.rows }}"
```
## Acceptance Criteria
### Functional
- [ ] `LakebaseConnectionInterface` mints a fresh OAuth token per connection via the Databricks SDK
- [ ] `Query`, `Batch` tasks implemented
- [ ] `Trigger` polling trigger implemented
- [ ] Unit + integration tests pass (`./gradlew test`)
- [ ] Build passes (`./gradlew build`)
### Kestra Plugin Coding Standards
- [ ] HTTP calls (if any outside the Databricks SDK) use Kestra's internal HTTP client — no OkHttp/Apache HttpClient
- [ ] All new properties use `Property`
- [ ] `clientSecret` annotated with `@PluginProperty(secret = true)`
- [ ] Every property and output has a `@Schema` annotation
- [ ] Task classes carry the five mandatory Lombok annotations (`@SuperBuilder`, `@ToString`, `@EqualsAndHashCode`, `@Getter`, `@NoArgsConstructor`)
- [ ] Logging via `runContext.logger()` only
- [ ] JSON serialization uses Jackson mappers from `io.kestra.core.serializers`
- [ ] All `Property` fields support Kestra expression language (template rendering)
### Documentation & Structure
- [ ] `@Plugin(examples = ...)` entries each set `full = true` with a complete runnable flow (id + namespace + tasks/triggers)
- [ ] Sensitive values in examples use `{{ secret('SECRET_NAME') }}`
- [ ] `package-info.java` with `@PluginSubGroup(category = PluginSubGroup.PluginCategory.DATA)`
- [ ] `metadata/lakebase.yaml` and plugin icon SVG present
- [ ] `io.kestra.plugin.databricks.md` updated with the Lakebase authentication flow
---
*[View as Artifact](https://claude.ai/code/artifact/eaaad748-e885-42a5-8dcf-e4fc2973498a)*
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with build.gradle and the PostgresConnectionInterface/PostgresService patterns referenced in plugin-jdbc-postgres. Then review the plugin-databricks structure, especially package-info.java, metadata/lakebase.yaml, and the existing how-to documentation. Done means Query, Batch, and Trigger support fresh OAuth credentials, tests and builds pass, and the examples and documentation are included.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- java, postgresql
- Domain
- backend, databases, documentation, testing-qa
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100