apache / apache/datafusion-ballista

[DISCUSSION] Restore Flight SQL support as a pluggable query frontend

Open
#2,298 2 comments 4 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
2.1k
Forks
320
Avg merge
1d 22h
Merged PRs (30d)
66

Description

> **Disclaimer:** this issue was drafted with an LLM (Claude Code) at my request. The code, git, and issue references were verified against `main` at 5e54503, but the design proposal is a starting point for discussion, not a settled plan.

## Is your feature request related to a problem or challenge? Please describe what you are trying to do.

Ballista has no way for a generic SQL client to connect to it. Every client must be a Rust program that builds a DataFusion `LogicalPlan`, serializes it, and speaks `SchedulerGrpc`. That rules out JDBC/ODBC tooling, BI tools, DBeaver, Python (`adbc`/`pyarrow`), and anything else that expects to send SQL text and get Arrow back.

Flight SQL used to fill that gap, and was removed in 46.0.0 (#1227, #1228). It's worth being precise about *why*, because it wasn't "Flight SQL is a bad fit for Ballista":

- Nobody was maintaining it, and it had accumulated unaddressed bugs (#1012, #941, #839, #756).
- It was substantially incomplete. Of ~40 `FlightSqlService` methods in `ballista/scheduler/src/flight_sql.rs`, well over half were `Status::unimplemented`, including `CommandGetSqlInfo` — which carried the comment `// TODO: implement for FlightSQL JDBC to work`. Catalog metadata (`GetCatalogs`, `GetDbSchemas`, `GetTables` as a real RPC, `GetXdbcTypeInfo`, keys) was stubbed, so driver-level introspection didn't work.
- It was structurally coupled to scheduler internals: it took a concrete `SchedulerServer`, reached into `server.state.config`, faked catalog responses by smuggling well-known strings (`"get_flight_info_tables"`) through the `job_id` field of a `FetchPartition` ticket, and carried its own inline Flight proxy.
- Authentication was `if user != "admin" || pass != "password"`, Basic-only.
- It had **zero tests** (`grep -c '#\[test\]' ` → 0).

So the removal was right, and reviving that file as-is would be wrong. But the capability is still valuable, and @avantgardnerio made the point directly in #2249: *"using this to restore FlightSQL could add a significant audience as well."* This issue is to design a new implementation properly.

## Why now

Several things have landed since 46.0.0 that make this materially easier than it was in 2023–2025:

1. **The scheduler has a real embedded Flight proxy** (#1351, `ballista/scheduler/src/flight_proxy_service.rs`), plus `FlightProxy::Local`/`External` advertisement in `GetJobStatusResult`. This fixes the old implementation's worst structural problem: it built `FlightEndpoint` locations out of per-executor `host:port` (and `make_local_fieps` hardcoded `127.0.0.1:50050` behind a `// TODO: use advertise host`), which is exactly why it broke behind NAT, Docker, Kubernetes, and load balancers (#1012, #1349). A new implementation can hand clients **one stable address** and let the scheduler fan out to executors.
2. **The scheduler already owns a per-session `SessionContext` built by a pluggable `SessionBuilder`** (`ballista/scheduler/src/state/session_manager.rs:48-85`). Server-side SQL planning needs a catalog, and this is an existing, documented extension point for supplying one — rather than the old UX of "re-register your tables with `CREATE EXTERNAL TABLE` on every new connection."
3. **A richer scheduler RPC surface**: `CreateUpdateSession`, `RemoveSession`, `ExecuteQueryPush` (streaming status, so no polling loop), `GetJobMetrics`, `CancelJob`, `CleanJobData`.
4. **`arrow-flight` 58.4 is already a workspace dependency with `flight-sql-experimental` enabled** (`Cargo.toml:41`), so this adds no new dependency — that feature flag is currently enabled for nothing.

## Relationship to #2249

This should be built **on top of** the pluggable-frontend layering @phillipleblanc proposed in #2249, not alongside it:

```text
Spark Connect / Ballista gRPC / Flight SQL / future protocols
|
mountable protocol frontend
|
plan decoder / relation translator
|
DataFusion LogicalPlan
|
transport-neutral QueryBackend
```

Flight SQL is arguably the *best first* consumer of that abstraction, and a useful forcing function for it:

- It is much smaller in scope than Spark Connect (no artifacts, no UDF upload, no Spark relation tree to translate) — SQL text in, Arrow out.
- It exercises every part of the contract anyway: session open/close, auth, plan decode, submit, status, cancel, result streaming, and metadata.
- It has an existing, external, standards-defined client population, so "does the abstraction actually work" gets tested by third-party drivers rather than by our own client.

If #2249's `QueryBackend`/`GrpcQueryFrontend` traits land first, this issue implements a `FlightSqlFrontend` against them. If Flight SQL work starts first, it should define its needs in those terms so the traits are shaped by two real consumers instead of one. Concretely, this frontend must not depend on `SchedulerServer` internals the way the old one did — that coupling is what made it unmaintainable.

## Describe the solution you'd like

### Placement

A new optional crate, `ballista/flight-sql`, exposed through a non-default `flight-sql` feature on the scheduler (and mountable by embedders into their own Tonic server per #2249). Rationale: it keeps the scheduler's default dependency surface unchanged, and it means the thing can be maintained — or deprecated — independently. That directly answers the maintainability objection in #1227.

### Core flow

`GetFlightInfo(CommandStatementQuery)`:
1. Resolve the Flight SQL session (from the handshake token) to a Ballista `session_id`; `CreateUpdateSession` if new.
2. Plan the SQL text server-side with that session's `SessionContext` → `LogicalPlan`.
3. Submit via the query backend (today: `ExecuteQuery`), obtaining a `job_id`.
4. Await completion via `ExecuteQueryPush`'s status stream rather than a poll loop.
5. Build `FlightEndpoint`s whose tickets are the existing `Action::FetchPartition` protobuf, and whose `location` is resolved through the **existing** `FlightProxy` logic — scheduler/LB by default, direct-to-executor only when explicitly configured.

`DoGet(ticket)` is then satisfied by the existing proxy path in `flight_proxy_service.rs`, which already decodes `Action::FetchPartition` and forwards to the owning executor. Ideally the frontend reuses that service rather than re-implementing a proxy inline, as the old code did.

### Open design decisions

These are the things I'd like input on; I don't think any of them should be settled unilaterally.

1. **Catalog / planning model.** SQL text has to be planned against *something*. Options: (a) session-scoped DDL only, as before — simple, but poor UX and no cross-connection persistence; (b) embedder-supplied catalog via `SessionBuilder` — my preference, reuses an existing extension point; (c) scheduler-side configured tables/object stores. Also: should `CommandStatementSubstraitPlan` be supported for plan-shaped clients, given `substrait` is already an optional scheduler feature?
2. **Long-running queries.** Blocking `GetFlightInfo` until the job completes (old behaviour) will hit client and proxy timeouts for TPC-H-scale queries. `PollFlightInfo` exists in `arrow-flight` 58 (the old impl stubbed it) and is designed for exactly this. Do we require it, and do the drivers we care about actually use it?
3. **Metadata surface.** What is the minimum viable set for real drivers? At minimum `CommandGetSqlInfo` and `GetXdbcTypeInfo`, which the old implementation never had — and without which the Arrow Flight JDBC driver doesn't work. Proposal: define the set by "the driver connects and introspects," not by "the trait method exists," and drive catalog answers from the session's real DataFusion catalog rather than hand-built `RecordBatch`es.
4. **Prepared statements.** Plan cache keyed by handle needs a TTL/eviction policy — the old `DashMap` only shed entries on explicit close, so any client that disconnected leaked. Parameter binding (`DoPutPreparedStatementQuery`, `parameter_schema`) was a `// TODO: parameters`; is bound-parameter support in scope for v1 or explicitly deferred?
5. **Auth.** Proposal: a pluggable authenticator trait with no built-in credentials, so embedders wire in their own, plus a clearly-labelled insecure dev default (or none at all). Reuse the existing `use_tls` and endpoint-customization hooks (#1400). Anything resembling `admin`/`password` should not come back.
6. **Lifecycle and cleanup.** Map `CancelFlightInfo`/`CancelQuery` → `CancelJob`, connection close → `RemoveSession`, and result cleanup → `CleanJobData`. The old implementation cleaned up none of this.
7. **Write path.** `CommandStatementUpdate` (DDL/DML/`INSERT INTO`) — in scope for v1 or deferred?
8. **Config naming.** `advertise_flight_sql_endpoint` currently has nothing to do with Flight SQL — it's the plain-Arrow-Flight result-proxy address, misnamed because it predates the Flight SQL removal (see #2281 and my comment on #2288). Either this frontend gives that name a legitimate meaning again, or the existing knob should be renamed and the Flight SQL frontend gets its own. Worth settling here rather than accumulating a third overlapping option.

### Suggested phasing

Each phase should be independently reviewable and independently useful:

1. `GetFlightInfo`/`DoGet` for `CommandStatementQuery` against a scheduler-configured catalog, endpoints resolved via the Flight proxy, with a Rust `FlightSqlServiceClient` integration test.
2. Session + handshake + pluggable auth; cancellation and cleanup wired to `CancelJob`/`RemoveSession`/`CleanJobData`.
3. Metadata RPCs sufficient for the Arrow Flight JDBC driver to connect and introspect, sourced from the real catalog.
4. Prepared statements with parameter binding; `PollFlightInfo` for long queries.
5. Docs (a replacement for the removed `docs/source/user-guide/flightsql.md`) and a `docker-compose` example that works from *outside* the container network — the scenario #1349 reported as broken.

### Definition of done (proposed)

- A third-party client works end-to-end: Arrow Flight SQL JDBC driver and `adbc`/`pyarrow` both connect, introspect, and run TPC-H queries against a multi-executor cluster.
- Results are fetchable from outside the cluster network without leaking executor addresses to the client.
- Integration tests in CI covering the query, metadata, prepared-statement, and cancellation paths — non-negotiable this time, given the old implementation shipped with none.
- Regression coverage for the specific historical failures: #1012 (endpoint location correctness), #941 (pyarrow decode), #756, #839.
- No new coupling to `SchedulerServer` internals.

## Describe alternatives you've considered

- **Resurrect `flight_sql.rs` from 45.0.0.** Cheapest start, and it's a genuinely useful *reference* for statement handling and endpoint construction (`git show 559bcf29^:ballista/scheduler/src/flight_sql.rs`). But it's coupled to concrete scheduler generics, predates the Flight proxy, targets an `arrow-flight` API that has drifted several major versions, and carries the incompleteness that got it removed. Reference, don't restore.
- **Reuse an existing DataFusion Flight SQL server implementation** (e.g. from `datafusion-contrib`, which #1227 floated as the home for the old code) and adapt it to submit through Ballista instead of executing locally. Worth evaluating seriously before writing a new server from scratch — I haven't assessed how well its abstractions fit distributed submission and multi-endpoint results. If it fits, most of this issue reduces to a `QueryBackend` implementation.
- **Do nothing; wait for Spark Connect (#2249).** Spark Connect is a much larger surface and reaches a different audience. It doesn't serve JDBC/ODBC/BI tools, so it isn't a substitute.
- **Client-side Flight SQL gateway** (a separate process holding a DataFusion `SessionContext` that forwards to Ballista). Keeps the scheduler clean, but duplicates session and catalog state and adds a hop; probably strictly worse than a mountable frontend once #2249 exists.

## Additional context

- Removal: #1227 → #1228 (559bcf29, released 46.0.0). Also removed `docs/source/user-guide/flightsql.md`.
- Historical bugs worth reading before designing: #1012, #941, #839, #756, #418, #230.
- The Flight proxy this would build on: #1349 → #1351 (ac49c18a), plus #2281/#2288 currently cleaning up its configuration.
- Pluggable frontend discussion: #2249. @phillipleblanc mentioned intending to open an epic for that work — if that lands, this should become a child of it.

I'm happy to help shepherd this, but I'd rather agree the layering and the open questions above first, particularly (1) the catalog model and (3) the driver-driven definition of the metadata surface.

Contributor guide

Open the contributing guide

Research direction

Start by reading #2249, ballista/scheduler/src/flight_proxy_service.rs, and ballista/scheduler/src/state/session_manager.rs, then compare the removed implementation with git show 559bcf29^:ballista/scheduler/src/flight_sql.rs. Define the frontend boundaries and design decisions before implementation. Done means third-party Flight SQL clients can query and introspect a multi-executor cluster, with integration and regression tests covering the proposed paths.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust, sql
Domain
backend-api-design, distributed-systems
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.