You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/09/03 15:11:38 UTC
[PR] feat: [WIP] Experimental Arrow Flight SQL frontend [datafusion-ballista]
andygrove opened a new pull request, #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416
> **Disclaimer:** this PR was implemented with an LLM (Claude Code) at my request.
> It is **WIP / experimental**, opened for design discussion on #2298 rather than
> as a merge candidate. Please read "Status" below before reviewing in depth —
> the end-to-end tests have not been run.
# Which issue does this PR close?
Part of #2298. Does not close it.
# Rationale for this change
Ballista has no way for a generic SQL client to connect. Flight SQL was removed
in 46.0.0 (#1227, #1228) because it was unmaintained, substantially incomplete,
structurally coupled to `SchedulerServer`, and had zero tests. This is a fresh
implementation aimed at those four objections, not a revival of the old
`flight_sql.rs`.
The concrete motivation was querying a Ballista cluster from Python via
`adbc_driver_flightsql`, bypassing the Ballista Python bindings.
# What changes are included in this PR?
**New crate `ballista/flight-sql`**, depending only on `ballista-core` — it
never references `SchedulerServer`.
- `backend.rs` — a `QueryBackend` trait (session open/close, submit-and-await,
cancel), the transport-neutral seam #2298 asked for.
- `service.rs` — `BallistaFlightSqlService<B: QueryBackend>`: handshake,
`CommandStatementQuery`, `DoGet`, prepared statements, `CommandStatementUpdate`,
`CancelQuery`.
- `metadata.rs` — `SqlInfo`, `XdbcTypeInfo`, and catalog introspection sourced
from the session's real DataFusion catalog.
- `auth.rs` — pluggable `Authenticator`, no built-in credentials.
- `session.rs` — TTL'd stores for sessions, prepared statements, and results.
**Scheduler wiring** — `ballista/scheduler/src/flight_sql.rs` (~100 lines)
implements `QueryBackend`, behind a non-default `flight-sql` feature and a
`--flight-sql` runtime flag.
**Supporting changes** — `flight_proxy_service.rs` moves from
`ballista/scheduler/src` to `ballista/core/src` unchanged so both crates can use
it; `start_grpc_service_with_listener` extracted so callers can bind port 0;
`docs/source/user-guide/flightsql.md` replaces the page deleted in #1228;
`examples/python/adbc_flight_sql.py`; a CI step running the tests with the
feature on.
# Design decisions worth discussing
Detail and rationale for each is in #2298; summarising what this PR chose:
1. **Endpoints carry no location**, so clients redeem tickets on the connection
they already have and no executor address leaks. Direct fix for #1012/#1349.
2. **The frontend subsumes the plain proxy** (only one can serve
`FlightService`); `do_get_fallback` keeps Ballista's own tickets working.
3. **Catalog model: option (b)** — embedder-supplied via `SessionBuilder`, plus
session-scoped DDL.
4. **`advertise_flight_sql_endpoint` left alone** (#2298 item 8); a separate
`--flight-sql` flag was added instead.
5. **CTAS, `INSERT`/`UPDATE`/`DELETE`/`COPY` are refused** rather than silently
executed on the scheduler.
6. **`CleanJobData` not wired** — the scheduler's existing delayed cleanup owns
result retention.
Not implemented: `PollFlightInfo` (so `GetFlightInfo` blocks until the query
finishes), bound parameters, transactions, Substrait, key metadata.
**Security — the main reason this is marked experimental.** The default
authenticator accepts every handshake. Separately, the Flight proxy trusts the
executor address carried in a ticket, so a forged ticket can make the scheduler
dial an arbitrary host and relay the response. That second one is pre-existing
in `BallistaFlightProxyService`, but enabling Flight SQL is what makes the port
something you would consider exposing. Both are documented; neither is fixed
here, and the fix wants its own issue.
# Status
**Please do not read green CI as end-to-end validation.**
Passing locally: 28 tests — 16 unit, plus 12 protocol-level tests driving the
`FlightSqlService` trait against a stub backend and asserting on decoded wire
bytes rather than crate internals. The scheduler's 350 unit tests pass with the
feature enabled; clippy, rustfmt, `--locked` and rustdoc are clean.
**Not run:** the three end-to-end tests in
`ballista/scheduler/tests/flight_sql.rs` (real scheduler + executor, driven by
`FlightSqlServiceClient`). They compile, but the environment this was developed
in blocks TCP `bind`, so the live scheduler↔executor↔client path is unverified,
as is the Python ADBC example.
# Are there any user-facing changes?
Only opt-in. `flight-sql` is a non-default compile-time feature and
`--flight-sql` is off at runtime, so no existing deployment changes behaviour.
Adds a user guide page and a Python example.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#issuecomment-5645571027
not sure why i see a empty line after each line
<img width="867" height="356" alt="image" src="https://github.com/user-attachments/assets/589aec3c-febf-4210-8d04-2c29cb776ab7" />
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#discussion_r4116246953
##########
ballista/scheduler/src/state/mod.rs:
##########
@@ -305,23 +340,35 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
.await
{
Ok(executor) => {
- if let Err(e) = state
+ match state
.task_manager
.launch_multi_task(&executor, tasks, &state.executor_manager)
.await
{
- let err_msg = format!("Failed to launch new task: {e}");
- error!("{}", err_msg.clone());
-
- // It's OK to remove executor aggressively,
- // since if the executor is in healthy state, it will be registered again.
- state
- .remove_executor(&executor_id, Some(err_msg), &sender)
- .await;
-
- false
- } else {
- true
+ Ok(unpreparable) => {
Review Comment:
This is its own PR now, #2492, and it has merged. Unpreparable tasks go through the same path #2477 and #2016 use for rejected tasks, so their slots are refunded and the job fails. With main merged in, the change no longer shows up here.
##########
ballista/scheduler/src/state/task_manager.rs:
##########
@@ -1456,4 +1471,42 @@ mod tests {
assert_eq!(manager.running_job_number(), 0);
Ok(())
}
+
+ /// A task whose definition cannot be prepared is never sent to an executor,
+ /// so no status will ever come back for it. `launch_multi_task` has to name
+ /// the job it dropped, or the job stays Running for the life of the
+ /// scheduler and everything awaiting its status blocks with no timeout.
+ #[tokio::test]
+ async fn tasks_that_cannot_be_prepared_are_reported_against_their_job() -> Result<()>
Review Comment:
This moved to #2492 along with the fix. The test there binds two jobs, makes one job's tasks unpreparable, and checks that only that job is reported failed and that exactly its slots are freed. It fails without the fix.
##########
ballista/core/src/serde/scheduler/mod.rs:
##########
@@ -71,11 +71,27 @@ impl PartitionLocation {
/// vocabulary of one writer. The stored field keeps that name until the
/// construction sites are swept; the wire protocol already speaks layouts.
pub fn layout(&self) -> ShuffleLayout {
- if self.is_sort_shuffle {
- ShuffleLayout::Sort
- } else {
- ShuffleLayout::Passthrough
- }
+ shuffle_layout(self.is_sort_shuffle)
+ }
+}
+
+impl crate::serde::protobuf::PartitionLocation {
Review Comment:
Done in #2495, which has merged. `fetch_partition` now uses `layout()` too, so the rule lives in one place.
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
+ return Err(Status::unimplemented(reason));
+ }
+
+ if disposition(&plan) == Disposition::RunOnScheduler {
+ let (schema, batches) = execute_locally(&ctx, plan).await?;
+ let handle = Uuid::new_v4().to_string();
+ self.store.insert_result(
+ handle.clone(),
+ LocalResult {
+ schema: schema.clone(),
+ batches,
+ },
+ );
+
+ let endpoint = Self::endpoint(StatementHandle::Local(handle));
+ return build_flight_info(&schema, vec![endpoint], descriptor);
+ }
+
+ let result = self
+ .backend
+ .execute(job_name, ctx, plan)
+ .await
+ .map_err(query_failed)?;
+
+ let endpoints = result
+ .partitions
+ .into_iter()
+ .map(|location| partition_handle(location).map(Self::endpoint))
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| Status::internal(format!("invalid partition location: {e}")))?;
+
+ log::debug!(
+ "flight-sql: job {} produced {} endpoint(s)",
+ result.job_id,
+ endpoints.len()
+ );
+
+ build_flight_info(&result.schema, endpoints, descriptor)
+ }
+
+ /// Builds the endpoint a client redeems for one slice of the result.
+ ///
+ /// Endpoints carry no location, which Flight defines as "fetch from the
+ /// server that gave you this FlightInfo". That keeps every cluster-internal
+ /// address off the wire and lets the frontend work unchanged behind NAT, a
+ /// load balancer, or an ingress.
+ fn endpoint(handle: StatementHandle) -> FlightEndpoint {
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.encode().into(),
+ };
+ FlightEndpoint::new().with_ticket(Ticket {
+ ticket: ticket.as_any().encode_to_vec().into(),
+ })
+ }
+
+ /// Serves a metadata command by round-tripping the command itself as the
+ /// ticket, so `DoGet` lands back on the matching handler.
+ fn metadata_info<C: ProstMessageExt>(
+ command: C,
+ schema: SchemaRef,
+ descriptor: FlightDescriptor,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let endpoint = FlightEndpoint::new().with_ticket(Ticket {
+ ticket: command.as_any().encode_to_vec().into(),
+ });
+ build_flight_info(&schema, vec![endpoint], descriptor).map(Response::new)
+ }
+}
+
+/// What the frontend should do with a planned statement.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Disposition {
+ /// Execute on the scheduler: it mutates session state and produces no data
+ /// worth distributing.
+ RunOnScheduler,
+ /// Submit to the cluster.
+ Distribute,
+ /// Refuse, with an explanation for the client.
+ Unsupported(&'static str),
+}
+
+/// Classifies a plan.
+///
+/// One function rather than two predicates, because the interesting cases are
+/// the ones where the answers overlap: `CREATE TABLE AS SELECT` is DDL, and
+/// DDL runs on the scheduler, so a caller that asked "is this DDL?" before
+/// asking "is this supported?" would silently execute its query on one node.
+/// Returning a single verdict makes that ordering impossible to get wrong.
+fn disposition(plan: &LogicalPlan) -> Disposition {
+ match plan {
+ LogicalPlan::Dml(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support INSERT/UPDATE/DELETE; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Copy(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support COPY; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support CREATE TABLE AS SELECT, \
+ because it would execute on the scheduler rather than the cluster; \
+ use CREATE EXTERNAL TABLE over data the executors can read",
+ ),
+ // Other DDL only edits the session catalog, and `SET`-style statements
+ // only edit session config; neither has anything to distribute.
+ LogicalPlan::Ddl(_) | LogicalPlan::Statement(_) => Disposition::RunOnScheduler,
+ // `SHOW ...` and anything else reading `information_schema` describes
+ // the catalog held here. There is nothing to distribute, and the
+ // physical form of those scans cannot be serialized for an executor,
+ // so distributing one would hand the client a job that never runs.
+ _ if scans_only_local_tables(plan) => Disposition::RunOnScheduler,
+ _ => Disposition::Distribute,
+ }
+}
+
+async fn execute_locally(
+ ctx: &SessionContext,
+ plan: LogicalPlan,
+) -> Result<(SchemaRef, Vec<RecordBatch>), Status> {
+ let df = ctx
+ .execute_logical_plan(plan)
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+ let planned_schema: SchemaRef = Arc::new(df.schema().as_arrow().clone());
+ let batches = df
+ .collect()
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+
+ // Prefer the schema the data actually carries; DataFusion's DDL results
+ // are empty and their frame schema is not always the same object.
+ let schema = batches
+ .first()
+ .map(|batch| batch.schema())
+ .unwrap_or(planned_schema);
+
+ Ok((schema, batches))
+}
+
+/// Turns a shuffle partition into the ticket payload the Flight proxy already
+/// knows how to redeem.
+fn partition_handle(
+ location: PartitionLocation,
+) -> Result<StatementHandle, BallistaError> {
+ let layout = location.layout();
+ let partition_id = location.partition_id.ok_or_else(|| {
+ BallistaError::Internal("partition location has no partition id".to_string())
+ })?;
+ let executor = location.executor_meta.ok_or_else(|| {
+ BallistaError::Internal("partition location has no executor metadata".to_string())
+ })?;
+
+ let action = BallistaAction::FetchPartition {
+ job_id: partition_id.job_id.into(),
+ stage_id: partition_id.stage_id as usize,
+ partition_id: partition_id.partition_id as usize,
+ host: executor.host,
+ port: executor.port as u16,
+ file_id: location.file_id,
+ layout,
+ // A Flight client wants the whole partition: the data file, in full.
+ file_kind: ShuffleFileKind::Data,
+ byte_ranges: vec![],
+ };
+
+ let encoded: protobuf::Action = action.try_into()?;
+ Ok(StatementHandle::Partition(encoded.encode_to_vec()))
+}
+
+fn build_flight_info(
+ schema: &Schema,
+ endpoints: Vec<FlightEndpoint>,
+ descriptor: FlightDescriptor,
+) -> Result<FlightInfo, Status> {
+ FlightInfo::new()
+ .try_with_schema(schema)
+ .map_err(|e| Status::internal(format!("failed to encode result schema: {e}")))
+ .map(|info| info.with_descriptor(descriptor).with_endpoints(endpoints))
+}
+
+/// Streams a single metadata batch back to the client.
+fn one_batch_response(batch: RecordBatch) -> Response<DoGetStream> {
+ batch_response(batch.schema(), vec![batch])
+}
+
+fn batch_response(schema: SchemaRef, batches: Vec<RecordBatch>) -> Response<DoGetStream> {
+ let stream = FlightDataEncoderBuilder::new()
+ .with_schema(schema)
+ .build(futures::stream::iter(batches.into_iter().map(Ok)))
+ .map_err(|e| Status::internal(format!("failed to encode results: {e}")));
+
+ Response::new(Box::pin(stream) as DoGetStream)
+}
+
+fn bearer_token(metadata: &MetadataMap) -> Option<String> {
+ let value = metadata.get("authorization")?.to_str().ok()?;
+ value
+ .strip_prefix("Bearer ")
+ .or_else(|| value.strip_prefix("bearer "))
+ .map(str::to_string)
+}
+
+fn invalid_ticket(e: impl std::fmt::Display) -> Status {
+ Status::invalid_argument(format!("invalid ticket: {e}"))
+}
+
+fn query_failed(e: BallistaError) -> Status {
+ // A failed job is the client's problem to see, not an opaque 500.
+ Status::internal(format!("query execution failed: {e}"))
+}
+
+fn encode_schema(schema: &Schema) -> Result<Vec<u8>, Status> {
+ let message: IpcMessage = SchemaAsIpc::new(schema, &IpcWriteOptions::default())
+ .try_into()
+ .map_err(|e| Status::internal(format!("failed to encode schema: {e}")))?;
+ Ok(message.0.to_vec())
+}
+
+#[tonic::async_trait]
+impl<B: QueryBackend> FlightSqlService for BallistaFlightSqlService<B> {
+ type FlightService = BallistaFlightSqlService<B>;
+
+ async fn do_handshake(
+ &self,
+ request: Request<Streaming<HandshakeRequest>>,
+ ) -> Result<
+ Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
+ Status,
+ > {
+ let identity = self.auth.authenticate(request.metadata()).await?;
+
+ let token = Uuid::new_v4().to_string();
+ let session_id = format!("flight-sql-{}", Uuid::new_v4());
+
+ // Build the session eagerly so a failure surfaces at handshake time
+ // rather than on the client's first query, and so the first query does
+ // not pay for it.
+ self.open_session(&session_id).await?;
+ self.store.insert_session(token.clone(), session_id.clone());
+
+ log::debug!(
+ "flight-sql: handshake for {:?} bound to session {session_id}",
+ identity.user
+ );
+
+ let result = HandshakeResponse {
+ protocol_version: 0,
+ payload: token.clone().into(),
+ };
+ let stream = futures::stream::once(async move { Ok(result) });
+
+ let mut response = Response::new(Box::pin(stream) as _);
+ response.metadata_mut().insert(
+ "authorization",
+ format!("Bearer {token}")
+ .parse()
+ .map_err(|_| Status::internal("failed to encode session token"))?,
+ );
+ Ok(response)
+ }
+
+ /// Redeems tickets that are not Flight SQL commands.
+ ///
+ /// Ballista's own Rust client fetches shuffle output with a
+ /// `ballista.protobuf.Action` ticket. When the Flight SQL frontend is
+ /// mounted it replaces the standalone proxy on the scheduler's port, so it
+ /// has to keep serving those tickets.
+ async fn do_get_fallback(
Review Comment:
Done. The proxy now maps an undecodable ticket to `InvalidArgument` itself, and `do_get_fallback` just hands the ticket to it. The proxy half of that landed in #2495.
##########
ballista/core/src/planner.rs:
##########
@@ -165,6 +162,19 @@ impl<T: 'static + AsLogicalPlan> QueryPlanner for BallistaQueryPlanner<T> {
}
}
+/// Returns `true` if every table `plan` scans lives in `information_schema`.
+///
+/// Such a plan reads the catalog of the node it is planned on and has nothing
+/// to distribute; worse, the physical form of those scans (`StreamingTableExec`)
+/// cannot be serialized for an executor, so distributing one produces a job
+/// that can never run. Callers that hand plans to the cluster should answer
+/// these locally instead.
+pub fn scans_only_local_tables(plan: &LogicalPlan) -> bool {
+ let mut local_run = LocalRun::default();
+ let _ = plan.visit(&mut local_run);
Review Comment:
Good catch. I split this out as #2493, which switches `LocalRun` to `visit_with_subqueries` and adds a test with this query. It's still open, so the change stays in this diff until it merges.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "flyovers-bot (via GitHub)" <gi...@apache.org>.
flyovers-bot commented on PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#issuecomment-5590595496
<!-- flyovers:flyover -->
[](https://flyovers.dev/flyover/688feacb-e5e2-4c74-aa2e-06cbc93f593e)
**[Watch the flyover →](https://flyovers.dev/flyover/688feacb-e5e2-4c74-aa2e-06cbc93f593e)** · 4:05 · 3 chapters · 32 files
An animated walkthrough of this diff: the shape of the change, what it adds, and where the risk sits.
<sub>250500c · updated automatically by flyovers</sub>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "villebro (via GitHub)" <gi...@apache.org>.
villebro commented on PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#issuecomment-5799476324
> not sure why i see a empty line after each line
@milenkovicm I've been seeing this from time to time. It's a GitHub rendering issue, and goes away with a browser refresh. Very annoying.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#issuecomment-5858069257
> **Scope.** The `PushStaged` job-failure fix, the `EXPLAIN` nullability change, and the ADBC benchmark and example (about 430 lines of Python) could each land separately.
Good call. I split all three out, along with a couple of other self-contained pieces:
- #2492 fails jobs whose tasks can't be prepared. Merged.
- #2494 makes the distributed `EXPLAIN` columns not-null. Merged.
- #2495 moves the Flight proxy service into `ballista-core`. Merged.
- #2493 makes `scans_only_local_tables` look inside subqueries. Still open.
- #2496 has the ADBC example and TPC-H benchmark. It's a draft stacked on this PR.
I just merged main into this branch, so the three merged PRs have dropped out of the diff here. The planner change will drop out too once #2493 lands.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#issuecomment-5969361948
Thanks @milenkovicm. Given that this is a new feature and disabled by default, I think it is safe to merge now and we can refine for the next release.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#discussion_r4116283899
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
+ return Err(Status::unimplemented(reason));
+ }
+
+ if disposition(&plan) == Disposition::RunOnScheduler {
+ let (schema, batches) = execute_locally(&ctx, plan).await?;
+ let handle = Uuid::new_v4().to_string();
+ self.store.insert_result(
+ handle.clone(),
+ LocalResult {
+ schema: schema.clone(),
+ batches,
+ },
+ );
+
+ let endpoint = Self::endpoint(StatementHandle::Local(handle));
+ return build_flight_info(&schema, vec![endpoint], descriptor);
+ }
+
+ let result = self
+ .backend
+ .execute(job_name, ctx, plan)
+ .await
+ .map_err(query_failed)?;
+
+ let endpoints = result
+ .partitions
+ .into_iter()
+ .map(|location| partition_handle(location).map(Self::endpoint))
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| Status::internal(format!("invalid partition location: {e}")))?;
+
+ log::debug!(
+ "flight-sql: job {} produced {} endpoint(s)",
+ result.job_id,
+ endpoints.len()
+ );
+
+ build_flight_info(&result.schema, endpoints, descriptor)
+ }
+
+ /// Builds the endpoint a client redeems for one slice of the result.
+ ///
+ /// Endpoints carry no location, which Flight defines as "fetch from the
+ /// server that gave you this FlightInfo". That keeps every cluster-internal
+ /// address off the wire and lets the frontend work unchanged behind NAT, a
+ /// load balancer, or an ingress.
+ fn endpoint(handle: StatementHandle) -> FlightEndpoint {
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.encode().into(),
+ };
+ FlightEndpoint::new().with_ticket(Ticket {
+ ticket: ticket.as_any().encode_to_vec().into(),
+ })
+ }
+
+ /// Serves a metadata command by round-tripping the command itself as the
+ /// ticket, so `DoGet` lands back on the matching handler.
+ fn metadata_info<C: ProstMessageExt>(
+ command: C,
+ schema: SchemaRef,
+ descriptor: FlightDescriptor,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let endpoint = FlightEndpoint::new().with_ticket(Ticket {
+ ticket: command.as_any().encode_to_vec().into(),
+ });
+ build_flight_info(&schema, vec![endpoint], descriptor).map(Response::new)
+ }
+}
+
+/// What the frontend should do with a planned statement.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Disposition {
+ /// Execute on the scheduler: it mutates session state and produces no data
+ /// worth distributing.
+ RunOnScheduler,
+ /// Submit to the cluster.
+ Distribute,
+ /// Refuse, with an explanation for the client.
+ Unsupported(&'static str),
+}
+
+/// Classifies a plan.
+///
+/// One function rather than two predicates, because the interesting cases are
+/// the ones where the answers overlap: `CREATE TABLE AS SELECT` is DDL, and
+/// DDL runs on the scheduler, so a caller that asked "is this DDL?" before
+/// asking "is this supported?" would silently execute its query on one node.
+/// Returning a single verdict makes that ordering impossible to get wrong.
+fn disposition(plan: &LogicalPlan) -> Disposition {
+ match plan {
+ LogicalPlan::Dml(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support INSERT/UPDATE/DELETE; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Copy(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support COPY; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support CREATE TABLE AS SELECT, \
+ because it would execute on the scheduler rather than the cluster; \
+ use CREATE EXTERNAL TABLE over data the executors can read",
+ ),
+ // Other DDL only edits the session catalog, and `SET`-style statements
+ // only edit session config; neither has anything to distribute.
+ LogicalPlan::Ddl(_) | LogicalPlan::Statement(_) => Disposition::RunOnScheduler,
+ // `SHOW ...` and anything else reading `information_schema` describes
+ // the catalog held here. There is nothing to distribute, and the
+ // physical form of those scans cannot be serialized for an executor,
+ // so distributing one would hand the client a job that never runs.
+ _ if scans_only_local_tables(plan) => Disposition::RunOnScheduler,
+ _ => Disposition::Distribute,
+ }
+}
+
+async fn execute_locally(
+ ctx: &SessionContext,
+ plan: LogicalPlan,
+) -> Result<(SchemaRef, Vec<RecordBatch>), Status> {
+ let df = ctx
+ .execute_logical_plan(plan)
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+ let planned_schema: SchemaRef = Arc::new(df.schema().as_arrow().clone());
+ let batches = df
+ .collect()
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+
+ // Prefer the schema the data actually carries; DataFusion's DDL results
+ // are empty and their frame schema is not always the same object.
+ let schema = batches
+ .first()
+ .map(|batch| batch.schema())
+ .unwrap_or(planned_schema);
+
+ Ok((schema, batches))
+}
+
+/// Turns a shuffle partition into the ticket payload the Flight proxy already
+/// knows how to redeem.
+fn partition_handle(
+ location: PartitionLocation,
+) -> Result<StatementHandle, BallistaError> {
+ let layout = location.layout();
+ let partition_id = location.partition_id.ok_or_else(|| {
+ BallistaError::Internal("partition location has no partition id".to_string())
+ })?;
+ let executor = location.executor_meta.ok_or_else(|| {
+ BallistaError::Internal("partition location has no executor metadata".to_string())
+ })?;
+
+ let action = BallistaAction::FetchPartition {
+ job_id: partition_id.job_id.into(),
+ stage_id: partition_id.stage_id as usize,
+ partition_id: partition_id.partition_id as usize,
+ host: executor.host,
+ port: executor.port as u16,
+ file_id: location.file_id,
+ layout,
+ // A Flight client wants the whole partition: the data file, in full.
+ file_kind: ShuffleFileKind::Data,
+ byte_ranges: vec![],
+ };
+
+ let encoded: protobuf::Action = action.try_into()?;
+ Ok(StatementHandle::Partition(encoded.encode_to_vec()))
+}
+
+fn build_flight_info(
+ schema: &Schema,
+ endpoints: Vec<FlightEndpoint>,
+ descriptor: FlightDescriptor,
+) -> Result<FlightInfo, Status> {
+ FlightInfo::new()
+ .try_with_schema(schema)
+ .map_err(|e| Status::internal(format!("failed to encode result schema: {e}")))
+ .map(|info| info.with_descriptor(descriptor).with_endpoints(endpoints))
+}
+
+/// Streams a single metadata batch back to the client.
+fn one_batch_response(batch: RecordBatch) -> Response<DoGetStream> {
+ batch_response(batch.schema(), vec![batch])
+}
+
+fn batch_response(schema: SchemaRef, batches: Vec<RecordBatch>) -> Response<DoGetStream> {
+ let stream = FlightDataEncoderBuilder::new()
+ .with_schema(schema)
+ .build(futures::stream::iter(batches.into_iter().map(Ok)))
+ .map_err(|e| Status::internal(format!("failed to encode results: {e}")));
+
+ Response::new(Box::pin(stream) as DoGetStream)
+}
+
+fn bearer_token(metadata: &MetadataMap) -> Option<String> {
+ let value = metadata.get("authorization")?.to_str().ok()?;
+ value
+ .strip_prefix("Bearer ")
+ .or_else(|| value.strip_prefix("bearer "))
+ .map(str::to_string)
+}
+
+fn invalid_ticket(e: impl std::fmt::Display) -> Status {
+ Status::invalid_argument(format!("invalid ticket: {e}"))
+}
+
+fn query_failed(e: BallistaError) -> Status {
+ // A failed job is the client's problem to see, not an opaque 500.
+ Status::internal(format!("query execution failed: {e}"))
+}
+
+fn encode_schema(schema: &Schema) -> Result<Vec<u8>, Status> {
+ let message: IpcMessage = SchemaAsIpc::new(schema, &IpcWriteOptions::default())
+ .try_into()
+ .map_err(|e| Status::internal(format!("failed to encode schema: {e}")))?;
+ Ok(message.0.to_vec())
+}
+
+#[tonic::async_trait]
+impl<B: QueryBackend> FlightSqlService for BallistaFlightSqlService<B> {
+ type FlightService = BallistaFlightSqlService<B>;
+
+ async fn do_handshake(
+ &self,
+ request: Request<Streaming<HandshakeRequest>>,
+ ) -> Result<
+ Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
+ Status,
+ > {
+ let identity = self.auth.authenticate(request.metadata()).await?;
+
+ let token = Uuid::new_v4().to_string();
+ let session_id = format!("flight-sql-{}", Uuid::new_v4());
+
+ // Build the session eagerly so a failure surfaces at handshake time
+ // rather than on the client's first query, and so the first query does
+ // not pay for it.
+ self.open_session(&session_id).await?;
+ self.store.insert_session(token.clone(), session_id.clone());
+
+ log::debug!(
+ "flight-sql: handshake for {:?} bound to session {session_id}",
+ identity.user
+ );
+
+ let result = HandshakeResponse {
+ protocol_version: 0,
+ payload: token.clone().into(),
+ };
+ let stream = futures::stream::once(async move { Ok(result) });
+
+ let mut response = Response::new(Box::pin(stream) as _);
+ response.metadata_mut().insert(
+ "authorization",
+ format!("Bearer {token}")
+ .parse()
+ .map_err(|_| Status::internal("failed to encode session token"))?,
+ );
+ Ok(response)
+ }
+
+ /// Redeems tickets that are not Flight SQL commands.
+ ///
+ /// Ballista's own Rust client fetches shuffle output with a
+ /// `ballista.protobuf.Action` ticket. When the Flight SQL frontend is
+ /// mounted it replaces the standalone proxy on the scheduler's port, so it
+ /// has to keep serving those tickets.
+ async fn do_get_fallback(
+ &self,
+ request: Request<Ticket>,
+ _message: Any,
+ ) -> Result<Response<DoGetStream>, Status> {
+ decode_protobuf(&request.get_ref().ticket).map_err(invalid_ticket)?;
+ self.proxy.do_get(request).await
+ }
+
+ async fn get_flight_info_statement(
+ &self,
+ query: CommandStatementQuery,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ let plan = Self::plan(&ctx, &query.query).await?;
+ let descriptor = request.into_inner();
+
+ self.flight_info_for(ctx, plan, descriptor, &job_name(&query.query))
+ .await
+ .map(Response::new)
+ }
+
+ async fn get_flight_info_prepared_statement(
+ &self,
+ query: CommandPreparedStatementQuery,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let handle = prepared_handle(&query.prepared_statement_handle)?;
+ let prepared = self.store.prepared(&handle).ok_or_else(|| {
+ Status::not_found("unknown or expired prepared statement handle")
+ })?;
+
+ let ctx = self.open_session(&prepared.session_id).await?;
+ let descriptor = request.into_inner();
+
+ self.flight_info_for(ctx, prepared.plan, descriptor, "flight-sql prepared")
+ .await
+ .map(Response::new)
+ }
+
+ async fn get_flight_info_catalogs(
+ &self,
+ query: CommandGetCatalogs,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_schemas(
+ &self,
+ query: CommandGetDbSchemas,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_tables(
+ &self,
+ query: CommandGetTables,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_table_types(
+ &self,
+ query: CommandGetTableTypes,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_sql_info(
+ &self,
+ query: CommandGetSqlInfo,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder(&self.sql_info).schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_xdbc_type_info(
+ &self,
+ query: CommandGetXdbcTypeInfo,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder(&self.xdbc_info).schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn do_get_statement(
+ &self,
+ ticket: TicketStatementQuery,
+ _request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let handle =
+ StatementHandle::decode(&ticket.statement_handle).map_err(invalid_ticket)?;
+
+ match handle {
+ StatementHandle::Partition(action) => {
+ // Unwrap back to the executor-facing ticket and let the proxy
+ // redeem it. The proxy dials the executor with a fresh request,
+ // so the client's credentials are not forwarded onwards.
+ self.proxy
+ .do_get(Request::new(Ticket {
+ ticket: action.into(),
+ }))
+ .await
+ }
+ StatementHandle::Local(handle) => {
+ let result = self.store.take_result(&handle).ok_or_else(|| {
+ Status::not_found("result already consumed or expired")
+ })?;
+ Ok(batch_response(result.schema, result.batches))
+ }
+ }
+ }
+
+ async fn do_get_catalogs(
+ &self,
+ query: CommandGetCatalogs,
+ request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ Ok(one_batch_response(metadata::catalogs(&ctx, query)?))
+ }
+
+ async fn do_get_schemas(
+ &self,
+ query: CommandGetDbSchemas,
+ request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ Ok(one_batch_response(metadata::db_schemas(&ctx, query)?))
+ }
+
+ async fn do_get_tables(
+ &self,
+ query: CommandGetTables,
+ request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ Ok(one_batch_response(metadata::tables(&ctx, query).await?))
+ }
+
+ async fn do_get_table_types(
+ &self,
+ query: CommandGetTableTypes,
+ _request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ Ok(one_batch_response(metadata::table_types(query)?))
+ }
+
+ async fn do_get_sql_info(
+ &self,
+ query: CommandGetSqlInfo,
+ _request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let batch = query
+ .into_builder(&self.sql_info)
+ .build()
+ .map_err(|e| Status::internal(format!("failed to build SqlInfo: {e}")))?;
+ Ok(one_batch_response(batch))
+ }
+
+ async fn do_get_xdbc_type_info(
+ &self,
+ query: CommandGetXdbcTypeInfo,
+ _request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let batch = query
+ .into_builder(&self.xdbc_info)
+ .build()
+ .map_err(|e| Status::internal(format!("failed to build type info: {e}")))?;
+ Ok(one_batch_response(batch))
+ }
+
+ async fn do_put_statement_update(
+ &self,
+ ticket: CommandStatementUpdate,
+ request: Request<PeekableFlightDataStream>,
+ ) -> Result<i64, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ let plan = Self::plan(&ctx, &ticket.query).await?;
+
+ match disposition(&plan) {
+ Disposition::RunOnScheduler => {}
+ Disposition::Unsupported(reason) => {
+ return Err(Status::unimplemented(reason));
+ }
+ Disposition::Distribute => {
+ return Err(Status::unimplemented(
+ "Ballista Flight SQL supports DDL and session statements on this \
+ path; the distributed write path is not implemented",
+ ));
+ }
+ }
+
+ execute_locally(&ctx, plan).await?;
+ // Only DDL and session statements reach here, and neither reports an
+ // affected-row count.
+ Ok(0)
+ }
+
+ async fn do_action_create_prepared_statement(
+ &self,
+ query: ActionCreatePreparedStatementRequest,
+ request: Request<Action>,
+ ) -> Result<ActionCreatePreparedStatementResult, Status> {
+ let session_id = self.session_id(request.metadata())?;
+ let ctx = self.open_session(&session_id).await?;
+ let plan = Self::plan(&ctx, &query.query).await?;
+ let schema = plan.schema().as_arrow().clone();
+
+ let handle = Uuid::new_v4().to_string();
+ self.store
+ .insert_prepared(handle.clone(), Prepared { session_id, plan });
+
+ Ok(ActionCreatePreparedStatementResult {
+ prepared_statement_handle: handle.into_bytes().into(),
+ dataset_schema: encode_schema(&schema)?.into(),
+ // Bound parameters are not supported yet, so the parameter schema
+ // is empty rather than absent: clients read "no parameters".
+ parameter_schema: encode_schema(&Schema::empty())?.into(),
+ })
+ }
+
+ async fn do_action_close_prepared_statement(
+ &self,
+ query: ActionClosePreparedStatementRequest,
+ _request: Request<Action>,
+ ) -> Result<(), Status> {
+ let handle = prepared_handle(&query.prepared_statement_handle)?;
+ self.store.remove_prepared(&handle);
+ Ok(())
+ }
+
+ async fn do_action_cancel_query(
Review Comment:
Removed. `do_action_cancel_query`, `job_id_from_flight_info`, `QueryBackend::cancel` and the test are gone, and `FlightSqlServerCancel` now reports `false`. The crate README lists query cancellation under not implemented. It can come back with `PollFlightInfo`.
##########
ballista/flight-sql/src/metadata.rs:
##########
@@ -0,0 +1,295 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Driver-facing metadata: server capabilities, type info, and catalog
+//! introspection.
+//!
+//! Catalog answers come from the session's real DataFusion catalog rather than
+//! hand-built record batches, so what a driver sees in its schema browser is
+//! what a query in the same session can actually reference.
+
+use arrow::array::RecordBatch;
+use arrow_flight::sql::metadata::{
+ GetCatalogsBuilder, GetDbSchemasBuilder, GetTablesBuilder, SqlInfoData,
+ SqlInfoDataBuilder, XdbcTypeInfo, XdbcTypeInfoData, XdbcTypeInfoDataBuilder,
+};
+use arrow_flight::sql::{
+ CommandGetCatalogs, CommandGetDbSchemas, CommandGetTableTypes, CommandGetTables,
+ Nullable, Searchable, SqlInfo, SqlSupportedTransaction, XdbcDataType,
+};
+use ballista_core::BALLISTA_VERSION;
+use datafusion::catalog::TableProvider;
+use datafusion::logical_expr::TableType;
+use datafusion::prelude::SessionContext;
+use tonic::Status;
+
+/// Table type strings reported to clients, matching the JDBC vocabulary that
+/// Flight SQL drivers expect.
+const TABLE: &str = "TABLE";
+const VIEW: &str = "VIEW";
+const LOCAL_TEMPORARY: &str = "LOCAL TEMPORARY";
+
+/// Describes the server to drivers deciding what SQL they may emit.
+///
+/// The old implementation left `CommandGetSqlInfo` unimplemented with a
+/// `// TODO: implement for FlightSQL JDBC to work` comment, which is precisely
+/// why the JDBC driver could not connect.
+pub(crate) fn sql_info() -> SqlInfoData {
+ let mut builder = SqlInfoDataBuilder::new();
+
+ builder.append(SqlInfo::FlightSqlServerName, "Apache DataFusion Ballista");
+ builder.append(SqlInfo::FlightSqlServerVersion, BALLISTA_VERSION);
+ // Arrow IPC format version, per format/Schema.fbs.
+ builder.append(SqlInfo::FlightSqlServerArrowVersion, "1.3");
+ builder.append(SqlInfo::FlightSqlServerReadOnly, false);
+ builder.append(SqlInfo::FlightSqlServerSql, true);
+ builder.append(SqlInfo::FlightSqlServerSubstrait, false);
+ builder.append(
+ SqlInfo::FlightSqlServerTransaction,
+ SqlSupportedTransaction::None as i32,
+ );
+ builder.append(SqlInfo::FlightSqlServerCancel, true);
+ builder.append(SqlInfo::FlightSqlServerBulkIngestion, false);
+
+ builder.append(SqlInfo::SqlDdlCatalog, false);
+ builder.append(SqlInfo::SqlDdlSchema, true);
+ builder.append(SqlInfo::SqlDdlTable, true);
+ // DataFusion folds unquoted identifiers to lower case and preserves the
+ // case of quoted ones: SQL_CASE_SENSITIVITY_LOWERCASE / _CASE_INSENSITIVE.
+ builder.append(SqlInfo::SqlIdentifierCase, 2i32);
Review Comment:
Fixed. Both use the named `SqlSupportedCaseSensitivity` values now. Unquoted identifiers report `SqlCaseSensitivityLowercase`. The enum has no value for quoted identifiers stored as written, so I followed the Arrow reference server (`FlightSqlExample`) and used `SqlCaseSensitivityCaseInsensitive`. The comment explains the choice.
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
+ return Err(Status::unimplemented(reason));
+ }
+
+ if disposition(&plan) == Disposition::RunOnScheduler {
+ let (schema, batches) = execute_locally(&ctx, plan).await?;
+ let handle = Uuid::new_v4().to_string();
+ self.store.insert_result(
+ handle.clone(),
+ LocalResult {
+ schema: schema.clone(),
+ batches,
+ },
+ );
+
+ let endpoint = Self::endpoint(StatementHandle::Local(handle));
+ return build_flight_info(&schema, vec![endpoint], descriptor);
+ }
+
+ let result = self
+ .backend
+ .execute(job_name, ctx, plan)
+ .await
+ .map_err(query_failed)?;
+
+ let endpoints = result
+ .partitions
+ .into_iter()
+ .map(|location| partition_handle(location).map(Self::endpoint))
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| Status::internal(format!("invalid partition location: {e}")))?;
+
+ log::debug!(
+ "flight-sql: job {} produced {} endpoint(s)",
+ result.job_id,
+ endpoints.len()
+ );
+
+ build_flight_info(&result.schema, endpoints, descriptor)
+ }
+
+ /// Builds the endpoint a client redeems for one slice of the result.
+ ///
+ /// Endpoints carry no location, which Flight defines as "fetch from the
+ /// server that gave you this FlightInfo". That keeps every cluster-internal
+ /// address off the wire and lets the frontend work unchanged behind NAT, a
+ /// load balancer, or an ingress.
+ fn endpoint(handle: StatementHandle) -> FlightEndpoint {
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.encode().into(),
+ };
+ FlightEndpoint::new().with_ticket(Ticket {
+ ticket: ticket.as_any().encode_to_vec().into(),
+ })
+ }
+
+ /// Serves a metadata command by round-tripping the command itself as the
+ /// ticket, so `DoGet` lands back on the matching handler.
+ fn metadata_info<C: ProstMessageExt>(
+ command: C,
+ schema: SchemaRef,
+ descriptor: FlightDescriptor,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let endpoint = FlightEndpoint::new().with_ticket(Ticket {
+ ticket: command.as_any().encode_to_vec().into(),
+ });
+ build_flight_info(&schema, vec![endpoint], descriptor).map(Response::new)
+ }
+}
+
+/// What the frontend should do with a planned statement.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Disposition {
+ /// Execute on the scheduler: it mutates session state and produces no data
+ /// worth distributing.
+ RunOnScheduler,
+ /// Submit to the cluster.
+ Distribute,
+ /// Refuse, with an explanation for the client.
+ Unsupported(&'static str),
+}
+
+/// Classifies a plan.
+///
+/// One function rather than two predicates, because the interesting cases are
+/// the ones where the answers overlap: `CREATE TABLE AS SELECT` is DDL, and
+/// DDL runs on the scheduler, so a caller that asked "is this DDL?" before
+/// asking "is this supported?" would silently execute its query on one node.
+/// Returning a single verdict makes that ordering impossible to get wrong.
+fn disposition(plan: &LogicalPlan) -> Disposition {
+ match plan {
+ LogicalPlan::Dml(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support INSERT/UPDATE/DELETE; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Copy(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support COPY; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support CREATE TABLE AS SELECT, \
+ because it would execute on the scheduler rather than the cluster; \
+ use CREATE EXTERNAL TABLE over data the executors can read",
+ ),
+ // Other DDL only edits the session catalog, and `SET`-style statements
+ // only edit session config; neither has anything to distribute.
+ LogicalPlan::Ddl(_) | LogicalPlan::Statement(_) => Disposition::RunOnScheduler,
Review Comment:
Fixed. `refusal` now walks the whole plan with `apply_with_subqueries` and refuses `Dml`, `Copy`, `CreateMemoryTable` and `Statement::Execute` wherever they appear. `EXECUTE q` and `EXPLAIN ANALYZE INSERT ...` are both refused now, and `statements_that_cannot_be_distributed_are_refused` covers them.
##########
ballista/scheduler/src/config.rs:
##########
@@ -397,6 +405,19 @@ pub struct SchedulerConfig {
pub task_max_failures: usize,
/// Number of failures attempts before stage is considered failed
pub stage_max_failures: usize,
+ /// Whether to serve Arrow Flight SQL on the scheduler's gRPC port.
+ ///
+ /// The frontend subsumes the plain Flight result proxy when enabled, so
+ /// Ballista's own clients keep working on the same port.
+ #[cfg(feature = "flight-sql")]
+ pub flight_sql: bool,
+ /// Authenticates Arrow Flight SQL handshakes.
+ ///
+ /// Without one the frontend accepts every client, so this is the only way
Review Comment:
Reworded. The field doc now says the authenticator isolates clients' sessions but doesn't secure the port, and to keep the scheduler on a trusted network. The crate README says the same. I also just fixed the `--flight-sql` help text, which still implied the authenticator secures the port.
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
Review Comment:
Fixed. `flight_info_for` now does a single `match disposition(&plan)` over the three variants, like `do_put_statement_update`.
##########
ballista/flight-sql/src/session.rs:
##########
@@ -0,0 +1,277 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Server-side state the Flight SQL frontend keeps between requests.
+//!
+//! Everything in here is keyed by an opaque handle the client is given and
+//! hands back, and everything expires. A client that disconnects without
+//! closing its prepared statements (or that never redeems a ticket) must not
+//! pin memory forever, which is what the previous Flight SQL implementation
+//! got wrong: its plan cache only shed entries on an explicit close.
+
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::SchemaRef;
+use dashmap::DashMap;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+
+/// A prepared statement's server-side state.
+#[derive(Clone)]
+pub(crate) struct Prepared {
+ /// Session the statement was prepared in; it is planned and executed
+ /// against that session's catalog.
+ pub session_id: String,
+ /// The plan as prepared; execution never re-plans it.
+ pub plan: LogicalPlan,
+}
+
+/// Result of a statement the frontend ran locally rather than distributing
+/// (DDL, and DML whose only output is an affected-row count).
+#[derive(Clone)]
+pub(crate) struct LocalResult {
+ pub schema: SchemaRef,
+ pub batches: Vec<RecordBatch>,
+}
+
+struct Tracked<T> {
+ value: T,
+ last_used: Instant,
+}
+
+impl<T> Tracked<T> {
+ fn new(value: T) -> Self {
+ Self {
+ value,
+ last_used: Instant::now(),
+ }
+ }
+}
+
+/// Handle stores for sessions, prepared statements, and locally-computed
+/// results, all with a shared idle TTL.
+pub(crate) struct SessionStore {
+ ttl: Duration,
+ /// Bearer token -> Ballista session id.
+ sessions: DashMap<String, Tracked<String>>,
+ /// Ballista session id -> its context.
+ contexts: DashMap<String, Tracked<Arc<SessionContext>>>,
+ /// Prepared statement handle -> prepared plan.
+ prepared: DashMap<String, Tracked<Prepared>>,
+ /// Local result handle -> materialized batches.
+ results: DashMap<String, Tracked<LocalResult>>,
+}
+
+impl SessionStore {
+ pub(crate) fn new(ttl: Duration) -> Self {
+ Self {
+ ttl,
+ sessions: DashMap::new(),
+ contexts: DashMap::new(),
+ prepared: DashMap::new(),
+ results: DashMap::new(),
+ }
+ }
+
+ pub(crate) fn insert_session(&self, token: String, session_id: String) {
+ self.sessions.insert(token, Tracked::new(session_id));
+ }
+
+ /// Resolves a bearer token to its session id, refreshing its idle timer.
+ pub(crate) fn session(&self, token: &str) -> Option<String> {
+ self.sessions.get_mut(token).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ /// Drops a token and the session it named, returning that session id.
+ ///
+ /// Each token gets its own session, so removing the token always releases
+ /// the session with it.
+ pub(crate) fn remove_session(&self, token: &str) -> Option<String> {
+ let (_, entry) = self.sessions.remove(token)?;
+ self.contexts.remove(&entry.value);
+ Some(entry.value)
+ }
+
+ /// Returns the cached context for a session, refreshing its idle timer.
+ pub(crate) fn context(&self, session_id: &str) -> Option<Arc<SessionContext>> {
+ self.contexts.get_mut(session_id).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ /// Caches a context, or returns the one another caller cached first.
+ ///
+ /// Two requests for a cold session can both build one; letting the first
+ /// insertion win keeps them looking at the same catalog.
+ pub(crate) fn insert_context(
+ &self,
+ session_id: String,
+ ctx: Arc<SessionContext>,
+ ) -> Arc<SessionContext> {
+ self.contexts
+ .entry(session_id)
+ .or_insert_with(|| Tracked::new(ctx))
+ .value
+ .clone()
+ }
+
+ pub(crate) fn insert_prepared(&self, handle: String, prepared: Prepared) {
+ self.prepared.insert(handle, Tracked::new(prepared));
+ }
+
+ pub(crate) fn prepared(&self, handle: &str) -> Option<Prepared> {
+ self.prepared.get_mut(handle).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ pub(crate) fn remove_prepared(&self, handle: &str) {
+ self.prepared.remove(handle);
+ }
+
+ pub(crate) fn insert_result(&self, handle: String, result: LocalResult) {
+ self.results.insert(handle, Tracked::new(result));
+ }
+
+ /// Takes a local result. Results are single-use: a ticket is redeemed once.
+ pub(crate) fn take_result(&self, handle: &str) -> Option<LocalResult> {
+ self.results.remove(handle).map(|(_, entry)| entry.value)
+ }
+
+ /// Evicts everything idle for longer than the TTL.
+ ///
+ /// Returns the session ids that no longer have any live token, so the
+ /// caller can release them in the backend.
+ pub(crate) fn sweep(&self) -> Vec<String> {
Review Comment:
Done. `sweep` collects the released session ids in a single `retain` pass and drops their contexts afterwards. `remove_session` is gone.
##########
ballista/flight-sql/src/session.rs:
##########
@@ -0,0 +1,277 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Server-side state the Flight SQL frontend keeps between requests.
+//!
+//! Everything in here is keyed by an opaque handle the client is given and
+//! hands back, and everything expires. A client that disconnects without
+//! closing its prepared statements (or that never redeems a ticket) must not
+//! pin memory forever, which is what the previous Flight SQL implementation
+//! got wrong: its plan cache only shed entries on an explicit close.
+
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::SchemaRef;
+use dashmap::DashMap;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+
+/// A prepared statement's server-side state.
+#[derive(Clone)]
+pub(crate) struct Prepared {
+ /// Session the statement was prepared in; it is planned and executed
+ /// against that session's catalog.
+ pub session_id: String,
+ /// The plan as prepared; execution never re-plans it.
+ pub plan: LogicalPlan,
+}
+
+/// Result of a statement the frontend ran locally rather than distributing
+/// (DDL, and DML whose only output is an affected-row count).
+#[derive(Clone)]
+pub(crate) struct LocalResult {
+ pub schema: SchemaRef,
+ pub batches: Vec<RecordBatch>,
+}
+
+struct Tracked<T> {
+ value: T,
+ last_used: Instant,
+}
+
+impl<T> Tracked<T> {
+ fn new(value: T) -> Self {
+ Self {
+ value,
+ last_used: Instant::now(),
+ }
+ }
+}
+
+/// Handle stores for sessions, prepared statements, and locally-computed
+/// results, all with a shared idle TTL.
+pub(crate) struct SessionStore {
+ ttl: Duration,
+ /// Bearer token -> Ballista session id.
+ sessions: DashMap<String, Tracked<String>>,
+ /// Ballista session id -> its context.
+ contexts: DashMap<String, Tracked<Arc<SessionContext>>>,
+ /// Prepared statement handle -> prepared plan.
+ prepared: DashMap<String, Tracked<Prepared>>,
+ /// Local result handle -> materialized batches.
+ results: DashMap<String, Tracked<LocalResult>>,
+}
+
+impl SessionStore {
+ pub(crate) fn new(ttl: Duration) -> Self {
+ Self {
+ ttl,
+ sessions: DashMap::new(),
+ contexts: DashMap::new(),
+ prepared: DashMap::new(),
+ results: DashMap::new(),
+ }
+ }
+
+ pub(crate) fn insert_session(&self, token: String, session_id: String) {
+ self.sessions.insert(token, Tracked::new(session_id));
+ }
+
+ /// Resolves a bearer token to its session id, refreshing its idle timer.
+ pub(crate) fn session(&self, token: &str) -> Option<String> {
+ self.sessions.get_mut(token).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ /// Drops a token and the session it named, returning that session id.
+ ///
+ /// Each token gets its own session, so removing the token always releases
+ /// the session with it.
+ pub(crate) fn remove_session(&self, token: &str) -> Option<String> {
+ let (_, entry) = self.sessions.remove(token)?;
+ self.contexts.remove(&entry.value);
+ Some(entry.value)
+ }
+
+ /// Returns the cached context for a session, refreshing its idle timer.
+ pub(crate) fn context(&self, session_id: &str) -> Option<Arc<SessionContext>> {
+ self.contexts.get_mut(session_id).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ /// Caches a context, or returns the one another caller cached first.
+ ///
+ /// Two requests for a cold session can both build one; letting the first
+ /// insertion win keeps them looking at the same catalog.
+ pub(crate) fn insert_context(
+ &self,
+ session_id: String,
+ ctx: Arc<SessionContext>,
+ ) -> Arc<SessionContext> {
+ self.contexts
+ .entry(session_id)
+ .or_insert_with(|| Tracked::new(ctx))
+ .value
+ .clone()
+ }
+
+ pub(crate) fn insert_prepared(&self, handle: String, prepared: Prepared) {
+ self.prepared.insert(handle, Tracked::new(prepared));
+ }
+
+ pub(crate) fn prepared(&self, handle: &str) -> Option<Prepared> {
+ self.prepared.get_mut(handle).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ pub(crate) fn remove_prepared(&self, handle: &str) {
+ self.prepared.remove(handle);
+ }
+
+ pub(crate) fn insert_result(&self, handle: String, result: LocalResult) {
+ self.results.insert(handle, Tracked::new(result));
+ }
+
+ /// Takes a local result. Results are single-use: a ticket is redeemed once.
+ pub(crate) fn take_result(&self, handle: &str) -> Option<LocalResult> {
+ self.results.remove(handle).map(|(_, entry)| entry.value)
+ }
+
+ /// Evicts everything idle for longer than the TTL.
+ ///
+ /// Returns the session ids that no longer have any live token, so the
+ /// caller can release them in the backend.
+ pub(crate) fn sweep(&self) -> Vec<String> {
+ let ttl = self.ttl;
+ let expired: Vec<String> = self
+ .sessions
+ .iter()
+ .filter(|entry| entry.value().last_used.elapsed() > ttl)
+ .map(|entry| entry.key().clone())
+ .collect();
+
+ let mut released = Vec::new();
+ for token in expired {
+ if let Some(session_id) = self.remove_session(&token) {
+ released.push(session_id);
+ }
+ }
+
+ self.contexts
+ .retain(|_, entry| entry.last_used.elapsed() <= ttl);
+ self.prepared
+ .retain(|_, entry| entry.last_used.elapsed() <= ttl);
+ self.results
+ .retain(|_, entry| entry.last_used.elapsed() <= ttl);
+
+ released
+ }
+
+ /// Starts a background task that sweeps at `interval`, closing sessions it
+ /// evicts. The task ends when the last reference to the store is dropped.
+ pub(crate) fn spawn_reaper<F, Fut>(self: &Arc<Self>, interval: Duration, close: F)
+ where
+ F: Fn(String) -> Fut + Send + Sync + 'static,
+ Fut: std::future::Future<Output = ()> + Send,
+ {
+ let store = Arc::downgrade(self);
+ tokio::spawn(async move {
+ let mut ticker = tokio::time::interval(interval);
+ // The first tick completes immediately; skip it so we do not sweep
+ // a store that was created a moment ago.
+ ticker.tick().await;
+ loop {
+ ticker.tick().await;
+ let Some(store) = store.upgrade() else {
+ return;
+ };
+ for session_id in store.sweep() {
+ log::debug!("flight-sql: expiring idle session {session_id}");
+ close(session_id).await;
+ }
+ }
+ });
+ }
+}
+
+#[cfg(test)]
+mod test {
+ use super::*;
+
+ #[test]
+ fn session_survives_while_touched_and_expires_when_idle() {
+ let store = SessionStore::new(Duration::from_millis(50));
+ store.insert_session("token".to_string(), "session".to_string());
+
+ assert_eq!(store.session("token"), Some("session".to_string()));
+ assert!(store.sweep().is_empty());
+
+ std::thread::sleep(Duration::from_millis(60));
Review Comment:
Fixed. The test backdates `last_used` instead of sleeping, like the idle-eviction tests in `client_pool.rs`.
##########
.github/workflows/rust.yml:
##########
@@ -99,6 +99,12 @@ jobs:
export ARROW_TEST_DATA=$(pwd)/testing/data
export PARQUET_TEST_DATA=$(pwd)/parquet-testing/data
cargo test --profile ci --features=testcontainers
+ - name: Run Arrow Flight SQL tests
+ # The frontend is a non-default feature, so the workspace run above
+ # compiles its integration test away.
+ run: |
+ export PATH=$PATH:$HOME/d/protoc/bin
+ cargo test --profile ci -p ballista-scheduler --features flight-sql --test flight_sql
Review Comment:
Both are sorted. The separate step is gone and `ballista-scheduler/flight-sql` is part of `--features` in the main test step, so it's one build. Separately, #2473 switched `rust_clippy.sh` to `--workspace`, so `ballista-flight-sql` gets linted now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#discussion_r4112938218
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
+ return Err(Status::unimplemented(reason));
+ }
+
+ if disposition(&plan) == Disposition::RunOnScheduler {
+ let (schema, batches) = execute_locally(&ctx, plan).await?;
+ let handle = Uuid::new_v4().to_string();
+ self.store.insert_result(
+ handle.clone(),
+ LocalResult {
+ schema: schema.clone(),
+ batches,
+ },
+ );
+
+ let endpoint = Self::endpoint(StatementHandle::Local(handle));
+ return build_flight_info(&schema, vec![endpoint], descriptor);
+ }
+
+ let result = self
+ .backend
+ .execute(job_name, ctx, plan)
+ .await
+ .map_err(query_failed)?;
+
+ let endpoints = result
+ .partitions
+ .into_iter()
+ .map(|location| partition_handle(location).map(Self::endpoint))
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| Status::internal(format!("invalid partition location: {e}")))?;
+
+ log::debug!(
+ "flight-sql: job {} produced {} endpoint(s)",
+ result.job_id,
+ endpoints.len()
+ );
+
+ build_flight_info(&result.schema, endpoints, descriptor)
+ }
+
+ /// Builds the endpoint a client redeems for one slice of the result.
+ ///
+ /// Endpoints carry no location, which Flight defines as "fetch from the
+ /// server that gave you this FlightInfo". That keeps every cluster-internal
+ /// address off the wire and lets the frontend work unchanged behind NAT, a
+ /// load balancer, or an ingress.
+ fn endpoint(handle: StatementHandle) -> FlightEndpoint {
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.encode().into(),
+ };
+ FlightEndpoint::new().with_ticket(Ticket {
+ ticket: ticket.as_any().encode_to_vec().into(),
+ })
+ }
+
+ /// Serves a metadata command by round-tripping the command itself as the
+ /// ticket, so `DoGet` lands back on the matching handler.
+ fn metadata_info<C: ProstMessageExt>(
+ command: C,
+ schema: SchemaRef,
+ descriptor: FlightDescriptor,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let endpoint = FlightEndpoint::new().with_ticket(Ticket {
+ ticket: command.as_any().encode_to_vec().into(),
+ });
+ build_flight_info(&schema, vec![endpoint], descriptor).map(Response::new)
+ }
+}
+
+/// What the frontend should do with a planned statement.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Disposition {
+ /// Execute on the scheduler: it mutates session state and produces no data
+ /// worth distributing.
+ RunOnScheduler,
+ /// Submit to the cluster.
+ Distribute,
+ /// Refuse, with an explanation for the client.
+ Unsupported(&'static str),
+}
+
+/// Classifies a plan.
+///
+/// One function rather than two predicates, because the interesting cases are
+/// the ones where the answers overlap: `CREATE TABLE AS SELECT` is DDL, and
+/// DDL runs on the scheduler, so a caller that asked "is this DDL?" before
+/// asking "is this supported?" would silently execute its query on one node.
+/// Returning a single verdict makes that ordering impossible to get wrong.
+fn disposition(plan: &LogicalPlan) -> Disposition {
+ match plan {
+ LogicalPlan::Dml(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support INSERT/UPDATE/DELETE; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Copy(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support COPY; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support CREATE TABLE AS SELECT, \
+ because it would execute on the scheduler rather than the cluster; \
+ use CREATE EXTERNAL TABLE over data the executors can read",
+ ),
+ // Other DDL only edits the session catalog, and `SET`-style statements
+ // only edit session config; neither has anything to distribute.
+ LogicalPlan::Ddl(_) | LogicalPlan::Statement(_) => Disposition::RunOnScheduler,
+ // `SHOW ...` and anything else reading `information_schema` describes
+ // the catalog held here. There is nothing to distribute, and the
+ // physical form of those scans cannot be serialized for an executor,
+ // so distributing one would hand the client a job that never runs.
+ _ if scans_only_local_tables(plan) => Disposition::RunOnScheduler,
+ _ => Disposition::Distribute,
+ }
+}
+
+async fn execute_locally(
+ ctx: &SessionContext,
+ plan: LogicalPlan,
+) -> Result<(SchemaRef, Vec<RecordBatch>), Status> {
+ let df = ctx
+ .execute_logical_plan(plan)
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+ let planned_schema: SchemaRef = Arc::new(df.schema().as_arrow().clone());
+ let batches = df
+ .collect()
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+
+ // Prefer the schema the data actually carries; DataFusion's DDL results
+ // are empty and their frame schema is not always the same object.
+ let schema = batches
+ .first()
+ .map(|batch| batch.schema())
+ .unwrap_or(planned_schema);
+
+ Ok((schema, batches))
+}
+
+/// Turns a shuffle partition into the ticket payload the Flight proxy already
+/// knows how to redeem.
+fn partition_handle(
+ location: PartitionLocation,
+) -> Result<StatementHandle, BallistaError> {
+ let layout = location.layout();
+ let partition_id = location.partition_id.ok_or_else(|| {
+ BallistaError::Internal("partition location has no partition id".to_string())
+ })?;
+ let executor = location.executor_meta.ok_or_else(|| {
+ BallistaError::Internal("partition location has no executor metadata".to_string())
+ })?;
+
+ let action = BallistaAction::FetchPartition {
+ job_id: partition_id.job_id.into(),
+ stage_id: partition_id.stage_id as usize,
+ partition_id: partition_id.partition_id as usize,
+ host: executor.host,
+ port: executor.port as u16,
+ file_id: location.file_id,
+ layout,
+ // A Flight client wants the whole partition: the data file, in full.
+ file_kind: ShuffleFileKind::Data,
+ byte_ranges: vec![],
+ };
+
+ let encoded: protobuf::Action = action.try_into()?;
+ Ok(StatementHandle::Partition(encoded.encode_to_vec()))
+}
+
+fn build_flight_info(
+ schema: &Schema,
+ endpoints: Vec<FlightEndpoint>,
+ descriptor: FlightDescriptor,
+) -> Result<FlightInfo, Status> {
+ FlightInfo::new()
+ .try_with_schema(schema)
+ .map_err(|e| Status::internal(format!("failed to encode result schema: {e}")))
+ .map(|info| info.with_descriptor(descriptor).with_endpoints(endpoints))
+}
+
+/// Streams a single metadata batch back to the client.
+fn one_batch_response(batch: RecordBatch) -> Response<DoGetStream> {
+ batch_response(batch.schema(), vec![batch])
+}
+
+fn batch_response(schema: SchemaRef, batches: Vec<RecordBatch>) -> Response<DoGetStream> {
+ let stream = FlightDataEncoderBuilder::new()
+ .with_schema(schema)
+ .build(futures::stream::iter(batches.into_iter().map(Ok)))
+ .map_err(|e| Status::internal(format!("failed to encode results: {e}")));
+
+ Response::new(Box::pin(stream) as DoGetStream)
+}
+
+fn bearer_token(metadata: &MetadataMap) -> Option<String> {
+ let value = metadata.get("authorization")?.to_str().ok()?;
+ value
+ .strip_prefix("Bearer ")
+ .or_else(|| value.strip_prefix("bearer "))
+ .map(str::to_string)
+}
+
+fn invalid_ticket(e: impl std::fmt::Display) -> Status {
+ Status::invalid_argument(format!("invalid ticket: {e}"))
+}
+
+fn query_failed(e: BallistaError) -> Status {
+ // A failed job is the client's problem to see, not an opaque 500.
+ Status::internal(format!("query execution failed: {e}"))
+}
+
+fn encode_schema(schema: &Schema) -> Result<Vec<u8>, Status> {
+ let message: IpcMessage = SchemaAsIpc::new(schema, &IpcWriteOptions::default())
+ .try_into()
+ .map_err(|e| Status::internal(format!("failed to encode schema: {e}")))?;
+ Ok(message.0.to_vec())
+}
+
+#[tonic::async_trait]
+impl<B: QueryBackend> FlightSqlService for BallistaFlightSqlService<B> {
+ type FlightService = BallistaFlightSqlService<B>;
+
+ async fn do_handshake(
+ &self,
+ request: Request<Streaming<HandshakeRequest>>,
+ ) -> Result<
+ Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
+ Status,
+ > {
+ let identity = self.auth.authenticate(request.metadata()).await?;
+
+ let token = Uuid::new_v4().to_string();
+ let session_id = format!("flight-sql-{}", Uuid::new_v4());
+
+ // Build the session eagerly so a failure surfaces at handshake time
+ // rather than on the client's first query, and so the first query does
+ // not pay for it.
+ self.open_session(&session_id).await?;
+ self.store.insert_session(token.clone(), session_id.clone());
+
+ log::debug!(
+ "flight-sql: handshake for {:?} bound to session {session_id}",
+ identity.user
+ );
+
+ let result = HandshakeResponse {
+ protocol_version: 0,
+ payload: token.clone().into(),
+ };
+ let stream = futures::stream::once(async move { Ok(result) });
+
+ let mut response = Response::new(Box::pin(stream) as _);
+ response.metadata_mut().insert(
+ "authorization",
+ format!("Bearer {token}")
+ .parse()
+ .map_err(|_| Status::internal("failed to encode session token"))?,
+ );
+ Ok(response)
+ }
+
+ /// Redeems tickets that are not Flight SQL commands.
+ ///
+ /// Ballista's own Rust client fetches shuffle output with a
+ /// `ballista.protobuf.Action` ticket. When the Flight SQL frontend is
+ /// mounted it replaces the standalone proxy on the scheduler's port, so it
+ /// has to keep serving those tickets.
+ async fn do_get_fallback(
+ &self,
+ request: Request<Ticket>,
+ _message: Any,
+ ) -> Result<Response<DoGetStream>, Status> {
+ decode_protobuf(&request.get_ref().ticket).map_err(invalid_ticket)?;
+ self.proxy.do_get(request).await
+ }
+
+ async fn get_flight_info_statement(
+ &self,
+ query: CommandStatementQuery,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ let plan = Self::plan(&ctx, &query.query).await?;
+ let descriptor = request.into_inner();
+
+ self.flight_info_for(ctx, plan, descriptor, &job_name(&query.query))
+ .await
+ .map(Response::new)
+ }
+
+ async fn get_flight_info_prepared_statement(
+ &self,
+ query: CommandPreparedStatementQuery,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let handle = prepared_handle(&query.prepared_statement_handle)?;
+ let prepared = self.store.prepared(&handle).ok_or_else(|| {
+ Status::not_found("unknown or expired prepared statement handle")
+ })?;
+
+ let ctx = self.open_session(&prepared.session_id).await?;
+ let descriptor = request.into_inner();
+
+ self.flight_info_for(ctx, prepared.plan, descriptor, "flight-sql prepared")
+ .await
+ .map(Response::new)
+ }
+
+ async fn get_flight_info_catalogs(
+ &self,
+ query: CommandGetCatalogs,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_schemas(
+ &self,
+ query: CommandGetDbSchemas,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_tables(
+ &self,
+ query: CommandGetTables,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_table_types(
+ &self,
+ query: CommandGetTableTypes,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_sql_info(
+ &self,
+ query: CommandGetSqlInfo,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder(&self.sql_info).schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_xdbc_type_info(
+ &self,
+ query: CommandGetXdbcTypeInfo,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder(&self.xdbc_info).schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn do_get_statement(
Review Comment:
Every Flight SQL handler now checks the token, and prepared statements and local results are bound to the session that created them.
I didn't sign the partition tickets though. `do_get_fallback` has to serve raw `FetchPartition` tickets to native clients with no token, so signed Flight SQL tickets wouldn't stop anyone. I added a note in the code, and the docs now say the authenticator isolates sessions but doesn't secure the port.
I kept `Identity` for now. Handles are bound to the session rather than the identity, which also works for anonymous clients.
##########
ballista/scheduler/tests/flight_sql.rs:
##########
@@ -0,0 +1,339 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! End-to-end coverage for the Arrow Flight SQL frontend against a real
+//! (in-process) cluster: scheduler plus one executor, driven entirely through
+//! the Flight SQL protocol.
+//!
+//! The pre-46.0.0 implementation shipped with no tests at all, and the bugs
+//! that got it removed (#1012, #941, #839, #756) were exactly the ones an
+//! end-to-end test catches: endpoints pointing somewhere the client cannot
+//! reach, and results that fail to decode.
+
+#![cfg(feature = "flight-sql")]
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow_flight::sql::client::FlightSqlServiceClient;
+use arrow_flight::sql::{CommandGetTables, SqlInfo};
+use ballista_core::config::TaskSchedulingPolicy;
+use ballista_core::error::BallistaError;
+use ballista_core::serde::BallistaCodec;
+use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
+use ballista_core::utils::{
+ GrpcClientConfig, create_grpc_client_connection, default_config_producer,
+ default_session_builder,
+};
+use ballista_scheduler::cluster::BallistaCluster;
+use ballista_scheduler::config::SchedulerConfig;
+use ballista_scheduler::metrics::default_metrics_collector;
+use ballista_scheduler::scheduler_process::start_grpc_service_with_listener;
+use ballista_scheduler::scheduler_server::SchedulerServer;
+use datafusion::arrow::array::RecordBatch;
+use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+use futures::TryStreamExt;
+use tonic::transport::Channel;
+
+/// Test errors are only ever printed, and the client, transport, and Ballista
+/// error types all implement `Error` — so one boxed type spares every call site
+/// a `map_err`.
+type Result<T = ()> = std::result::Result<T, Box<dyn std::error::Error>>;
+
+/// Boots a scheduler serving Flight SQL plus one executor, and returns the
+/// scheduler's URL.
+async fn start_cluster() -> Result<String> {
Review Comment:
For the push path I went with a unit test instead. `launch_tasks_fails_jobs_whose_tasks_cannot_be_prepared` in `state/mod.rs` checks that the job is reported failed and its slots are freed, and it fails if the fix is reverted. After merging main the fix is also much smaller, since #2477 and #2016 already handle rejected jobs and refunds.
The test now uses `with_flight_sql(true)`. I couldn't reuse `new_standalone_scheduler_with_builder` because it hardcodes its `SchedulerConfig` and only mounts `SchedulerGrpc`, so there's no way to turn Flight SQL on through it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#discussion_r4088010457
##########
ballista/flight-sql/src/metadata.rs:
##########
@@ -0,0 +1,295 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Driver-facing metadata: server capabilities, type info, and catalog
+//! introspection.
+//!
+//! Catalog answers come from the session's real DataFusion catalog rather than
+//! hand-built record batches, so what a driver sees in its schema browser is
+//! what a query in the same session can actually reference.
+
+use arrow::array::RecordBatch;
+use arrow_flight::sql::metadata::{
+ GetCatalogsBuilder, GetDbSchemasBuilder, GetTablesBuilder, SqlInfoData,
+ SqlInfoDataBuilder, XdbcTypeInfo, XdbcTypeInfoData, XdbcTypeInfoDataBuilder,
+};
+use arrow_flight::sql::{
+ CommandGetCatalogs, CommandGetDbSchemas, CommandGetTableTypes, CommandGetTables,
+ Nullable, Searchable, SqlInfo, SqlSupportedTransaction, XdbcDataType,
+};
+use ballista_core::BALLISTA_VERSION;
+use datafusion::catalog::TableProvider;
+use datafusion::logical_expr::TableType;
+use datafusion::prelude::SessionContext;
+use tonic::Status;
+
+/// Table type strings reported to clients, matching the JDBC vocabulary that
+/// Flight SQL drivers expect.
+const TABLE: &str = "TABLE";
+const VIEW: &str = "VIEW";
+const LOCAL_TEMPORARY: &str = "LOCAL TEMPORARY";
+
+/// Describes the server to drivers deciding what SQL they may emit.
+///
+/// The old implementation left `CommandGetSqlInfo` unimplemented with a
+/// `// TODO: implement for FlightSQL JDBC to work` comment, which is precisely
+/// why the JDBC driver could not connect.
+pub(crate) fn sql_info() -> SqlInfoData {
+ let mut builder = SqlInfoDataBuilder::new();
+
+ builder.append(SqlInfo::FlightSqlServerName, "Apache DataFusion Ballista");
+ builder.append(SqlInfo::FlightSqlServerVersion, BALLISTA_VERSION);
+ // Arrow IPC format version, per format/Schema.fbs.
+ builder.append(SqlInfo::FlightSqlServerArrowVersion, "1.3");
+ builder.append(SqlInfo::FlightSqlServerReadOnly, false);
+ builder.append(SqlInfo::FlightSqlServerSql, true);
+ builder.append(SqlInfo::FlightSqlServerSubstrait, false);
+ builder.append(
+ SqlInfo::FlightSqlServerTransaction,
+ SqlSupportedTransaction::None as i32,
+ );
+ builder.append(SqlInfo::FlightSqlServerCancel, true);
+ builder.append(SqlInfo::FlightSqlServerBulkIngestion, false);
+
+ builder.append(SqlInfo::SqlDdlCatalog, false);
+ builder.append(SqlInfo::SqlDdlSchema, true);
+ builder.append(SqlInfo::SqlDdlTable, true);
+ // DataFusion folds unquoted identifiers to lower case and preserves the
+ // case of quoted ones: SQL_CASE_SENSITIVITY_LOWERCASE / _CASE_INSENSITIVE.
+ builder.append(SqlInfo::SqlIdentifierCase, 2i32);
Review Comment:
**Major.** These numbers don't match the comment. In `SqlSupportedCaseSensitivity`, 2 is `SQL_CASE_SENSITIVITY_UPPERCASE` and 3 is `SQL_CASE_SENSITIVITY_LOWERCASE`. So this reports unquoted identifiers as stored in upper case and quoted ones as lower case, and clients that read these values will get DataFusion's identifier rules backwards.
`arrow_flight::sql` re-exports the enum, which makes the intent explicit:
```rust
builder.append(
SqlInfo::SqlIdentifierCase,
SqlSupportedCaseSensitivity::SqlCaseSensitivityLowercase as i32,
);
```
Same for `SqlQuotedIdentifierCase` on line 75, with the value picked deliberately.
##########
ballista/core/src/planner.rs:
##########
@@ -165,6 +162,19 @@ impl<T: 'static + AsLogicalPlan> QueryPlanner for BallistaQueryPlanner<T> {
}
}
+/// Returns `true` if every table `plan` scans lives in `information_schema`.
+///
+/// Such a plan reads the catalog of the node it is planned on and has nothing
+/// to distribute; worse, the physical form of those scans (`StreamingTableExec`)
+/// cannot be serialized for an executor, so distributing one produces a job
+/// that can never run. Callers that hand plans to the cluster should answer
+/// these locally instead.
+pub fn scans_only_local_tables(plan: &LogicalPlan) -> bool {
+ let mut local_run = LocalRun::default();
+ let _ = plan.visit(&mut local_run);
Review Comment:
**Major.** `visit` doesn't descend into subquery expressions, while `visit_with_subqueries` does (DataFusion 55, `logical_plan/tree_node.rs`). The Flight SQL frontend calls this on the unoptimized plan from `create_logical_plan`, so for
```sql
SELECT table_name FROM information_schema.tables
WHERE table_name IN (SELECT name FROM big)
```
the only scan it sees is `information_schema.tables`. The statement is classified `RunOnScheduler` and `big` is scanned on the scheduler. `BallistaQueryPlanner` receives the optimized plan, where most subqueries are already decorrelated into joins, which mostly hides this there. `plan.visit_with_subqueries(&mut local_run)` fixes both callers.
##########
ballista/scheduler/src/config.rs:
##########
@@ -397,6 +405,19 @@ pub struct SchedulerConfig {
pub task_max_failures: usize,
/// Number of failures attempts before stage is considered failed
pub stage_max_failures: usize,
+ /// Whether to serve Arrow Flight SQL on the scheduler's gRPC port.
+ ///
+ /// The frontend subsumes the plain Flight result proxy when enabled, so
+ /// Ballista's own clients keep working on the same port.
+ #[cfg(feature = "flight-sql")]
+ pub flight_sql: bool,
+ /// Authenticates Arrow Flight SQL handshakes.
+ ///
+ /// Without one the frontend accepts every client, so this is the only way
Review Comment:
**Major (docs).** This overstates what the authenticator gives. The same port also serves `SchedulerGrpc` and the REST API (including `PATCH /api/job/{id}`) with no authentication, and several Flight SQL paths skip the token check (see the comment on `do_get_statement`). Suggest describing it as per-client session isolation and keeping the "trusted network only" guidance from the user guide.
##########
ballista/core/src/serde/scheduler/mod.rs:
##########
@@ -71,11 +71,27 @@ impl PartitionLocation {
/// vocabulary of one writer. The stored field keeps that name until the
/// construction sites are swept; the wire protocol already speaks layouts.
pub fn layout(&self) -> ShuffleLayout {
- if self.is_sort_shuffle {
- ShuffleLayout::Sort
- } else {
- ShuffleLayout::Passthrough
- }
+ shuffle_layout(self.is_sort_shuffle)
+ }
+}
+
+impl crate::serde::protobuf::PartitionLocation {
Review Comment:
**Minor (reuse).** The only caller is `partition_handle` in the new crate, while `fetch_partition` in `execution_plans/distributed_query.rs` (line 899) still inlines the same `is_sort_shuffle` rule on the same protobuf type. Either use this there too, or add one constructor for `Action::FetchPartition` from a `PartitionLocation` that both `partition_handle` and `BallistaClient::fetch_partition_proxied` use.
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
+ return Err(Status::unimplemented(reason));
+ }
+
+ if disposition(&plan) == Disposition::RunOnScheduler {
+ let (schema, batches) = execute_locally(&ctx, plan).await?;
+ let handle = Uuid::new_v4().to_string();
+ self.store.insert_result(
+ handle.clone(),
+ LocalResult {
+ schema: schema.clone(),
+ batches,
+ },
+ );
+
+ let endpoint = Self::endpoint(StatementHandle::Local(handle));
+ return build_flight_info(&schema, vec![endpoint], descriptor);
+ }
+
+ let result = self
+ .backend
+ .execute(job_name, ctx, plan)
+ .await
+ .map_err(query_failed)?;
+
+ let endpoints = result
+ .partitions
+ .into_iter()
+ .map(|location| partition_handle(location).map(Self::endpoint))
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| Status::internal(format!("invalid partition location: {e}")))?;
+
+ log::debug!(
+ "flight-sql: job {} produced {} endpoint(s)",
+ result.job_id,
+ endpoints.len()
+ );
+
+ build_flight_info(&result.schema, endpoints, descriptor)
+ }
+
+ /// Builds the endpoint a client redeems for one slice of the result.
+ ///
+ /// Endpoints carry no location, which Flight defines as "fetch from the
+ /// server that gave you this FlightInfo". That keeps every cluster-internal
+ /// address off the wire and lets the frontend work unchanged behind NAT, a
+ /// load balancer, or an ingress.
+ fn endpoint(handle: StatementHandle) -> FlightEndpoint {
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.encode().into(),
+ };
+ FlightEndpoint::new().with_ticket(Ticket {
+ ticket: ticket.as_any().encode_to_vec().into(),
+ })
+ }
+
+ /// Serves a metadata command by round-tripping the command itself as the
+ /// ticket, so `DoGet` lands back on the matching handler.
+ fn metadata_info<C: ProstMessageExt>(
+ command: C,
+ schema: SchemaRef,
+ descriptor: FlightDescriptor,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let endpoint = FlightEndpoint::new().with_ticket(Ticket {
+ ticket: command.as_any().encode_to_vec().into(),
+ });
+ build_flight_info(&schema, vec![endpoint], descriptor).map(Response::new)
+ }
+}
+
+/// What the frontend should do with a planned statement.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Disposition {
+ /// Execute on the scheduler: it mutates session state and produces no data
+ /// worth distributing.
+ RunOnScheduler,
+ /// Submit to the cluster.
+ Distribute,
+ /// Refuse, with an explanation for the client.
+ Unsupported(&'static str),
+}
+
+/// Classifies a plan.
+///
+/// One function rather than two predicates, because the interesting cases are
+/// the ones where the answers overlap: `CREATE TABLE AS SELECT` is DDL, and
+/// DDL runs on the scheduler, so a caller that asked "is this DDL?" before
+/// asking "is this supported?" would silently execute its query on one node.
+/// Returning a single verdict makes that ordering impossible to get wrong.
+fn disposition(plan: &LogicalPlan) -> Disposition {
+ match plan {
+ LogicalPlan::Dml(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support INSERT/UPDATE/DELETE; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Copy(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support COPY; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support CREATE TABLE AS SELECT, \
+ because it would execute on the scheduler rather than the cluster; \
+ use CREATE EXTERNAL TABLE over data the executors can read",
+ ),
+ // Other DDL only edits the session catalog, and `SET`-style statements
+ // only edit session config; neither has anything to distribute.
+ LogicalPlan::Ddl(_) | LogicalPlan::Statement(_) => Disposition::RunOnScheduler,
Review Comment:
**Major.** `disposition` only looks at the root node, so a few statements still run on the scheduler or get past the write refusal:
- `EXECUTE q` is a `LogicalPlan::Statement`, so it lands here. In DataFusion 55, `execute_logical_plan` turns it into the prepared query (`execute_prepared`), and `execute_locally` then runs that `SELECT` on the scheduler (the default `SessionBuilder` uses DataFusion's local planner) and keeps the whole result in memory as a `LocalResult`. That is the same reason CTAS is refused.
- `EXPLAIN ANALYZE INSERT ...` plans as `Analyze` over `Dml` (`explain_to_plan` plans any statement), so it is classified `Distribute` instead of being refused.
Suggest refusing `Statement::Execute` and checking the whole plan, subqueries included, for `Dml`, `Copy` and `CreateMemoryTable`. Related: the comment on `scans_only_local_tables`.
##########
ballista/flight-sql/src/session.rs:
##########
@@ -0,0 +1,277 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Server-side state the Flight SQL frontend keeps between requests.
+//!
+//! Everything in here is keyed by an opaque handle the client is given and
+//! hands back, and everything expires. A client that disconnects without
+//! closing its prepared statements (or that never redeems a ticket) must not
+//! pin memory forever, which is what the previous Flight SQL implementation
+//! got wrong: its plan cache only shed entries on an explicit close.
+
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::SchemaRef;
+use dashmap::DashMap;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+
+/// A prepared statement's server-side state.
+#[derive(Clone)]
+pub(crate) struct Prepared {
+ /// Session the statement was prepared in; it is planned and executed
+ /// against that session's catalog.
+ pub session_id: String,
+ /// The plan as prepared; execution never re-plans it.
+ pub plan: LogicalPlan,
+}
+
+/// Result of a statement the frontend ran locally rather than distributing
+/// (DDL, and DML whose only output is an affected-row count).
+#[derive(Clone)]
+pub(crate) struct LocalResult {
+ pub schema: SchemaRef,
+ pub batches: Vec<RecordBatch>,
+}
+
+struct Tracked<T> {
+ value: T,
+ last_used: Instant,
+}
+
+impl<T> Tracked<T> {
+ fn new(value: T) -> Self {
+ Self {
+ value,
+ last_used: Instant::now(),
+ }
+ }
+}
+
+/// Handle stores for sessions, prepared statements, and locally-computed
+/// results, all with a shared idle TTL.
+pub(crate) struct SessionStore {
+ ttl: Duration,
+ /// Bearer token -> Ballista session id.
+ sessions: DashMap<String, Tracked<String>>,
+ /// Ballista session id -> its context.
+ contexts: DashMap<String, Tracked<Arc<SessionContext>>>,
+ /// Prepared statement handle -> prepared plan.
+ prepared: DashMap<String, Tracked<Prepared>>,
+ /// Local result handle -> materialized batches.
+ results: DashMap<String, Tracked<LocalResult>>,
+}
+
+impl SessionStore {
+ pub(crate) fn new(ttl: Duration) -> Self {
+ Self {
+ ttl,
+ sessions: DashMap::new(),
+ contexts: DashMap::new(),
+ prepared: DashMap::new(),
+ results: DashMap::new(),
+ }
+ }
+
+ pub(crate) fn insert_session(&self, token: String, session_id: String) {
+ self.sessions.insert(token, Tracked::new(session_id));
+ }
+
+ /// Resolves a bearer token to its session id, refreshing its idle timer.
+ pub(crate) fn session(&self, token: &str) -> Option<String> {
+ self.sessions.get_mut(token).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ /// Drops a token and the session it named, returning that session id.
+ ///
+ /// Each token gets its own session, so removing the token always releases
+ /// the session with it.
+ pub(crate) fn remove_session(&self, token: &str) -> Option<String> {
+ let (_, entry) = self.sessions.remove(token)?;
+ self.contexts.remove(&entry.value);
+ Some(entry.value)
+ }
+
+ /// Returns the cached context for a session, refreshing its idle timer.
+ pub(crate) fn context(&self, session_id: &str) -> Option<Arc<SessionContext>> {
+ self.contexts.get_mut(session_id).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ /// Caches a context, or returns the one another caller cached first.
+ ///
+ /// Two requests for a cold session can both build one; letting the first
+ /// insertion win keeps them looking at the same catalog.
+ pub(crate) fn insert_context(
+ &self,
+ session_id: String,
+ ctx: Arc<SessionContext>,
+ ) -> Arc<SessionContext> {
+ self.contexts
+ .entry(session_id)
+ .or_insert_with(|| Tracked::new(ctx))
+ .value
+ .clone()
+ }
+
+ pub(crate) fn insert_prepared(&self, handle: String, prepared: Prepared) {
+ self.prepared.insert(handle, Tracked::new(prepared));
+ }
+
+ pub(crate) fn prepared(&self, handle: &str) -> Option<Prepared> {
+ self.prepared.get_mut(handle).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ pub(crate) fn remove_prepared(&self, handle: &str) {
+ self.prepared.remove(handle);
+ }
+
+ pub(crate) fn insert_result(&self, handle: String, result: LocalResult) {
+ self.results.insert(handle, Tracked::new(result));
+ }
+
+ /// Takes a local result. Results are single-use: a ticket is redeemed once.
+ pub(crate) fn take_result(&self, handle: &str) -> Option<LocalResult> {
+ self.results.remove(handle).map(|(_, entry)| entry.value)
+ }
+
+ /// Evicts everything idle for longer than the TTL.
+ ///
+ /// Returns the session ids that no longer have any live token, so the
+ /// caller can release them in the backend.
+ pub(crate) fn sweep(&self) -> Vec<String> {
Review Comment:
**Minor.** This iterates `sessions` to collect the expired tokens and then removes them one by one. A single `retain` pass can collect the released session ids as it goes, and their contexts can be dropped afterwards. That also leaves `remove_session` used only by a test.
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
+ return Err(Status::unimplemented(reason));
+ }
+
+ if disposition(&plan) == Disposition::RunOnScheduler {
+ let (schema, batches) = execute_locally(&ctx, plan).await?;
+ let handle = Uuid::new_v4().to_string();
+ self.store.insert_result(
+ handle.clone(),
+ LocalResult {
+ schema: schema.clone(),
+ batches,
+ },
+ );
+
+ let endpoint = Self::endpoint(StatementHandle::Local(handle));
+ return build_flight_info(&schema, vec![endpoint], descriptor);
+ }
+
+ let result = self
+ .backend
+ .execute(job_name, ctx, plan)
+ .await
+ .map_err(query_failed)?;
+
+ let endpoints = result
+ .partitions
+ .into_iter()
+ .map(|location| partition_handle(location).map(Self::endpoint))
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| Status::internal(format!("invalid partition location: {e}")))?;
+
+ log::debug!(
+ "flight-sql: job {} produced {} endpoint(s)",
+ result.job_id,
+ endpoints.len()
+ );
+
+ build_flight_info(&result.schema, endpoints, descriptor)
+ }
+
+ /// Builds the endpoint a client redeems for one slice of the result.
+ ///
+ /// Endpoints carry no location, which Flight defines as "fetch from the
+ /// server that gave you this FlightInfo". That keeps every cluster-internal
+ /// address off the wire and lets the frontend work unchanged behind NAT, a
+ /// load balancer, or an ingress.
+ fn endpoint(handle: StatementHandle) -> FlightEndpoint {
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.encode().into(),
+ };
+ FlightEndpoint::new().with_ticket(Ticket {
+ ticket: ticket.as_any().encode_to_vec().into(),
+ })
+ }
+
+ /// Serves a metadata command by round-tripping the command itself as the
+ /// ticket, so `DoGet` lands back on the matching handler.
+ fn metadata_info<C: ProstMessageExt>(
+ command: C,
+ schema: SchemaRef,
+ descriptor: FlightDescriptor,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let endpoint = FlightEndpoint::new().with_ticket(Ticket {
+ ticket: command.as_any().encode_to_vec().into(),
+ });
+ build_flight_info(&schema, vec![endpoint], descriptor).map(Response::new)
+ }
+}
+
+/// What the frontend should do with a planned statement.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Disposition {
+ /// Execute on the scheduler: it mutates session state and produces no data
+ /// worth distributing.
+ RunOnScheduler,
+ /// Submit to the cluster.
+ Distribute,
+ /// Refuse, with an explanation for the client.
+ Unsupported(&'static str),
+}
+
+/// Classifies a plan.
+///
+/// One function rather than two predicates, because the interesting cases are
+/// the ones where the answers overlap: `CREATE TABLE AS SELECT` is DDL, and
+/// DDL runs on the scheduler, so a caller that asked "is this DDL?" before
+/// asking "is this supported?" would silently execute its query on one node.
+/// Returning a single verdict makes that ordering impossible to get wrong.
+fn disposition(plan: &LogicalPlan) -> Disposition {
+ match plan {
+ LogicalPlan::Dml(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support INSERT/UPDATE/DELETE; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Copy(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support COPY; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support CREATE TABLE AS SELECT, \
+ because it would execute on the scheduler rather than the cluster; \
+ use CREATE EXTERNAL TABLE over data the executors can read",
+ ),
+ // Other DDL only edits the session catalog, and `SET`-style statements
+ // only edit session config; neither has anything to distribute.
+ LogicalPlan::Ddl(_) | LogicalPlan::Statement(_) => Disposition::RunOnScheduler,
+ // `SHOW ...` and anything else reading `information_schema` describes
+ // the catalog held here. There is nothing to distribute, and the
+ // physical form of those scans cannot be serialized for an executor,
+ // so distributing one would hand the client a job that never runs.
+ _ if scans_only_local_tables(plan) => Disposition::RunOnScheduler,
+ _ => Disposition::Distribute,
+ }
+}
+
+async fn execute_locally(
+ ctx: &SessionContext,
+ plan: LogicalPlan,
+) -> Result<(SchemaRef, Vec<RecordBatch>), Status> {
+ let df = ctx
+ .execute_logical_plan(plan)
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+ let planned_schema: SchemaRef = Arc::new(df.schema().as_arrow().clone());
+ let batches = df
+ .collect()
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+
+ // Prefer the schema the data actually carries; DataFusion's DDL results
+ // are empty and their frame schema is not always the same object.
+ let schema = batches
+ .first()
+ .map(|batch| batch.schema())
+ .unwrap_or(planned_schema);
+
+ Ok((schema, batches))
+}
+
+/// Turns a shuffle partition into the ticket payload the Flight proxy already
+/// knows how to redeem.
+fn partition_handle(
+ location: PartitionLocation,
+) -> Result<StatementHandle, BallistaError> {
+ let layout = location.layout();
+ let partition_id = location.partition_id.ok_or_else(|| {
+ BallistaError::Internal("partition location has no partition id".to_string())
+ })?;
+ let executor = location.executor_meta.ok_or_else(|| {
+ BallistaError::Internal("partition location has no executor metadata".to_string())
+ })?;
+
+ let action = BallistaAction::FetchPartition {
+ job_id: partition_id.job_id.into(),
+ stage_id: partition_id.stage_id as usize,
+ partition_id: partition_id.partition_id as usize,
+ host: executor.host,
+ port: executor.port as u16,
+ file_id: location.file_id,
+ layout,
+ // A Flight client wants the whole partition: the data file, in full.
+ file_kind: ShuffleFileKind::Data,
+ byte_ranges: vec![],
+ };
+
+ let encoded: protobuf::Action = action.try_into()?;
+ Ok(StatementHandle::Partition(encoded.encode_to_vec()))
+}
+
+fn build_flight_info(
+ schema: &Schema,
+ endpoints: Vec<FlightEndpoint>,
+ descriptor: FlightDescriptor,
+) -> Result<FlightInfo, Status> {
+ FlightInfo::new()
+ .try_with_schema(schema)
+ .map_err(|e| Status::internal(format!("failed to encode result schema: {e}")))
+ .map(|info| info.with_descriptor(descriptor).with_endpoints(endpoints))
+}
+
+/// Streams a single metadata batch back to the client.
+fn one_batch_response(batch: RecordBatch) -> Response<DoGetStream> {
+ batch_response(batch.schema(), vec![batch])
+}
+
+fn batch_response(schema: SchemaRef, batches: Vec<RecordBatch>) -> Response<DoGetStream> {
+ let stream = FlightDataEncoderBuilder::new()
+ .with_schema(schema)
+ .build(futures::stream::iter(batches.into_iter().map(Ok)))
+ .map_err(|e| Status::internal(format!("failed to encode results: {e}")));
+
+ Response::new(Box::pin(stream) as DoGetStream)
+}
+
+fn bearer_token(metadata: &MetadataMap) -> Option<String> {
+ let value = metadata.get("authorization")?.to_str().ok()?;
+ value
+ .strip_prefix("Bearer ")
+ .or_else(|| value.strip_prefix("bearer "))
+ .map(str::to_string)
+}
+
+fn invalid_ticket(e: impl std::fmt::Display) -> Status {
+ Status::invalid_argument(format!("invalid ticket: {e}"))
+}
+
+fn query_failed(e: BallistaError) -> Status {
+ // A failed job is the client's problem to see, not an opaque 500.
+ Status::internal(format!("query execution failed: {e}"))
+}
+
+fn encode_schema(schema: &Schema) -> Result<Vec<u8>, Status> {
+ let message: IpcMessage = SchemaAsIpc::new(schema, &IpcWriteOptions::default())
+ .try_into()
+ .map_err(|e| Status::internal(format!("failed to encode schema: {e}")))?;
+ Ok(message.0.to_vec())
+}
+
+#[tonic::async_trait]
+impl<B: QueryBackend> FlightSqlService for BallistaFlightSqlService<B> {
+ type FlightService = BallistaFlightSqlService<B>;
+
+ async fn do_handshake(
+ &self,
+ request: Request<Streaming<HandshakeRequest>>,
+ ) -> Result<
+ Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
+ Status,
+ > {
+ let identity = self.auth.authenticate(request.metadata()).await?;
+
+ let token = Uuid::new_v4().to_string();
+ let session_id = format!("flight-sql-{}", Uuid::new_v4());
+
+ // Build the session eagerly so a failure surfaces at handshake time
+ // rather than on the client's first query, and so the first query does
+ // not pay for it.
+ self.open_session(&session_id).await?;
+ self.store.insert_session(token.clone(), session_id.clone());
+
+ log::debug!(
+ "flight-sql: handshake for {:?} bound to session {session_id}",
+ identity.user
+ );
+
+ let result = HandshakeResponse {
+ protocol_version: 0,
+ payload: token.clone().into(),
+ };
+ let stream = futures::stream::once(async move { Ok(result) });
+
+ let mut response = Response::new(Box::pin(stream) as _);
+ response.metadata_mut().insert(
+ "authorization",
+ format!("Bearer {token}")
+ .parse()
+ .map_err(|_| Status::internal("failed to encode session token"))?,
+ );
+ Ok(response)
+ }
+
+ /// Redeems tickets that are not Flight SQL commands.
+ ///
+ /// Ballista's own Rust client fetches shuffle output with a
+ /// `ballista.protobuf.Action` ticket. When the Flight SQL frontend is
+ /// mounted it replaces the standalone proxy on the scheduler's port, so it
+ /// has to keep serving those tickets.
+ async fn do_get_fallback(
+ &self,
+ request: Request<Ticket>,
+ _message: Any,
+ ) -> Result<Response<DoGetStream>, Status> {
+ decode_protobuf(&request.get_ref().ticket).map_err(invalid_ticket)?;
+ self.proxy.do_get(request).await
+ }
+
+ async fn get_flight_info_statement(
+ &self,
+ query: CommandStatementQuery,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ let plan = Self::plan(&ctx, &query.query).await?;
+ let descriptor = request.into_inner();
+
+ self.flight_info_for(ctx, plan, descriptor, &job_name(&query.query))
+ .await
+ .map(Response::new)
+ }
+
+ async fn get_flight_info_prepared_statement(
+ &self,
+ query: CommandPreparedStatementQuery,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let handle = prepared_handle(&query.prepared_statement_handle)?;
+ let prepared = self.store.prepared(&handle).ok_or_else(|| {
+ Status::not_found("unknown or expired prepared statement handle")
+ })?;
+
+ let ctx = self.open_session(&prepared.session_id).await?;
+ let descriptor = request.into_inner();
+
+ self.flight_info_for(ctx, prepared.plan, descriptor, "flight-sql prepared")
+ .await
+ .map(Response::new)
+ }
+
+ async fn get_flight_info_catalogs(
+ &self,
+ query: CommandGetCatalogs,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_schemas(
+ &self,
+ query: CommandGetDbSchemas,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_tables(
+ &self,
+ query: CommandGetTables,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_table_types(
+ &self,
+ query: CommandGetTableTypes,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_sql_info(
+ &self,
+ query: CommandGetSqlInfo,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder(&self.sql_info).schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_xdbc_type_info(
+ &self,
+ query: CommandGetXdbcTypeInfo,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder(&self.xdbc_info).schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn do_get_statement(
+ &self,
+ ticket: TicketStatementQuery,
+ _request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let handle =
+ StatementHandle::decode(&ticket.statement_handle).map_err(invalid_ticket)?;
+
+ match handle {
+ StatementHandle::Partition(action) => {
+ // Unwrap back to the executor-facing ticket and let the proxy
+ // redeem it. The proxy dials the executor with a fresh request,
+ // so the client's credentials are not forwarded onwards.
+ self.proxy
+ .do_get(Request::new(Ticket {
+ ticket: action.into(),
+ }))
+ .await
+ }
+ StatementHandle::Local(handle) => {
+ let result = self.store.take_result(&handle).ok_or_else(|| {
+ Status::not_found("result already consumed or expired")
+ })?;
+ Ok(batch_response(result.schema, result.batches))
+ }
+ }
+ }
+
+ async fn do_get_catalogs(
+ &self,
+ query: CommandGetCatalogs,
+ request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ Ok(one_batch_response(metadata::catalogs(&ctx, query)?))
+ }
+
+ async fn do_get_schemas(
+ &self,
+ query: CommandGetDbSchemas,
+ request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ Ok(one_batch_response(metadata::db_schemas(&ctx, query)?))
+ }
+
+ async fn do_get_tables(
+ &self,
+ query: CommandGetTables,
+ request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ Ok(one_batch_response(metadata::tables(&ctx, query).await?))
+ }
+
+ async fn do_get_table_types(
+ &self,
+ query: CommandGetTableTypes,
+ _request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ Ok(one_batch_response(metadata::table_types(query)?))
+ }
+
+ async fn do_get_sql_info(
+ &self,
+ query: CommandGetSqlInfo,
+ _request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let batch = query
+ .into_builder(&self.sql_info)
+ .build()
+ .map_err(|e| Status::internal(format!("failed to build SqlInfo: {e}")))?;
+ Ok(one_batch_response(batch))
+ }
+
+ async fn do_get_xdbc_type_info(
+ &self,
+ query: CommandGetXdbcTypeInfo,
+ _request: Request<Ticket>,
+ ) -> Result<Response<DoGetStream>, Status> {
+ let batch = query
+ .into_builder(&self.xdbc_info)
+ .build()
+ .map_err(|e| Status::internal(format!("failed to build type info: {e}")))?;
+ Ok(one_batch_response(batch))
+ }
+
+ async fn do_put_statement_update(
+ &self,
+ ticket: CommandStatementUpdate,
+ request: Request<PeekableFlightDataStream>,
+ ) -> Result<i64, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ let plan = Self::plan(&ctx, &ticket.query).await?;
+
+ match disposition(&plan) {
+ Disposition::RunOnScheduler => {}
+ Disposition::Unsupported(reason) => {
+ return Err(Status::unimplemented(reason));
+ }
+ Disposition::Distribute => {
+ return Err(Status::unimplemented(
+ "Ballista Flight SQL supports DDL and session statements on this \
+ path; the distributed write path is not implemented",
+ ));
+ }
+ }
+
+ execute_locally(&ctx, plan).await?;
+ // Only DDL and session statements reach here, and neither reports an
+ // affected-row count.
+ Ok(0)
+ }
+
+ async fn do_action_create_prepared_statement(
+ &self,
+ query: ActionCreatePreparedStatementRequest,
+ request: Request<Action>,
+ ) -> Result<ActionCreatePreparedStatementResult, Status> {
+ let session_id = self.session_id(request.metadata())?;
+ let ctx = self.open_session(&session_id).await?;
+ let plan = Self::plan(&ctx, &query.query).await?;
+ let schema = plan.schema().as_arrow().clone();
+
+ let handle = Uuid::new_v4().to_string();
+ self.store
+ .insert_prepared(handle.clone(), Prepared { session_id, plan });
+
+ Ok(ActionCreatePreparedStatementResult {
+ prepared_statement_handle: handle.into_bytes().into(),
+ dataset_schema: encode_schema(&schema)?.into(),
+ // Bound parameters are not supported yet, so the parameter schema
+ // is empty rather than absent: clients read "no parameters".
+ parameter_schema: encode_schema(&Schema::empty())?.into(),
+ })
+ }
+
+ async fn do_action_close_prepared_statement(
+ &self,
+ query: ActionClosePreparedStatementRequest,
+ _request: Request<Action>,
+ ) -> Result<(), Status> {
+ let handle = prepared_handle(&query.prepared_statement_handle)?;
+ self.store.remove_prepared(&handle);
+ Ok(())
+ }
+
+ async fn do_action_cancel_query(
Review Comment:
**Blocker.** This can never cancel a running query, and calling it deletes the results the client was just given.
- A client only holds the `FlightInfo` after `GetFlightInfo` returns, so by then the job has already succeeded (design note 7).
- `cancel_job` posts `JobCancel`, and that handler always calls `clean_up_failed_job`. That removes the job's data on the executors without the `finished_job_data_clean_up_interval_seconds` delay, so the endpoints in that `FlightInfo` stop working.
- `result: 1` is `CANCEL_RESULT_CANCELLED`. For a cancel that completes in the background, the proto asks for `CANCEL_RESULT_CANCELLING` (2).
- `ActionCancelQueryRequest` has been deprecated since Flight SQL 13.0.0 in favor of `CancelFlightInfo`, which arrow-flight 59.2.0 routes to `do_action_fallback`.
Suggest removing `do_action_cancel_query`, `job_id_from_flight_info`, `QueryBackend::cancel` and `cancelling_a_query_reaches_the_backend`, reporting `FlightSqlServerCancel` as `false` (`metadata.rs:65`), and dropping "query cancellation" from the crate README. Cancel could come back together with `PollFlightInfo`.
##########
ballista/scheduler/tests/flight_sql.rs:
##########
@@ -0,0 +1,339 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! End-to-end coverage for the Arrow Flight SQL frontend against a real
+//! (in-process) cluster: scheduler plus one executor, driven entirely through
+//! the Flight SQL protocol.
+//!
+//! The pre-46.0.0 implementation shipped with no tests at all, and the bugs
+//! that got it removed (#1012, #941, #839, #756) were exactly the ones an
+//! end-to-end test catches: endpoints pointing somewhere the client cannot
+//! reach, and results that fail to decode.
+
+#![cfg(feature = "flight-sql")]
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow_flight::sql::client::FlightSqlServiceClient;
+use arrow_flight::sql::{CommandGetTables, SqlInfo};
+use ballista_core::config::TaskSchedulingPolicy;
+use ballista_core::error::BallistaError;
+use ballista_core::serde::BallistaCodec;
+use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
+use ballista_core::utils::{
+ GrpcClientConfig, create_grpc_client_connection, default_config_producer,
+ default_session_builder,
+};
+use ballista_scheduler::cluster::BallistaCluster;
+use ballista_scheduler::config::SchedulerConfig;
+use ballista_scheduler::metrics::default_metrics_collector;
+use ballista_scheduler::scheduler_process::start_grpc_service_with_listener;
+use ballista_scheduler::scheduler_server::SchedulerServer;
+use datafusion::arrow::array::RecordBatch;
+use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+use futures::TryStreamExt;
+use tonic::transport::Channel;
+
+/// Test errors are only ever printed, and the client, transport, and Ballista
+/// error types all implement `Error` — so one boxed type spares every call site
+/// a `map_err`.
+type Result<T = ()> = std::result::Result<T, Box<dyn std::error::Error>>;
+
+/// Boots a scheduler serving Flight SQL plus one executor, and returns the
+/// scheduler's URL.
+async fn start_cluster() -> Result<String> {
+ let mut config = SchedulerConfig::default()
+ .with_scheduler_policy(TaskSchedulingPolicy::PullStaged);
+ config.flight_sql = true;
+
+ let cluster = BallistaCluster::new_memory(
+ "localhost:50050",
+ Arc::new(default_session_builder),
+ Arc::new(default_config_producer),
+ );
+
+ let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerServer::new(
+ "localhost:50050".to_owned(),
+ cluster,
+ BallistaCodec::default(),
+ Arc::new(config),
+ default_metrics_collector()?,
+ );
+ scheduler.init().await?;
+
+ // Bind first so the test learns the port without racing another process
+ // for it.
+ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+ let addr = listener.local_addr()?;
+ tokio::spawn(start_grpc_service_with_listener(listener, scheduler));
+
+ let url = format!("http://{addr}");
+
+ let scheduler_client = connect_with_retry(&url).await?;
+ ballista_executor::new_standalone_executor(
+ scheduler_client,
+ 2,
+ BallistaCodec::default(),
+ )
+ .await?;
+
+ Ok(url)
+}
+
+async fn connect_with_retry(url: &str) -> Result<SchedulerGrpcClient<Channel>> {
+ for _ in 0..100 {
+ if let Ok(client) = SchedulerGrpcClient::connect(url.to_string()).await {
+ return Ok(client);
+ }
+ tokio::time::sleep(Duration::from_millis(50)).await;
+ }
+ Err(BallistaError::General(format!("scheduler at {url} did not come up")).into())
+}
+
+/// Connects a Flight SQL client, authenticates, and registers a `people` table
+/// for it to query.
+///
+/// The `TempDir` is returned because it owns the CSV the table points at; drop
+/// it and the executors have nothing to read.
+async fn client_with_people(
+ url: &str,
+) -> Result<(FlightSqlServiceClient<Channel>, tempfile::TempDir)> {
+ let channel =
+ create_grpc_client_connection(url.to_string(), &GrpcClientConfig::default())
+ .await?;
+ let mut client = FlightSqlServiceClient::new(channel);
+ client.handshake("anonymous", "").await?;
+
+ let csv = tempfile::tempdir()?;
+ let path = csv.path().join("data.csv");
+ std::fs::write(&path, "id,name\n1,alice\n2,bob\n3,carol\n")?;
+ let path = path.to_string_lossy();
+
+ // DDL runs on the scheduler and registers into the session catalog.
+ let info = client
+ .execute(
+ format!(
+ "CREATE EXTERNAL TABLE people STORED AS CSV LOCATION '{path}' \
+ OPTIONS ('format.has_header' 'true')"
+ ),
+ None,
+ )
+ .await?;
+ collect(&mut client, info).await?;
+
+ Ok((client, csv))
+}
+
+/// Collects every endpoint of a `FlightInfo`, which is what a real client does
+/// and what proves the tickets we hand out are redeemable.
+async fn collect(
+ client: &mut FlightSqlServiceClient<Channel>,
+ info: arrow_flight::FlightInfo,
+) -> Result<Vec<RecordBatch>> {
+ let mut batches = Vec::new();
+ for endpoint in info.endpoint {
+ let ticket = endpoint
+ .ticket
+ .ok_or_else(|| BallistaError::General("endpoint has no ticket".into()))?;
+ let mut collected: Vec<RecordBatch> =
+ client.do_get(ticket).await?.try_collect().await?;
+ batches.append(&mut collected);
+ }
+ Ok(batches)
+}
+
+/// Total rows across every batch a query returned.
+fn row_count(batches: &[RecordBatch]) -> usize {
+ batches.iter().map(|b| b.num_rows()).sum()
+}
+
+/// The headline path: authenticate, register a table with DDL, run a query
+/// that the cluster actually distributes, and read the results back through
+/// the scheduler.
+#[tokio::test]
+async fn flight_sql_runs_a_distributed_query() -> Result {
+ let url = start_cluster().await?;
+ let (mut client, _csv) = client_with_people(&url).await?;
+
+ let info = client
+ .execute(
+ "SELECT name FROM people WHERE id > 1 ORDER BY name".to_string(),
+ None,
+ )
+ .await?;
+
+ // Every endpoint must be redeemable on the connection we already have:
+ // no executor address is allowed to leak to the client.
+ assert!(
+ info.endpoint.iter().all(|e| e.location.is_empty()),
+ "endpoints must not advertise cluster-internal addresses: {:?}",
+ info.endpoint
+ );
+
+ let batches = collect(&mut client, info).await?;
+ assert_eq!(
+ row_count(&batches),
+ 2,
+ "expected bob and carol, got {batches:?}"
+ );
+
+ Ok(())
+}
+
+/// A table registered through DDL must still be there for the next request.
+/// `QueryBackend::session` builds a fresh `SessionContext` every call, so this
+/// only holds because the frontend caches it per session.
+#[tokio::test]
+async fn a_session_keeps_its_catalog_across_requests() -> Result {
Review Comment:
**Test.** This is already covered by `flight_sql_runs_a_distributed_query` and `flight_sql_supports_prepared_statements`, which both register the table in one request and query it in another. Each test boots its own cluster, so dropping it also saves CI time.
##########
ballista/flight-sql/tests/service.rs:
##########
@@ -0,0 +1,593 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Protocol-level tests for the Flight SQL frontend, driven through the
+//! `FlightSqlService` trait against a stub backend.
+//!
+//! These deliberately avoid a live cluster so they run anywhere; the
+//! scheduler's `tests/flight_sql.rs` covers the same surface against real
+//! executors. Everything here is asserted the way a client sees it — tickets
+//! are decoded from the wire bytes rather than through crate internals — so
+//! the tests fail if the wire format changes.
+
+use std::sync::Arc;
+
+use arrow_flight::decode::FlightRecordBatchStream;
+use arrow_flight::error::FlightError;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCreatePreparedStatementRequest, Any, CommandGetCatalogs, CommandGetTables,
+ CommandStatementQuery, TicketStatementQuery,
+};
+use arrow_flight::{Action, FlightDescriptor, FlightInfo, Ticket};
+use async_trait::async_trait;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::serde::decode_protobuf;
+use ballista_core::serde::protobuf::{ExecutorMetadata, PartitionId, PartitionLocation};
+use ballista_core::serde::scheduler::{
+ Action as BallistaAction, ShuffleFileKind, ShuffleLayout,
+};
+use ballista_flight_sql::backend::{QueryBackend, QueryResult};
+use ballista_flight_sql::{
+ ANONYMOUS_SESSION, AnonymousAuthenticator, Authenticator, BallistaFlightSqlService,
+ Identity,
+};
+use dashmap::DashMap;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::empty::EmptyTable;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+use futures::TryStreamExt;
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Status};
+
+/// A backend that never contacts a cluster: it hands out real
+/// `SessionContext`s so planning is genuine, and fabricates the partition
+/// locations a completed job would have produced.
+#[derive(Default)]
+struct StubBackend {
+ sessions: DashMap<String, Arc<SessionContext>>,
+ executed: DashMap<String, usize>,
+ cancelled: DashMap<String, usize>,
+}
+
+const JOB_ID: &str = "job-1";
+
+#[async_trait]
+impl QueryBackend for StubBackend {
+ async fn session(&self, session_id: &str) -> Result<Arc<SessionContext>> {
+ Ok(self
+ .sessions
+ .entry(session_id.to_string())
+ .or_insert_with(|| Arc::new(SessionContext::new()))
+ .clone())
+ }
+
+ async fn close_session(&self, session_id: &str) -> Result<()> {
+ self.sessions.remove(session_id);
+ Ok(())
+ }
+
+ async fn execute(
+ &self,
+ _job_name: &str,
+ _ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ ) -> Result<QueryResult> {
+ *self.executed.entry(JOB_ID.to_string()).or_insert(0) += 1;
+
+ Ok(QueryResult {
+ job_id: JOB_ID.to_string(),
+ schema: Arc::new(plan.schema().as_arrow().clone()),
+ partitions: vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: Some(PartitionId {
+ job_id: JOB_ID.to_string(),
+ stage_id: 3,
+ partition_id: 7,
+ }),
+ executor_meta: Some(ExecutorMetadata {
+ id: "executor-1".to_string(),
+ host: "executor-host".to_string(),
+ port: 50051,
+ grpc_port: 50052,
+ specification: None,
+ os_info: None,
+ }),
+ partition_stats: None,
+ file_id: None,
+ is_sort_shuffle: false,
+ }],
+ })
+ }
+
+ async fn cancel(&self, job_id: &str) -> Result<()> {
+ *self.cancelled.entry(job_id.to_string()).or_insert(0) += 1;
+ Ok(())
+ }
+}
+
+/// Rejects everything, to check that the frontend actually consults the
+/// authenticator and refuses tokenless requests when one is installed.
+struct DenyAll;
+
+#[async_trait]
+impl Authenticator for DenyAll {
+ async fn authenticate(
+ &self,
+ _headers: &MetadataMap,
+ ) -> std::result::Result<Identity, Status> {
+ Err(Status::unauthenticated("nope"))
+ }
+}
+
+fn make_service(backend: Arc<StubBackend>) -> BallistaFlightSqlService<StubBackend> {
+ BallistaFlightSqlService::new(
+ backend,
+ BallistaFlightProxyService::new(4_194_304, 4_194_304, false, None),
+ )
+}
+
+fn descriptor() -> Request<FlightDescriptor> {
+ Request::new(FlightDescriptor::new_cmd(vec![]))
+}
+
+/// `expect_err` needs `Debug` on the success type, and Flight's stream
+/// responses do not have it.
+fn expect_err<T>(result: std::result::Result<T, Status>, msg: &str) -> Status {
+ match result {
+ Ok(_) => panic!("{msg}"),
+ Err(status) => status,
+ }
+}
+
+/// Pulls the ticket out of an endpoint exactly as a Flight client would.
+fn statement_handle(info: &FlightInfo, index: usize) -> Vec<u8> {
+ let ticket = info.endpoint[index].ticket.as_ref().expect("ticket");
+ let any = Any::decode(&*ticket.ticket).expect("ticket is an Any");
+ let statement: TicketStatementQuery =
+ any.unpack().expect("unpackable").expect("statement ticket");
+ statement.statement_handle.to_vec()
+}
+
+fn ticket_for(info: &FlightInfo, index: usize) -> Request<Ticket> {
+ Request::new(info.endpoint[index].ticket.clone().expect("ticket"))
+}
+
+type DoGetStream = <BallistaFlightSqlService<StubBackend> as FlightService>::DoGetStream;
+
+async fn collect(response: tonic::Response<DoGetStream>) -> Vec<RecordBatch> {
+ let stream = response
+ .into_inner()
+ .map_err(|e| FlightError::Tonic(Box::new(e)));
+ FlightRecordBatchStream::new_from_flight_data(stream)
+ .try_collect()
+ .await
+ .expect("results decode")
+}
+
+async fn register_table(backend: &StubBackend, session_id: &str) {
+ let ctx = backend.session(session_id).await.unwrap();
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, true),
+ Field::new("name", DataType::Utf8, true),
+ ]));
+ ctx.register_table("people", Arc::new(EmptyTable::new(schema)))
+ .unwrap();
+}
+
+#[tokio::test]
+async fn select_produces_one_endpoint_per_partition_with_no_location() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query planned and submitted")
+ .into_inner();
+
+ assert_eq!(info.endpoint.len(), 1);
+ assert!(
+ info.endpoint[0].location.is_empty(),
+ "an endpoint with no location tells the client to reuse its existing \
+ connection; advertising the executor is what broke #1012"
+ );
+
+ // The ticket must carry the executor-facing fetch action, so the proxy can
+ // redeem it without any extra server-side lookup.
+ let handle = statement_handle(&info, 0);
+ let action = decode_protobuf(&handle[1..]).expect("handle wraps a Ballista action");
+ match action {
+ BallistaAction::FetchPartition {
+ job_id,
+ stage_id,
+ partition_id,
+ host,
+ port,
+ ..
+ } => {
+ assert_eq!(job_id.to_string(), JOB_ID);
+ assert_eq!(stage_id, 3);
+ assert_eq!(partition_id, 7);
+ assert_eq!(host, "executor-host");
+ assert_eq!(port, 50051);
+ }
+ }
+}
+
+#[tokio::test]
+async fn ddl_runs_on_the_scheduler_and_its_result_is_single_use() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE SCHEMA reporting".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("DDL executes")
+ .into_inner();
+
+ // DDL never reaches the cluster.
+ assert!(backend.executed.is_empty());
+
+ // The schema really was created in the session the client will query.
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ ctx.catalog("datafusion")
+ .unwrap()
+ .schema("reporting")
+ .is_some()
+ );
+
+ let handle = statement_handle(&info, 0);
+ assert_eq!(handle[0], 1, "DDL results use the local-result tag");
+
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.into(),
+ };
+ service
+ .do_get_statement(ticket.clone(), ticket_for(&info, 0))
+ .await
+ .expect("first fetch succeeds");
+
+ let err = expect_err(
+ service.do_get_statement(ticket, ticket_for(&info, 0)).await,
+ "a ticket is redeemable once",
+ );
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn writes_are_refused_rather_than_run_on_the_scheduler() {
Review Comment:
**Test.** This and `ctas_is_refused_rather_than_run_on_one_node` exercise the same path. One test that loops over `INSERT`, CTAS, `COPY` and `EXECUTE` (or an `rstest` case list) would cover more statements with less code. The `EXECUTE` case fails today (see the comment on `disposition`).
##########
ballista/scheduler/src/state/mod.rs:
##########
@@ -305,23 +340,35 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
.await
{
Ok(executor) => {
- if let Err(e) = state
+ match state
.task_manager
.launch_multi_task(&executor, tasks, &state.executor_manager)
.await
{
- let err_msg = format!("Failed to launch new task: {e}");
- error!("{}", err_msg.clone());
-
- // It's OK to remove executor aggressively,
- // since if the executor is in healthy state, it will be registered again.
- state
- .remove_executor(&executor_id, Some(err_msg), &sender)
- .await;
-
- false
- } else {
- true
+ Ok(unpreparable) => {
Review Comment:
**Major.** Failing the job is right, but the vcores bound to these tasks never come back. This branch reports success, so the spawned task returns `vec![]` and `unbind_tasks` never refunds what `bind_one` reserved. `abort_job` then sends `CancelTasks` for tasks the executor never received, so no status arrives to refund them either. Under `PushStaged`, each such job permanently lowers that executor's capacity, which is the wedge described in #2418. I found this by reading the code and haven't reproduced it.
Suggest returning the unpreparable tasks' `vcores_consumed` as `ExecutorSlot`s along with failing the job.
Two related points:
- The pull path already fails jobs for the same reason through `SchedulerServer::fail_job` (`scheduler_server/grpc.rs:174`). One helper shared by both paths would keep them consistent, including `queued_at`, which is looked up here and set to now in `fail_job`.
- This changes scheduling for every `PushStaged` user, not only Flight SQL, so it may be easier to review and backport as its own PR.
##########
ballista/scheduler/src/state/task_manager.rs:
##########
@@ -1456,4 +1471,42 @@ mod tests {
assert_eq!(manager.running_job_number(), 0);
Ok(())
}
+
+ /// A task whose definition cannot be prepared is never sent to an executor,
+ /// so no status will ever come back for it. `launch_multi_task` has to name
+ /// the job it dropped, or the job stays Running for the life of the
+ /// scheduler and everything awaiting its status blocks with no timeout.
+ #[tokio::test]
+ async fn tasks_that_cannot_be_prepared_are_reported_against_their_job() -> Result<()>
Review Comment:
**Test.** This checks what `launch_multi_task` returns, not the behavior the fix is about, so it still passes if the `fail_unlaunchable_job` call in `state/mod.rs` is reverted. A `SchedulerTest` case under `PushStaged` that asserts the job ends up `Failed`, and that the executor's vcores are back, would fail without the fix.
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
Review Comment:
**Minor.** `disposition(&plan)` runs twice here, and each call walks the plan through `scans_only_local_tables`. One `match disposition(&plan) { ... }` over the three variants, as `do_put_statement_update` already does, avoids the second pass.
##########
.github/workflows/rust.yml:
##########
@@ -99,6 +99,12 @@ jobs:
export ARROW_TEST_DATA=$(pwd)/testing/data
export PARQUET_TEST_DATA=$(pwd)/parquet-testing/data
cargo test --profile ci --features=testcontainers
+ - name: Run Arrow Flight SQL tests
+ # The frontend is a non-default feature, so the workspace run above
+ # compiles its integration test away.
+ run: |
+ export PATH=$PATH:$HOME/d/protoc/bin
+ cargo test --profile ci -p ballista-scheduler --features flight-sql --test flight_sql
Review Comment:
**Minor (CI).** `ci/scripts/rust_clippy.sh` lints an explicit list of packages and doesn't include `ballista-flight-sql`. `--package ballista-scheduler --all-features` only lints the scheduler's own code, so the new crate is never linted. Suggest adding it to that script.
Also, this step compiles the scheduler a second time with a different feature set. Adding `ballista-scheduler/flight-sql` to `--features` in the step above should cover it in one build. I haven't measured the difference.
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
+ return Err(Status::unimplemented(reason));
+ }
+
+ if disposition(&plan) == Disposition::RunOnScheduler {
+ let (schema, batches) = execute_locally(&ctx, plan).await?;
+ let handle = Uuid::new_v4().to_string();
+ self.store.insert_result(
+ handle.clone(),
+ LocalResult {
+ schema: schema.clone(),
+ batches,
+ },
+ );
+
+ let endpoint = Self::endpoint(StatementHandle::Local(handle));
+ return build_flight_info(&schema, vec![endpoint], descriptor);
+ }
+
+ let result = self
+ .backend
+ .execute(job_name, ctx, plan)
+ .await
+ .map_err(query_failed)?;
+
+ let endpoints = result
+ .partitions
+ .into_iter()
+ .map(|location| partition_handle(location).map(Self::endpoint))
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| Status::internal(format!("invalid partition location: {e}")))?;
+
+ log::debug!(
+ "flight-sql: job {} produced {} endpoint(s)",
+ result.job_id,
+ endpoints.len()
+ );
+
+ build_flight_info(&result.schema, endpoints, descriptor)
+ }
+
+ /// Builds the endpoint a client redeems for one slice of the result.
+ ///
+ /// Endpoints carry no location, which Flight defines as "fetch from the
+ /// server that gave you this FlightInfo". That keeps every cluster-internal
+ /// address off the wire and lets the frontend work unchanged behind NAT, a
+ /// load balancer, or an ingress.
+ fn endpoint(handle: StatementHandle) -> FlightEndpoint {
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.encode().into(),
+ };
+ FlightEndpoint::new().with_ticket(Ticket {
+ ticket: ticket.as_any().encode_to_vec().into(),
+ })
+ }
+
+ /// Serves a metadata command by round-tripping the command itself as the
+ /// ticket, so `DoGet` lands back on the matching handler.
+ fn metadata_info<C: ProstMessageExt>(
+ command: C,
+ schema: SchemaRef,
+ descriptor: FlightDescriptor,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let endpoint = FlightEndpoint::new().with_ticket(Ticket {
+ ticket: command.as_any().encode_to_vec().into(),
+ });
+ build_flight_info(&schema, vec![endpoint], descriptor).map(Response::new)
+ }
+}
+
+/// What the frontend should do with a planned statement.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Disposition {
+ /// Execute on the scheduler: it mutates session state and produces no data
+ /// worth distributing.
+ RunOnScheduler,
+ /// Submit to the cluster.
+ Distribute,
+ /// Refuse, with an explanation for the client.
+ Unsupported(&'static str),
+}
+
+/// Classifies a plan.
+///
+/// One function rather than two predicates, because the interesting cases are
+/// the ones where the answers overlap: `CREATE TABLE AS SELECT` is DDL, and
+/// DDL runs on the scheduler, so a caller that asked "is this DDL?" before
+/// asking "is this supported?" would silently execute its query on one node.
+/// Returning a single verdict makes that ordering impossible to get wrong.
+fn disposition(plan: &LogicalPlan) -> Disposition {
+ match plan {
+ LogicalPlan::Dml(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support INSERT/UPDATE/DELETE; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Copy(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support COPY; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support CREATE TABLE AS SELECT, \
+ because it would execute on the scheduler rather than the cluster; \
+ use CREATE EXTERNAL TABLE over data the executors can read",
+ ),
+ // Other DDL only edits the session catalog, and `SET`-style statements
+ // only edit session config; neither has anything to distribute.
+ LogicalPlan::Ddl(_) | LogicalPlan::Statement(_) => Disposition::RunOnScheduler,
+ // `SHOW ...` and anything else reading `information_schema` describes
+ // the catalog held here. There is nothing to distribute, and the
+ // physical form of those scans cannot be serialized for an executor,
+ // so distributing one would hand the client a job that never runs.
+ _ if scans_only_local_tables(plan) => Disposition::RunOnScheduler,
+ _ => Disposition::Distribute,
+ }
+}
+
+async fn execute_locally(
+ ctx: &SessionContext,
+ plan: LogicalPlan,
+) -> Result<(SchemaRef, Vec<RecordBatch>), Status> {
+ let df = ctx
+ .execute_logical_plan(plan)
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+ let planned_schema: SchemaRef = Arc::new(df.schema().as_arrow().clone());
+ let batches = df
+ .collect()
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+
+ // Prefer the schema the data actually carries; DataFusion's DDL results
+ // are empty and their frame schema is not always the same object.
+ let schema = batches
+ .first()
+ .map(|batch| batch.schema())
+ .unwrap_or(planned_schema);
+
+ Ok((schema, batches))
+}
+
+/// Turns a shuffle partition into the ticket payload the Flight proxy already
+/// knows how to redeem.
+fn partition_handle(
+ location: PartitionLocation,
+) -> Result<StatementHandle, BallistaError> {
+ let layout = location.layout();
+ let partition_id = location.partition_id.ok_or_else(|| {
+ BallistaError::Internal("partition location has no partition id".to_string())
+ })?;
+ let executor = location.executor_meta.ok_or_else(|| {
+ BallistaError::Internal("partition location has no executor metadata".to_string())
+ })?;
+
+ let action = BallistaAction::FetchPartition {
+ job_id: partition_id.job_id.into(),
+ stage_id: partition_id.stage_id as usize,
+ partition_id: partition_id.partition_id as usize,
+ host: executor.host,
+ port: executor.port as u16,
+ file_id: location.file_id,
+ layout,
+ // A Flight client wants the whole partition: the data file, in full.
+ file_kind: ShuffleFileKind::Data,
+ byte_ranges: vec![],
+ };
+
+ let encoded: protobuf::Action = action.try_into()?;
+ Ok(StatementHandle::Partition(encoded.encode_to_vec()))
+}
+
+fn build_flight_info(
+ schema: &Schema,
+ endpoints: Vec<FlightEndpoint>,
+ descriptor: FlightDescriptor,
+) -> Result<FlightInfo, Status> {
+ FlightInfo::new()
+ .try_with_schema(schema)
+ .map_err(|e| Status::internal(format!("failed to encode result schema: {e}")))
+ .map(|info| info.with_descriptor(descriptor).with_endpoints(endpoints))
+}
+
+/// Streams a single metadata batch back to the client.
+fn one_batch_response(batch: RecordBatch) -> Response<DoGetStream> {
+ batch_response(batch.schema(), vec![batch])
+}
+
+fn batch_response(schema: SchemaRef, batches: Vec<RecordBatch>) -> Response<DoGetStream> {
+ let stream = FlightDataEncoderBuilder::new()
+ .with_schema(schema)
+ .build(futures::stream::iter(batches.into_iter().map(Ok)))
+ .map_err(|e| Status::internal(format!("failed to encode results: {e}")));
+
+ Response::new(Box::pin(stream) as DoGetStream)
+}
+
+fn bearer_token(metadata: &MetadataMap) -> Option<String> {
+ let value = metadata.get("authorization")?.to_str().ok()?;
+ value
+ .strip_prefix("Bearer ")
+ .or_else(|| value.strip_prefix("bearer "))
+ .map(str::to_string)
+}
+
+fn invalid_ticket(e: impl std::fmt::Display) -> Status {
+ Status::invalid_argument(format!("invalid ticket: {e}"))
+}
+
+fn query_failed(e: BallistaError) -> Status {
+ // A failed job is the client's problem to see, not an opaque 500.
+ Status::internal(format!("query execution failed: {e}"))
+}
+
+fn encode_schema(schema: &Schema) -> Result<Vec<u8>, Status> {
+ let message: IpcMessage = SchemaAsIpc::new(schema, &IpcWriteOptions::default())
+ .try_into()
+ .map_err(|e| Status::internal(format!("failed to encode schema: {e}")))?;
+ Ok(message.0.to_vec())
+}
+
+#[tonic::async_trait]
+impl<B: QueryBackend> FlightSqlService for BallistaFlightSqlService<B> {
+ type FlightService = BallistaFlightSqlService<B>;
+
+ async fn do_handshake(
+ &self,
+ request: Request<Streaming<HandshakeRequest>>,
+ ) -> Result<
+ Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
+ Status,
+ > {
+ let identity = self.auth.authenticate(request.metadata()).await?;
+
+ let token = Uuid::new_v4().to_string();
+ let session_id = format!("flight-sql-{}", Uuid::new_v4());
+
+ // Build the session eagerly so a failure surfaces at handshake time
+ // rather than on the client's first query, and so the first query does
+ // not pay for it.
+ self.open_session(&session_id).await?;
+ self.store.insert_session(token.clone(), session_id.clone());
+
+ log::debug!(
+ "flight-sql: handshake for {:?} bound to session {session_id}",
+ identity.user
+ );
+
+ let result = HandshakeResponse {
+ protocol_version: 0,
+ payload: token.clone().into(),
+ };
+ let stream = futures::stream::once(async move { Ok(result) });
+
+ let mut response = Response::new(Box::pin(stream) as _);
+ response.metadata_mut().insert(
+ "authorization",
+ format!("Bearer {token}")
+ .parse()
+ .map_err(|_| Status::internal("failed to encode session token"))?,
+ );
+ Ok(response)
+ }
+
+ /// Redeems tickets that are not Flight SQL commands.
+ ///
+ /// Ballista's own Rust client fetches shuffle output with a
+ /// `ballista.protobuf.Action` ticket. When the Flight SQL frontend is
+ /// mounted it replaces the standalone proxy on the scheduler's port, so it
+ /// has to keep serving those tickets.
+ async fn do_get_fallback(
+ &self,
+ request: Request<Ticket>,
+ _message: Any,
+ ) -> Result<Response<DoGetStream>, Status> {
+ decode_protobuf(&request.get_ref().ticket).map_err(invalid_ticket)?;
+ self.proxy.do_get(request).await
+ }
+
+ async fn get_flight_info_statement(
+ &self,
+ query: CommandStatementQuery,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let ctx = self.context(request.metadata()).await?;
+ let plan = Self::plan(&ctx, &query.query).await?;
+ let descriptor = request.into_inner();
+
+ self.flight_info_for(ctx, plan, descriptor, &job_name(&query.query))
+ .await
+ .map(Response::new)
+ }
+
+ async fn get_flight_info_prepared_statement(
+ &self,
+ query: CommandPreparedStatementQuery,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let handle = prepared_handle(&query.prepared_statement_handle)?;
+ let prepared = self.store.prepared(&handle).ok_or_else(|| {
+ Status::not_found("unknown or expired prepared statement handle")
+ })?;
+
+ let ctx = self.open_session(&prepared.session_id).await?;
+ let descriptor = request.into_inner();
+
+ self.flight_info_for(ctx, prepared.plan, descriptor, "flight-sql prepared")
+ .await
+ .map(Response::new)
+ }
+
+ async fn get_flight_info_catalogs(
+ &self,
+ query: CommandGetCatalogs,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_schemas(
+ &self,
+ query: CommandGetDbSchemas,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_tables(
+ &self,
+ query: CommandGetTables,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_table_types(
+ &self,
+ query: CommandGetTableTypes,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder().schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_sql_info(
+ &self,
+ query: CommandGetSqlInfo,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.clone().into_builder(&self.sql_info).schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn get_flight_info_xdbc_type_info(
+ &self,
+ query: CommandGetXdbcTypeInfo,
+ request: Request<FlightDescriptor>,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let schema = query.into_builder(&self.xdbc_info).schema();
+ Self::metadata_info(query, schema, request.into_inner())
+ }
+
+ async fn do_get_statement(
Review Comment:
**Major (security).** With an authenticator installed, the token is only checked where a session context is needed. `do_get_statement`, `get_flight_info_prepared_statement`, `do_action_close_prepared_statement`, `do_action_cancel_query` and `do_get_fallback` never call `session_id`.
Partition tickets describe themselves fully and aren't signed, and job IDs are listed at `GET /api/jobs` on the same port (with the default `rest-api` feature). So a client with no token can forge a ticket and read any job's result partitions.
Suggestions:
- Check the token in every handler that reads data or changes state.
- Bind prepared handles and local results to the session that created them.
- Sign partition tickets, or swap them for a server-side handle.
`do_get_fallback` has to stay open for native Ballista clients, which is one more reason the docs shouldn't present the authenticator as a security boundary (see the comment in `config.rs`). `Identity` is currently only logged. It's the natural thing to bind handles to, otherwise it can go.
##########
ballista/flight-sql/src/session.rs:
##########
@@ -0,0 +1,277 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Server-side state the Flight SQL frontend keeps between requests.
+//!
+//! Everything in here is keyed by an opaque handle the client is given and
+//! hands back, and everything expires. A client that disconnects without
+//! closing its prepared statements (or that never redeems a ticket) must not
+//! pin memory forever, which is what the previous Flight SQL implementation
+//! got wrong: its plan cache only shed entries on an explicit close.
+
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::SchemaRef;
+use dashmap::DashMap;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+
+/// A prepared statement's server-side state.
+#[derive(Clone)]
+pub(crate) struct Prepared {
+ /// Session the statement was prepared in; it is planned and executed
+ /// against that session's catalog.
+ pub session_id: String,
+ /// The plan as prepared; execution never re-plans it.
+ pub plan: LogicalPlan,
+}
+
+/// Result of a statement the frontend ran locally rather than distributing
+/// (DDL, and DML whose only output is an affected-row count).
+#[derive(Clone)]
+pub(crate) struct LocalResult {
+ pub schema: SchemaRef,
+ pub batches: Vec<RecordBatch>,
+}
+
+struct Tracked<T> {
+ value: T,
+ last_used: Instant,
+}
+
+impl<T> Tracked<T> {
+ fn new(value: T) -> Self {
+ Self {
+ value,
+ last_used: Instant::now(),
+ }
+ }
+}
+
+/// Handle stores for sessions, prepared statements, and locally-computed
+/// results, all with a shared idle TTL.
+pub(crate) struct SessionStore {
+ ttl: Duration,
+ /// Bearer token -> Ballista session id.
+ sessions: DashMap<String, Tracked<String>>,
+ /// Ballista session id -> its context.
+ contexts: DashMap<String, Tracked<Arc<SessionContext>>>,
+ /// Prepared statement handle -> prepared plan.
+ prepared: DashMap<String, Tracked<Prepared>>,
+ /// Local result handle -> materialized batches.
+ results: DashMap<String, Tracked<LocalResult>>,
+}
+
+impl SessionStore {
+ pub(crate) fn new(ttl: Duration) -> Self {
+ Self {
+ ttl,
+ sessions: DashMap::new(),
+ contexts: DashMap::new(),
+ prepared: DashMap::new(),
+ results: DashMap::new(),
+ }
+ }
+
+ pub(crate) fn insert_session(&self, token: String, session_id: String) {
+ self.sessions.insert(token, Tracked::new(session_id));
+ }
+
+ /// Resolves a bearer token to its session id, refreshing its idle timer.
+ pub(crate) fn session(&self, token: &str) -> Option<String> {
+ self.sessions.get_mut(token).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ /// Drops a token and the session it named, returning that session id.
+ ///
+ /// Each token gets its own session, so removing the token always releases
+ /// the session with it.
+ pub(crate) fn remove_session(&self, token: &str) -> Option<String> {
+ let (_, entry) = self.sessions.remove(token)?;
+ self.contexts.remove(&entry.value);
+ Some(entry.value)
+ }
+
+ /// Returns the cached context for a session, refreshing its idle timer.
+ pub(crate) fn context(&self, session_id: &str) -> Option<Arc<SessionContext>> {
+ self.contexts.get_mut(session_id).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ /// Caches a context, or returns the one another caller cached first.
+ ///
+ /// Two requests for a cold session can both build one; letting the first
+ /// insertion win keeps them looking at the same catalog.
+ pub(crate) fn insert_context(
+ &self,
+ session_id: String,
+ ctx: Arc<SessionContext>,
+ ) -> Arc<SessionContext> {
+ self.contexts
+ .entry(session_id)
+ .or_insert_with(|| Tracked::new(ctx))
+ .value
+ .clone()
+ }
+
+ pub(crate) fn insert_prepared(&self, handle: String, prepared: Prepared) {
+ self.prepared.insert(handle, Tracked::new(prepared));
+ }
+
+ pub(crate) fn prepared(&self, handle: &str) -> Option<Prepared> {
+ self.prepared.get_mut(handle).map(|mut entry| {
+ entry.last_used = Instant::now();
+ entry.value.clone()
+ })
+ }
+
+ pub(crate) fn remove_prepared(&self, handle: &str) {
+ self.prepared.remove(handle);
+ }
+
+ pub(crate) fn insert_result(&self, handle: String, result: LocalResult) {
+ self.results.insert(handle, Tracked::new(result));
+ }
+
+ /// Takes a local result. Results are single-use: a ticket is redeemed once.
+ pub(crate) fn take_result(&self, handle: &str) -> Option<LocalResult> {
+ self.results.remove(handle).map(|(_, entry)| entry.value)
+ }
+
+ /// Evicts everything idle for longer than the TTL.
+ ///
+ /// Returns the session ids that no longer have any live token, so the
+ /// caller can release them in the backend.
+ pub(crate) fn sweep(&self) -> Vec<String> {
+ let ttl = self.ttl;
+ let expired: Vec<String> = self
+ .sessions
+ .iter()
+ .filter(|entry| entry.value().last_used.elapsed() > ttl)
+ .map(|entry| entry.key().clone())
+ .collect();
+
+ let mut released = Vec::new();
+ for token in expired {
+ if let Some(session_id) = self.remove_session(&token) {
+ released.push(session_id);
+ }
+ }
+
+ self.contexts
+ .retain(|_, entry| entry.last_used.elapsed() <= ttl);
+ self.prepared
+ .retain(|_, entry| entry.last_used.elapsed() <= ttl);
+ self.results
+ .retain(|_, entry| entry.last_used.elapsed() <= ttl);
+
+ released
+ }
+
+ /// Starts a background task that sweeps at `interval`, closing sessions it
+ /// evicts. The task ends when the last reference to the store is dropped.
+ pub(crate) fn spawn_reaper<F, Fut>(self: &Arc<Self>, interval: Duration, close: F)
+ where
+ F: Fn(String) -> Fut + Send + Sync + 'static,
+ Fut: std::future::Future<Output = ()> + Send,
+ {
+ let store = Arc::downgrade(self);
+ tokio::spawn(async move {
+ let mut ticker = tokio::time::interval(interval);
+ // The first tick completes immediately; skip it so we do not sweep
+ // a store that was created a moment ago.
+ ticker.tick().await;
+ loop {
+ ticker.tick().await;
+ let Some(store) = store.upgrade() else {
+ return;
+ };
+ for session_id in store.sweep() {
+ log::debug!("flight-sql: expiring idle session {session_id}");
+ close(session_id).await;
+ }
+ }
+ });
+ }
+}
+
+#[cfg(test)]
+mod test {
+ use super::*;
+
+ #[test]
+ fn session_survives_while_touched_and_expires_when_idle() {
+ let store = SessionStore::new(Duration::from_millis(50));
+ store.insert_session("token".to_string(), "session".to_string());
+
+ assert_eq!(store.session("token"), Some("session".to_string()));
+ assert!(store.sweep().is_empty());
+
+ std::thread::sleep(Duration::from_millis(60));
Review Comment:
**Minor (flaky test).** The first `sweep()` has to run within 50 ms of `insert_session`, which a loaded CI runner can miss. The idle-eviction tests in `executor/src/client_pool.rs` backdate the timestamp (`Instant::now() - age`) instead of sleeping. Setting `last_used` in the past here would make this deterministic.
##########
ballista/scheduler/src/flight_sql.rs:
##########
@@ -0,0 +1,111 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Binds the [`ballista_flight_sql`] frontend to this scheduler.
+//!
+//! The frontend does not know about `SchedulerServer`; this is the only place
+//! the two meet, which is what keeps the Flight SQL code from re-acquiring the
+//! coupling that made the pre-46.0.0 implementation unmaintainable.
+
+use std::sync::Arc;
+
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::serde::protobuf::{JobStatus, SuccessfulJob, job_status};
+use ballista_flight_sql::backend::{QueryBackend, QueryResult};
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+use datafusion_proto::logical_plan::AsLogicalPlan;
+use datafusion_proto::physical_plan::AsExecutionPlan;
+
+use crate::scheduler_server::SchedulerServer;
+
+/// Buffer for the job status stream. Statuses are consumed as fast as they
+/// arrive, so this only has to absorb bursts.
+const STATUS_BUFFER: usize = 16;
+
+#[async_trait::async_trait]
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> QueryBackend
+ for SchedulerServer<T, U>
+{
+ async fn session(&self, session_id: &str) -> Result<Arc<SessionContext>> {
+ let config = self.state.session_manager.produce_config();
+ self.state
+ .session_manager
+ .create_or_update_session(session_id, &config)
+ .await
+ }
+
+ async fn close_session(&self, session_id: &str) -> Result<()> {
+ self.state.session_manager.remove_session(session_id).await
+ }
+
+ async fn execute(
+ &self,
+ job_name: &str,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ ) -> Result<QueryResult> {
+ // The schema handed to the client has to be the one the shuffle files
+ // actually carry, which is the physical plan's, not the logical plan's:
+ // the two disagree about nullability often enough to matter (`version()`
+ // is nullable logically and not-null in the data; a list literal goes the
+ // other way), and a Flight SQL client that is promised one schema and
+ // handed another -- ADBC among them -- rejects the result outright.
+ // Planning twice costs a few milliseconds against a job that is about to
+ // run on the cluster.
+ let schema = ctx.state().create_physical_plan(&plan).await?.schema();
Review Comment:
**Question.** This builds the physical plan once only to read its schema, and the job is planned again on submit. For listing tables on object storage, physical planning lists files (and may collect statistics), so that work happens twice per query. Could the schema come from the job's final stage plan once it succeeds, since that is the plan that wrote the files?
##########
ballista/scheduler/tests/flight_sql.rs:
##########
@@ -0,0 +1,339 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! End-to-end coverage for the Arrow Flight SQL frontend against a real
+//! (in-process) cluster: scheduler plus one executor, driven entirely through
+//! the Flight SQL protocol.
+//!
+//! The pre-46.0.0 implementation shipped with no tests at all, and the bugs
+//! that got it removed (#1012, #941, #839, #756) were exactly the ones an
+//! end-to-end test catches: endpoints pointing somewhere the client cannot
+//! reach, and results that fail to decode.
+
+#![cfg(feature = "flight-sql")]
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow_flight::sql::client::FlightSqlServiceClient;
+use arrow_flight::sql::{CommandGetTables, SqlInfo};
+use ballista_core::config::TaskSchedulingPolicy;
+use ballista_core::error::BallistaError;
+use ballista_core::serde::BallistaCodec;
+use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
+use ballista_core::utils::{
+ GrpcClientConfig, create_grpc_client_connection, default_config_producer,
+ default_session_builder,
+};
+use ballista_scheduler::cluster::BallistaCluster;
+use ballista_scheduler::config::SchedulerConfig;
+use ballista_scheduler::metrics::default_metrics_collector;
+use ballista_scheduler::scheduler_process::start_grpc_service_with_listener;
+use ballista_scheduler::scheduler_server::SchedulerServer;
+use datafusion::arrow::array::RecordBatch;
+use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+use futures::TryStreamExt;
+use tonic::transport::Channel;
+
+/// Test errors are only ever printed, and the client, transport, and Ballista
+/// error types all implement `Error` — so one boxed type spares every call site
+/// a `map_err`.
+type Result<T = ()> = std::result::Result<T, Box<dyn std::error::Error>>;
+
+/// Boots a scheduler serving Flight SQL plus one executor, and returns the
+/// scheduler's URL.
+async fn start_cluster() -> Result<String> {
Review Comment:
**Test.** Every end-to-end test runs under `PullStaged`, while the hang fix in `state/mod.rs` is on the `PushStaged` path. Covering push would need a push-mode executor, since `new_standalone_executor` only runs `poll_loop`. If that's too much for this PR, the `SchedulerTest` case suggested on `task_manager.rs` covers the fix itself.
Smaller points: this repeats most of `standalone::new_standalone_scheduler_with_builder`, and it sets `config.flight_sql = true` directly instead of using `with_flight_sql(true)`.
##########
ballista/flight-sql/src/service.rs:
##########
@@ -0,0 +1,852 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! The Arrow Flight SQL frontend itself.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
+use arrow_flight::encode::FlightDataEncoderBuilder;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCancelQueryRequest, ActionCancelQueryResult,
+ ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
+ ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetDbSchemas,
+ CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo,
+ CommandPreparedStatementQuery, CommandStatementQuery, CommandStatementUpdate,
+ ProstMessageExt, SqlInfo, TicketStatementQuery,
+ metadata::{SqlInfoData, XdbcTypeInfoData},
+ server::PeekableFlightDataStream,
+};
+use arrow_flight::{
+ Action, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest,
+ HandshakeResponse, IpcMessage, SchemaAsIpc, Ticket,
+};
+use ballista_core::error::BallistaError;
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::planner::scans_only_local_tables;
+use ballista_core::serde::protobuf::PartitionLocation;
+use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind};
+use ballista_core::serde::{decode_protobuf, protobuf};
+use datafusion::logical_expr::{DdlStatement, LogicalPlan};
+use datafusion::prelude::SessionContext;
+use futures::{Stream, TryStreamExt};
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Response, Status, Streaming};
+use uuid::Uuid;
+
+use crate::auth::{AnonymousAuthenticator, Authenticator};
+use crate::backend::QueryBackend;
+use crate::metadata;
+use crate::session::{LocalResult, Prepared, SessionStore};
+use crate::ticket::StatementHandle;
+
+/// Session shared by every client that connects without authenticating.
+///
+/// Anonymous clients cannot be told apart, so they necessarily share catalog
+/// state. Configure an [`Authenticator`] to get a session per connection.
+pub const ANONYMOUS_SESSION: &str = "flight-sql-anonymous";
+
+/// How long a session, prepared statement, or unredeemed local result may sit
+/// idle before it is discarded.
+const DEFAULT_TTL: Duration = Duration::from_secs(30 * 60);
+
+/// How often expired handles are swept.
+const REAP_INTERVAL: Duration = Duration::from_secs(60);
+
+type DoGetStream =
+ Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
+
+/// Serves Arrow Flight SQL on behalf of a Ballista cluster.
+///
+/// Clients send SQL text; the frontend plans it against the session's catalog,
+/// submits the plan through a [`QueryBackend`], and hands back one
+/// `FlightEndpoint` per output partition. `DoGet` on those tickets is proxied
+/// to the executor holding the partition, so clients never need to reach
+/// executors themselves — the failure mode that made the pre-46.0.0
+/// implementation unusable behind NAT, Docker, and Kubernetes.
+pub struct BallistaFlightSqlService<B: QueryBackend> {
+ backend: Arc<B>,
+ proxy: BallistaFlightProxyService,
+ auth: Arc<dyn Authenticator>,
+ store: Arc<SessionStore>,
+ sql_info: SqlInfoData,
+ xdbc_info: XdbcTypeInfoData,
+}
+
+impl<B: QueryBackend> BallistaFlightSqlService<B> {
+ /// Builds a frontend over `backend`, using `proxy` to stream partition
+ /// data back from executors.
+ ///
+ /// The service authenticates nobody until an [`Authenticator`] is supplied
+ /// via [`with_authenticator`](Self::with_authenticator).
+ pub fn new(backend: Arc<B>, proxy: BallistaFlightProxyService) -> Self {
+ let store = Arc::new(SessionStore::new(DEFAULT_TTL));
+
+ let backend_for_reaper = backend.clone();
+ store.spawn_reaper(REAP_INTERVAL, move |session_id| {
+ let backend = backend_for_reaper.clone();
+ async move {
+ if let Err(e) = backend.close_session(&session_id).await {
+ log::warn!("flight-sql: failed to close session {session_id}: {e}");
+ }
+ }
+ });
+
+ Self {
+ backend,
+ proxy,
+ auth: Arc::new(AnonymousAuthenticator),
+ store,
+ sql_info: metadata::sql_info(),
+ xdbc_info: metadata::xdbc_type_info(),
+ }
+ }
+
+ /// Installs an authenticator. Without one, every handshake is accepted and
+ /// unauthenticated clients share a single session.
+ pub fn with_authenticator(mut self, auth: Arc<dyn Authenticator>) -> Self {
+ self.auth = auth;
+ self
+ }
+
+ /// True when the service will accept unauthenticated clients, which the
+ /// scheduler logs at startup.
+ pub fn allows_anonymous(&self) -> bool {
+ self.auth.allows_anonymous()
+ }
+
+ /// Resolves the Ballista session for a request from its bearer token.
+ fn session_id(&self, metadata: &MetadataMap) -> Result<String, Status> {
+ match bearer_token(metadata) {
+ Some(token) => self.store.session(&token).ok_or_else(|| {
+ Status::unauthenticated(
+ "unknown or expired session token; re-run the Flight handshake",
+ )
+ }),
+ None if self.auth.allows_anonymous() => Ok(ANONYMOUS_SESSION.to_string()),
+ None => Err(Status::unauthenticated(
+ "missing bearer token; authenticate with the Flight handshake first",
+ )),
+ }
+ }
+
+ /// Returns the context for `session_id`, building it on first use.
+ ///
+ /// The cache is what makes a session a session: [`QueryBackend::session`]
+ /// builds a fresh `SessionContext` every call, so without it a table
+ /// created by one request would be invisible to the next — and every
+ /// request would pay for a full DataFusion session to be constructed.
+ async fn open_session(
+ &self,
+ session_id: &str,
+ ) -> Result<Arc<SessionContext>, Status> {
+ if let Some(ctx) = self.store.context(session_id) {
+ return Ok(ctx);
+ }
+
+ let ctx = self
+ .backend
+ .session(session_id)
+ .await
+ .map_err(|e| Status::internal(format!("failed to open session: {e}")))?;
+
+ Ok(self.store.insert_context(session_id.to_string(), ctx))
+ }
+
+ /// Resolves a request to the context it should be planned against.
+ async fn context(
+ &self,
+ metadata: &MetadataMap,
+ ) -> Result<Arc<SessionContext>, Status> {
+ self.open_session(&self.session_id(metadata)?).await
+ }
+
+ /// Plans `sql` against the session's catalog.
+ async fn plan(ctx: &SessionContext, sql: &str) -> Result<LogicalPlan, Status> {
+ ctx.state()
+ .create_logical_plan(sql)
+ .await
+ .map_err(|e| Status::invalid_argument(format!("failed to plan query: {e}")))
+ }
+
+ /// Runs a planned statement and describes where to collect its results.
+ async fn flight_info_for(
+ &self,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ descriptor: FlightDescriptor,
+ job_name: &str,
+ ) -> Result<FlightInfo, Status> {
+ if let Disposition::Unsupported(reason) = disposition(&plan) {
+ return Err(Status::unimplemented(reason));
+ }
+
+ if disposition(&plan) == Disposition::RunOnScheduler {
+ let (schema, batches) = execute_locally(&ctx, plan).await?;
+ let handle = Uuid::new_v4().to_string();
+ self.store.insert_result(
+ handle.clone(),
+ LocalResult {
+ schema: schema.clone(),
+ batches,
+ },
+ );
+
+ let endpoint = Self::endpoint(StatementHandle::Local(handle));
+ return build_flight_info(&schema, vec![endpoint], descriptor);
+ }
+
+ let result = self
+ .backend
+ .execute(job_name, ctx, plan)
+ .await
+ .map_err(query_failed)?;
+
+ let endpoints = result
+ .partitions
+ .into_iter()
+ .map(|location| partition_handle(location).map(Self::endpoint))
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| Status::internal(format!("invalid partition location: {e}")))?;
+
+ log::debug!(
+ "flight-sql: job {} produced {} endpoint(s)",
+ result.job_id,
+ endpoints.len()
+ );
+
+ build_flight_info(&result.schema, endpoints, descriptor)
+ }
+
+ /// Builds the endpoint a client redeems for one slice of the result.
+ ///
+ /// Endpoints carry no location, which Flight defines as "fetch from the
+ /// server that gave you this FlightInfo". That keeps every cluster-internal
+ /// address off the wire and lets the frontend work unchanged behind NAT, a
+ /// load balancer, or an ingress.
+ fn endpoint(handle: StatementHandle) -> FlightEndpoint {
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.encode().into(),
+ };
+ FlightEndpoint::new().with_ticket(Ticket {
+ ticket: ticket.as_any().encode_to_vec().into(),
+ })
+ }
+
+ /// Serves a metadata command by round-tripping the command itself as the
+ /// ticket, so `DoGet` lands back on the matching handler.
+ fn metadata_info<C: ProstMessageExt>(
+ command: C,
+ schema: SchemaRef,
+ descriptor: FlightDescriptor,
+ ) -> Result<Response<FlightInfo>, Status> {
+ let endpoint = FlightEndpoint::new().with_ticket(Ticket {
+ ticket: command.as_any().encode_to_vec().into(),
+ });
+ build_flight_info(&schema, vec![endpoint], descriptor).map(Response::new)
+ }
+}
+
+/// What the frontend should do with a planned statement.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Disposition {
+ /// Execute on the scheduler: it mutates session state and produces no data
+ /// worth distributing.
+ RunOnScheduler,
+ /// Submit to the cluster.
+ Distribute,
+ /// Refuse, with an explanation for the client.
+ Unsupported(&'static str),
+}
+
+/// Classifies a plan.
+///
+/// One function rather than two predicates, because the interesting cases are
+/// the ones where the answers overlap: `CREATE TABLE AS SELECT` is DDL, and
+/// DDL runs on the scheduler, so a caller that asked "is this DDL?" before
+/// asking "is this supported?" would silently execute its query on one node.
+/// Returning a single verdict makes that ordering impossible to get wrong.
+fn disposition(plan: &LogicalPlan) -> Disposition {
+ match plan {
+ LogicalPlan::Dml(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support INSERT/UPDATE/DELETE; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Copy(_) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support COPY; \
+ the distributed write path is not implemented",
+ ),
+ LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Disposition::Unsupported(
+ "Ballista Flight SQL does not support CREATE TABLE AS SELECT, \
+ because it would execute on the scheduler rather than the cluster; \
+ use CREATE EXTERNAL TABLE over data the executors can read",
+ ),
+ // Other DDL only edits the session catalog, and `SET`-style statements
+ // only edit session config; neither has anything to distribute.
+ LogicalPlan::Ddl(_) | LogicalPlan::Statement(_) => Disposition::RunOnScheduler,
+ // `SHOW ...` and anything else reading `information_schema` describes
+ // the catalog held here. There is nothing to distribute, and the
+ // physical form of those scans cannot be serialized for an executor,
+ // so distributing one would hand the client a job that never runs.
+ _ if scans_only_local_tables(plan) => Disposition::RunOnScheduler,
+ _ => Disposition::Distribute,
+ }
+}
+
+async fn execute_locally(
+ ctx: &SessionContext,
+ plan: LogicalPlan,
+) -> Result<(SchemaRef, Vec<RecordBatch>), Status> {
+ let df = ctx
+ .execute_logical_plan(plan)
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+ let planned_schema: SchemaRef = Arc::new(df.schema().as_arrow().clone());
+ let batches = df
+ .collect()
+ .await
+ .map_err(|e| Status::internal(format!("failed to execute statement: {e}")))?;
+
+ // Prefer the schema the data actually carries; DataFusion's DDL results
+ // are empty and their frame schema is not always the same object.
+ let schema = batches
+ .first()
+ .map(|batch| batch.schema())
+ .unwrap_or(planned_schema);
+
+ Ok((schema, batches))
+}
+
+/// Turns a shuffle partition into the ticket payload the Flight proxy already
+/// knows how to redeem.
+fn partition_handle(
+ location: PartitionLocation,
+) -> Result<StatementHandle, BallistaError> {
+ let layout = location.layout();
+ let partition_id = location.partition_id.ok_or_else(|| {
+ BallistaError::Internal("partition location has no partition id".to_string())
+ })?;
+ let executor = location.executor_meta.ok_or_else(|| {
+ BallistaError::Internal("partition location has no executor metadata".to_string())
+ })?;
+
+ let action = BallistaAction::FetchPartition {
+ job_id: partition_id.job_id.into(),
+ stage_id: partition_id.stage_id as usize,
+ partition_id: partition_id.partition_id as usize,
+ host: executor.host,
+ port: executor.port as u16,
+ file_id: location.file_id,
+ layout,
+ // A Flight client wants the whole partition: the data file, in full.
+ file_kind: ShuffleFileKind::Data,
+ byte_ranges: vec![],
+ };
+
+ let encoded: protobuf::Action = action.try_into()?;
+ Ok(StatementHandle::Partition(encoded.encode_to_vec()))
+}
+
+fn build_flight_info(
+ schema: &Schema,
+ endpoints: Vec<FlightEndpoint>,
+ descriptor: FlightDescriptor,
+) -> Result<FlightInfo, Status> {
+ FlightInfo::new()
+ .try_with_schema(schema)
+ .map_err(|e| Status::internal(format!("failed to encode result schema: {e}")))
+ .map(|info| info.with_descriptor(descriptor).with_endpoints(endpoints))
+}
+
+/// Streams a single metadata batch back to the client.
+fn one_batch_response(batch: RecordBatch) -> Response<DoGetStream> {
+ batch_response(batch.schema(), vec![batch])
+}
+
+fn batch_response(schema: SchemaRef, batches: Vec<RecordBatch>) -> Response<DoGetStream> {
+ let stream = FlightDataEncoderBuilder::new()
+ .with_schema(schema)
+ .build(futures::stream::iter(batches.into_iter().map(Ok)))
+ .map_err(|e| Status::internal(format!("failed to encode results: {e}")));
+
+ Response::new(Box::pin(stream) as DoGetStream)
+}
+
+fn bearer_token(metadata: &MetadataMap) -> Option<String> {
+ let value = metadata.get("authorization")?.to_str().ok()?;
+ value
+ .strip_prefix("Bearer ")
+ .or_else(|| value.strip_prefix("bearer "))
+ .map(str::to_string)
+}
+
+fn invalid_ticket(e: impl std::fmt::Display) -> Status {
+ Status::invalid_argument(format!("invalid ticket: {e}"))
+}
+
+fn query_failed(e: BallistaError) -> Status {
+ // A failed job is the client's problem to see, not an opaque 500.
+ Status::internal(format!("query execution failed: {e}"))
+}
+
+fn encode_schema(schema: &Schema) -> Result<Vec<u8>, Status> {
+ let message: IpcMessage = SchemaAsIpc::new(schema, &IpcWriteOptions::default())
+ .try_into()
+ .map_err(|e| Status::internal(format!("failed to encode schema: {e}")))?;
+ Ok(message.0.to_vec())
+}
+
+#[tonic::async_trait]
+impl<B: QueryBackend> FlightSqlService for BallistaFlightSqlService<B> {
+ type FlightService = BallistaFlightSqlService<B>;
+
+ async fn do_handshake(
+ &self,
+ request: Request<Streaming<HandshakeRequest>>,
+ ) -> Result<
+ Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
+ Status,
+ > {
+ let identity = self.auth.authenticate(request.metadata()).await?;
+
+ let token = Uuid::new_v4().to_string();
+ let session_id = format!("flight-sql-{}", Uuid::new_v4());
+
+ // Build the session eagerly so a failure surfaces at handshake time
+ // rather than on the client's first query, and so the first query does
+ // not pay for it.
+ self.open_session(&session_id).await?;
+ self.store.insert_session(token.clone(), session_id.clone());
+
+ log::debug!(
+ "flight-sql: handshake for {:?} bound to session {session_id}",
+ identity.user
+ );
+
+ let result = HandshakeResponse {
+ protocol_version: 0,
+ payload: token.clone().into(),
+ };
+ let stream = futures::stream::once(async move { Ok(result) });
+
+ let mut response = Response::new(Box::pin(stream) as _);
+ response.metadata_mut().insert(
+ "authorization",
+ format!("Bearer {token}")
+ .parse()
+ .map_err(|_| Status::internal("failed to encode session token"))?,
+ );
+ Ok(response)
+ }
+
+ /// Redeems tickets that are not Flight SQL commands.
+ ///
+ /// Ballista's own Rust client fetches shuffle output with a
+ /// `ballista.protobuf.Action` ticket. When the Flight SQL frontend is
+ /// mounted it replaces the standalone proxy on the scheduler's port, so it
+ /// has to keep serving those tickets.
+ async fn do_get_fallback(
Review Comment:
**Minor.** The ticket is decoded here only to return `InvalidArgument`, and then `BallistaFlightProxyService::do_get` decodes it again. Mapping the decode error to `Status::invalid_argument` inside the proxy removes the second decode and gives the plain proxy the same error code.
##########
ballista/flight-sql/tests/service.rs:
##########
@@ -0,0 +1,593 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Protocol-level tests for the Flight SQL frontend, driven through the
+//! `FlightSqlService` trait against a stub backend.
+//!
+//! These deliberately avoid a live cluster so they run anywhere; the
+//! scheduler's `tests/flight_sql.rs` covers the same surface against real
+//! executors. Everything here is asserted the way a client sees it — tickets
+//! are decoded from the wire bytes rather than through crate internals — so
+//! the tests fail if the wire format changes.
+
+use std::sync::Arc;
+
+use arrow_flight::decode::FlightRecordBatchStream;
+use arrow_flight::error::FlightError;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCreatePreparedStatementRequest, Any, CommandGetCatalogs, CommandGetTables,
+ CommandStatementQuery, TicketStatementQuery,
+};
+use arrow_flight::{Action, FlightDescriptor, FlightInfo, Ticket};
+use async_trait::async_trait;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::serde::decode_protobuf;
+use ballista_core::serde::protobuf::{ExecutorMetadata, PartitionId, PartitionLocation};
+use ballista_core::serde::scheduler::{
+ Action as BallistaAction, ShuffleFileKind, ShuffleLayout,
+};
+use ballista_flight_sql::backend::{QueryBackend, QueryResult};
+use ballista_flight_sql::{
+ ANONYMOUS_SESSION, AnonymousAuthenticator, Authenticator, BallistaFlightSqlService,
+ Identity,
+};
+use dashmap::DashMap;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::empty::EmptyTable;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+use futures::TryStreamExt;
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Status};
+
+/// A backend that never contacts a cluster: it hands out real
+/// `SessionContext`s so planning is genuine, and fabricates the partition
+/// locations a completed job would have produced.
+#[derive(Default)]
+struct StubBackend {
+ sessions: DashMap<String, Arc<SessionContext>>,
+ executed: DashMap<String, usize>,
+ cancelled: DashMap<String, usize>,
+}
+
+const JOB_ID: &str = "job-1";
+
+#[async_trait]
+impl QueryBackend for StubBackend {
+ async fn session(&self, session_id: &str) -> Result<Arc<SessionContext>> {
+ Ok(self
+ .sessions
+ .entry(session_id.to_string())
+ .or_insert_with(|| Arc::new(SessionContext::new()))
+ .clone())
+ }
+
+ async fn close_session(&self, session_id: &str) -> Result<()> {
+ self.sessions.remove(session_id);
+ Ok(())
+ }
+
+ async fn execute(
+ &self,
+ _job_name: &str,
+ _ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ ) -> Result<QueryResult> {
+ *self.executed.entry(JOB_ID.to_string()).or_insert(0) += 1;
+
+ Ok(QueryResult {
+ job_id: JOB_ID.to_string(),
+ schema: Arc::new(plan.schema().as_arrow().clone()),
+ partitions: vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: Some(PartitionId {
+ job_id: JOB_ID.to_string(),
+ stage_id: 3,
+ partition_id: 7,
+ }),
+ executor_meta: Some(ExecutorMetadata {
+ id: "executor-1".to_string(),
+ host: "executor-host".to_string(),
+ port: 50051,
+ grpc_port: 50052,
+ specification: None,
+ os_info: None,
+ }),
+ partition_stats: None,
+ file_id: None,
+ is_sort_shuffle: false,
+ }],
+ })
+ }
+
+ async fn cancel(&self, job_id: &str) -> Result<()> {
+ *self.cancelled.entry(job_id.to_string()).or_insert(0) += 1;
+ Ok(())
+ }
+}
+
+/// Rejects everything, to check that the frontend actually consults the
+/// authenticator and refuses tokenless requests when one is installed.
+struct DenyAll;
+
+#[async_trait]
+impl Authenticator for DenyAll {
+ async fn authenticate(
+ &self,
+ _headers: &MetadataMap,
+ ) -> std::result::Result<Identity, Status> {
+ Err(Status::unauthenticated("nope"))
+ }
+}
+
+fn make_service(backend: Arc<StubBackend>) -> BallistaFlightSqlService<StubBackend> {
+ BallistaFlightSqlService::new(
+ backend,
+ BallistaFlightProxyService::new(4_194_304, 4_194_304, false, None),
+ )
+}
+
+fn descriptor() -> Request<FlightDescriptor> {
+ Request::new(FlightDescriptor::new_cmd(vec![]))
+}
+
+/// `expect_err` needs `Debug` on the success type, and Flight's stream
+/// responses do not have it.
+fn expect_err<T>(result: std::result::Result<T, Status>, msg: &str) -> Status {
+ match result {
+ Ok(_) => panic!("{msg}"),
+ Err(status) => status,
+ }
+}
+
+/// Pulls the ticket out of an endpoint exactly as a Flight client would.
+fn statement_handle(info: &FlightInfo, index: usize) -> Vec<u8> {
+ let ticket = info.endpoint[index].ticket.as_ref().expect("ticket");
+ let any = Any::decode(&*ticket.ticket).expect("ticket is an Any");
+ let statement: TicketStatementQuery =
+ any.unpack().expect("unpackable").expect("statement ticket");
+ statement.statement_handle.to_vec()
+}
+
+fn ticket_for(info: &FlightInfo, index: usize) -> Request<Ticket> {
+ Request::new(info.endpoint[index].ticket.clone().expect("ticket"))
+}
+
+type DoGetStream = <BallistaFlightSqlService<StubBackend> as FlightService>::DoGetStream;
+
+async fn collect(response: tonic::Response<DoGetStream>) -> Vec<RecordBatch> {
+ let stream = response
+ .into_inner()
+ .map_err(|e| FlightError::Tonic(Box::new(e)));
+ FlightRecordBatchStream::new_from_flight_data(stream)
+ .try_collect()
+ .await
+ .expect("results decode")
+}
+
+async fn register_table(backend: &StubBackend, session_id: &str) {
+ let ctx = backend.session(session_id).await.unwrap();
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, true),
+ Field::new("name", DataType::Utf8, true),
+ ]));
+ ctx.register_table("people", Arc::new(EmptyTable::new(schema)))
+ .unwrap();
+}
+
+#[tokio::test]
+async fn select_produces_one_endpoint_per_partition_with_no_location() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query planned and submitted")
+ .into_inner();
+
+ assert_eq!(info.endpoint.len(), 1);
+ assert!(
+ info.endpoint[0].location.is_empty(),
+ "an endpoint with no location tells the client to reuse its existing \
+ connection; advertising the executor is what broke #1012"
+ );
+
+ // The ticket must carry the executor-facing fetch action, so the proxy can
+ // redeem it without any extra server-side lookup.
+ let handle = statement_handle(&info, 0);
+ let action = decode_protobuf(&handle[1..]).expect("handle wraps a Ballista action");
+ match action {
+ BallistaAction::FetchPartition {
+ job_id,
+ stage_id,
+ partition_id,
+ host,
+ port,
+ ..
+ } => {
+ assert_eq!(job_id.to_string(), JOB_ID);
+ assert_eq!(stage_id, 3);
+ assert_eq!(partition_id, 7);
+ assert_eq!(host, "executor-host");
+ assert_eq!(port, 50051);
+ }
+ }
+}
+
+#[tokio::test]
+async fn ddl_runs_on_the_scheduler_and_its_result_is_single_use() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE SCHEMA reporting".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("DDL executes")
+ .into_inner();
+
+ // DDL never reaches the cluster.
+ assert!(backend.executed.is_empty());
+
+ // The schema really was created in the session the client will query.
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ ctx.catalog("datafusion")
+ .unwrap()
+ .schema("reporting")
+ .is_some()
+ );
+
+ let handle = statement_handle(&info, 0);
+ assert_eq!(handle[0], 1, "DDL results use the local-result tag");
+
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.into(),
+ };
+ service
+ .do_get_statement(ticket.clone(), ticket_for(&info, 0))
+ .await
+ .expect("first fetch succeeds");
+
+ let err = expect_err(
+ service.do_get_statement(ticket, ticket_for(&info, 0)).await,
+ "a ticket is redeemable once",
+ );
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn writes_are_refused_rather_than_run_on_the_scheduler() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "INSERT INTO people VALUES (1, 'x')".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("the distributed write path is not implemented");
+
+ assert_eq!(err.code(), tonic::Code::Unimplemented);
+ assert!(backend.executed.is_empty());
+}
+
+/// `CREATE TABLE AS SELECT` arrives as DDL, and DDL runs on the scheduler. It
+/// must be refused rather than quietly executed on a single node.
+#[tokio::test]
+async fn ctas_is_refused_rather_than_run_on_one_node() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let err = expect_err(
+ service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE TABLE big AS SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await,
+ "CTAS would give up distribution",
+ );
+
+ assert_eq!(err.code(), tonic::Code::Unimplemented);
+
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ !ctx.table_exist("big").unwrap(),
+ "the refused statement must not have taken effect"
+ );
+}
+
+#[tokio::test]
+async fn catalog_metadata_reflects_the_session() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let command = CommandGetTables {
+ catalog: None,
+ db_schema_filter_pattern: None,
+ table_name_filter_pattern: None,
+ table_types: vec![],
+ include_schema: true,
+ };
+ let batches = collect(
+ service
+ .do_get_tables(command, Request::new(Ticket::default()))
+ .await
+ .expect("tables listed"),
+ )
+ .await;
+
+ let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(
+ rows, 1,
+ "the table registered in the session must be listed"
+ );
+
+ let batches = collect(
+ service
+ .do_get_catalogs(CommandGetCatalogs {}, Request::new(Ticket::default()))
+ .await
+ .expect("catalogs listed"),
+ )
+ .await;
+ let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(rows, 1);
+}
+
+#[tokio::test]
+async fn prepared_statements_are_planned_once_and_expire_on_close() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let prepared = service
+ .do_action_create_prepared_statement(
+ ActionCreatePreparedStatementRequest {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("statement prepared");
+
+ assert!(
+ !prepared.dataset_schema.is_empty(),
+ "clients need the result schema before executing"
+ );
+
+ let handle = prepared.prepared_statement_handle.clone();
+ service
+ .get_flight_info_prepared_statement(
+ arrow_flight::sql::CommandPreparedStatementQuery {
+ prepared_statement_handle: handle.clone(),
+ },
+ descriptor(),
+ )
+ .await
+ .expect("prepared statement executes");
+
+ service
+ .do_action_close_prepared_statement(
+ arrow_flight::sql::ActionClosePreparedStatementRequest {
+ prepared_statement_handle: handle.clone(),
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("closes");
+
+ let err = service
+ .get_flight_info_prepared_statement(
+ arrow_flight::sql::CommandPreparedStatementQuery {
+ prepared_statement_handle: handle,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("a closed handle must not linger");
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn an_authenticator_makes_tokenless_requests_fail() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend).with_authenticator(Arc::new(DenyAll));
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT 1".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("no anonymous fallback when an authenticator is installed");
+
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+#[tokio::test]
+async fn the_default_service_allows_anonymous_access() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+ assert!(service.allows_anonymous());
+
+ let strict = make_service(Arc::new(StubBackend::default()))
+ .with_authenticator(Arc::new(DenyAll));
+ assert!(!strict.allows_anonymous());
+
+ // And the shipped default really is the permissive one, so the scheduler's
+ // startup warning is not dead code.
+ assert!(AnonymousAuthenticator.allows_anonymous());
+}
+
+#[tokio::test]
+async fn cancelling_a_query_reaches_the_backend() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query submitted")
+ .into_inner();
+
+ service
+ .do_action_cancel_query(
+ arrow_flight::sql::ActionCancelQueryRequest {
+ info: info.encode_to_vec().into(),
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("cancellation accepted");
+
+ assert_eq!(
+ backend.cancelled.get(JOB_ID).map(|v| *v),
+ Some(1),
+ "cancel must be routed to the job the FlightInfo describes"
+ );
+}
+
+#[tokio::test]
+async fn unknown_tokens_are_rejected() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+
+ let mut request = descriptor();
+ request
+ .metadata_mut()
+ .insert("authorization", "Bearer not-a-real-token".parse().unwrap());
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT 1".to_string(),
+ transaction_id: None,
+ },
+ request,
+ )
+ .await
+ .expect_err("a stale token must not silently fall back to a shared session");
+
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+/// The frontend replaces the standalone Flight proxy when mounted, so
+/// Ballista's own client tickets must still be recognised. The fallback is
+/// only reachable if arrow-flight's dispatcher does not mistake a Ballista
+/// action for a Flight SQL command.
+#[tokio::test]
+async fn ballista_client_tickets_are_recognised_by_the_fallback() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+
+ let garbage = Ticket {
+ ticket: vec![0xff, 0xff, 0xff].into(),
+ };
+ let err = expect_err(
+ service
+ .do_get_fallback(Request::new(garbage), Any::default())
+ .await,
+ "an unrecognized ticket is a client error",
+ );
+ assert_eq!(err.code(), tonic::Code::InvalidArgument);
+
+ // A well-formed Ballista ticket gets past validation and is handed to the
+ // proxy, which then fails to dial the (nonexistent) executor. Reaching a
+ // connection error is the assertion: it proves the ticket was accepted.
+ let action = BallistaAction::FetchPartition {
+ job_id: "job".into(),
+ stage_id: 1,
+ partition_id: 0,
+ host: "127.0.0.1".to_string(),
+ port: 1,
+ file_id: None,
+ layout: ShuffleLayout::Passthrough,
+ file_kind: ShuffleFileKind::Data,
+ byte_ranges: vec![],
+ };
+ let encoded: ballista_core::serde::protobuf::Action = action.try_into().unwrap();
+ let err = expect_err(
+ service
+ .do_get_fallback(
+ Request::new(Ticket {
+ ticket: encoded.encode_to_vec().into(),
+ }),
+ Any::default(),
+ )
+ .await,
+ "no executor is listening",
+ );
+ assert_ne!(
+ err.code(),
+ tonic::Code::InvalidArgument,
+ "the ticket itself must be accepted: {err}"
+ );
+}
+
+#[tokio::test]
+async fn sessions_are_isolated_per_token() -> std::result::Result<(), BallistaError> {
Review Comment:
**Test.** This only exercises `StubBackend` and `SessionContext`, and `_service` is never used. Either drop it or make it go through the frontend: two handshakes, DDL in one session, and the table absent from the other.
##########
ballista/flight-sql/tests/service.rs:
##########
@@ -0,0 +1,593 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Protocol-level tests for the Flight SQL frontend, driven through the
+//! `FlightSqlService` trait against a stub backend.
+//!
+//! These deliberately avoid a live cluster so they run anywhere; the
+//! scheduler's `tests/flight_sql.rs` covers the same surface against real
+//! executors. Everything here is asserted the way a client sees it — tickets
+//! are decoded from the wire bytes rather than through crate internals — so
+//! the tests fail if the wire format changes.
+
+use std::sync::Arc;
+
+use arrow_flight::decode::FlightRecordBatchStream;
+use arrow_flight::error::FlightError;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCreatePreparedStatementRequest, Any, CommandGetCatalogs, CommandGetTables,
+ CommandStatementQuery, TicketStatementQuery,
+};
+use arrow_flight::{Action, FlightDescriptor, FlightInfo, Ticket};
+use async_trait::async_trait;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::serde::decode_protobuf;
+use ballista_core::serde::protobuf::{ExecutorMetadata, PartitionId, PartitionLocation};
+use ballista_core::serde::scheduler::{
+ Action as BallistaAction, ShuffleFileKind, ShuffleLayout,
+};
+use ballista_flight_sql::backend::{QueryBackend, QueryResult};
+use ballista_flight_sql::{
+ ANONYMOUS_SESSION, AnonymousAuthenticator, Authenticator, BallistaFlightSqlService,
+ Identity,
+};
+use dashmap::DashMap;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::empty::EmptyTable;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+use futures::TryStreamExt;
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Status};
+
+/// A backend that never contacts a cluster: it hands out real
+/// `SessionContext`s so planning is genuine, and fabricates the partition
+/// locations a completed job would have produced.
+#[derive(Default)]
+struct StubBackend {
+ sessions: DashMap<String, Arc<SessionContext>>,
+ executed: DashMap<String, usize>,
+ cancelled: DashMap<String, usize>,
+}
+
+const JOB_ID: &str = "job-1";
+
+#[async_trait]
+impl QueryBackend for StubBackend {
+ async fn session(&self, session_id: &str) -> Result<Arc<SessionContext>> {
+ Ok(self
+ .sessions
+ .entry(session_id.to_string())
+ .or_insert_with(|| Arc::new(SessionContext::new()))
+ .clone())
+ }
+
+ async fn close_session(&self, session_id: &str) -> Result<()> {
+ self.sessions.remove(session_id);
+ Ok(())
+ }
+
+ async fn execute(
+ &self,
+ _job_name: &str,
+ _ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ ) -> Result<QueryResult> {
+ *self.executed.entry(JOB_ID.to_string()).or_insert(0) += 1;
+
+ Ok(QueryResult {
+ job_id: JOB_ID.to_string(),
+ schema: Arc::new(plan.schema().as_arrow().clone()),
+ partitions: vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: Some(PartitionId {
+ job_id: JOB_ID.to_string(),
+ stage_id: 3,
+ partition_id: 7,
+ }),
+ executor_meta: Some(ExecutorMetadata {
+ id: "executor-1".to_string(),
+ host: "executor-host".to_string(),
+ port: 50051,
+ grpc_port: 50052,
+ specification: None,
+ os_info: None,
+ }),
+ partition_stats: None,
+ file_id: None,
+ is_sort_shuffle: false,
+ }],
+ })
+ }
+
+ async fn cancel(&self, job_id: &str) -> Result<()> {
+ *self.cancelled.entry(job_id.to_string()).or_insert(0) += 1;
+ Ok(())
+ }
+}
+
+/// Rejects everything, to check that the frontend actually consults the
+/// authenticator and refuses tokenless requests when one is installed.
+struct DenyAll;
+
+#[async_trait]
+impl Authenticator for DenyAll {
+ async fn authenticate(
+ &self,
+ _headers: &MetadataMap,
+ ) -> std::result::Result<Identity, Status> {
+ Err(Status::unauthenticated("nope"))
+ }
+}
+
+fn make_service(backend: Arc<StubBackend>) -> BallistaFlightSqlService<StubBackend> {
+ BallistaFlightSqlService::new(
+ backend,
+ BallistaFlightProxyService::new(4_194_304, 4_194_304, false, None),
+ )
+}
+
+fn descriptor() -> Request<FlightDescriptor> {
+ Request::new(FlightDescriptor::new_cmd(vec![]))
+}
+
+/// `expect_err` needs `Debug` on the success type, and Flight's stream
+/// responses do not have it.
+fn expect_err<T>(result: std::result::Result<T, Status>, msg: &str) -> Status {
+ match result {
+ Ok(_) => panic!("{msg}"),
+ Err(status) => status,
+ }
+}
+
+/// Pulls the ticket out of an endpoint exactly as a Flight client would.
+fn statement_handle(info: &FlightInfo, index: usize) -> Vec<u8> {
+ let ticket = info.endpoint[index].ticket.as_ref().expect("ticket");
+ let any = Any::decode(&*ticket.ticket).expect("ticket is an Any");
+ let statement: TicketStatementQuery =
+ any.unpack().expect("unpackable").expect("statement ticket");
+ statement.statement_handle.to_vec()
+}
+
+fn ticket_for(info: &FlightInfo, index: usize) -> Request<Ticket> {
+ Request::new(info.endpoint[index].ticket.clone().expect("ticket"))
+}
+
+type DoGetStream = <BallistaFlightSqlService<StubBackend> as FlightService>::DoGetStream;
+
+async fn collect(response: tonic::Response<DoGetStream>) -> Vec<RecordBatch> {
+ let stream = response
+ .into_inner()
+ .map_err(|e| FlightError::Tonic(Box::new(e)));
+ FlightRecordBatchStream::new_from_flight_data(stream)
+ .try_collect()
+ .await
+ .expect("results decode")
+}
+
+async fn register_table(backend: &StubBackend, session_id: &str) {
+ let ctx = backend.session(session_id).await.unwrap();
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, true),
+ Field::new("name", DataType::Utf8, true),
+ ]));
+ ctx.register_table("people", Arc::new(EmptyTable::new(schema)))
+ .unwrap();
+}
+
+#[tokio::test]
+async fn select_produces_one_endpoint_per_partition_with_no_location() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query planned and submitted")
+ .into_inner();
+
+ assert_eq!(info.endpoint.len(), 1);
+ assert!(
+ info.endpoint[0].location.is_empty(),
+ "an endpoint with no location tells the client to reuse its existing \
+ connection; advertising the executor is what broke #1012"
+ );
+
+ // The ticket must carry the executor-facing fetch action, so the proxy can
+ // redeem it without any extra server-side lookup.
+ let handle = statement_handle(&info, 0);
+ let action = decode_protobuf(&handle[1..]).expect("handle wraps a Ballista action");
+ match action {
+ BallistaAction::FetchPartition {
+ job_id,
+ stage_id,
+ partition_id,
+ host,
+ port,
+ ..
+ } => {
+ assert_eq!(job_id.to_string(), JOB_ID);
+ assert_eq!(stage_id, 3);
+ assert_eq!(partition_id, 7);
+ assert_eq!(host, "executor-host");
+ assert_eq!(port, 50051);
+ }
+ }
+}
+
+#[tokio::test]
+async fn ddl_runs_on_the_scheduler_and_its_result_is_single_use() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE SCHEMA reporting".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("DDL executes")
+ .into_inner();
+
+ // DDL never reaches the cluster.
+ assert!(backend.executed.is_empty());
+
+ // The schema really was created in the session the client will query.
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ ctx.catalog("datafusion")
+ .unwrap()
+ .schema("reporting")
+ .is_some()
+ );
+
+ let handle = statement_handle(&info, 0);
+ assert_eq!(handle[0], 1, "DDL results use the local-result tag");
+
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.into(),
+ };
+ service
+ .do_get_statement(ticket.clone(), ticket_for(&info, 0))
+ .await
+ .expect("first fetch succeeds");
+
+ let err = expect_err(
+ service.do_get_statement(ticket, ticket_for(&info, 0)).await,
+ "a ticket is redeemable once",
+ );
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn writes_are_refused_rather_than_run_on_the_scheduler() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "INSERT INTO people VALUES (1, 'x')".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("the distributed write path is not implemented");
+
+ assert_eq!(err.code(), tonic::Code::Unimplemented);
+ assert!(backend.executed.is_empty());
+}
+
+/// `CREATE TABLE AS SELECT` arrives as DDL, and DDL runs on the scheduler. It
+/// must be refused rather than quietly executed on a single node.
+#[tokio::test]
+async fn ctas_is_refused_rather_than_run_on_one_node() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let err = expect_err(
+ service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE TABLE big AS SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await,
+ "CTAS would give up distribution",
+ );
+
+ assert_eq!(err.code(), tonic::Code::Unimplemented);
+
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ !ctx.table_exist("big").unwrap(),
+ "the refused statement must not have taken effect"
+ );
+}
+
+#[tokio::test]
+async fn catalog_metadata_reflects_the_session() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let command = CommandGetTables {
+ catalog: None,
+ db_schema_filter_pattern: None,
+ table_name_filter_pattern: None,
+ table_types: vec![],
+ include_schema: true,
+ };
+ let batches = collect(
+ service
+ .do_get_tables(command, Request::new(Ticket::default()))
+ .await
+ .expect("tables listed"),
+ )
+ .await;
+
+ let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(
+ rows, 1,
+ "the table registered in the session must be listed"
+ );
+
+ let batches = collect(
+ service
+ .do_get_catalogs(CommandGetCatalogs {}, Request::new(Ticket::default()))
+ .await
+ .expect("catalogs listed"),
+ )
+ .await;
+ let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(rows, 1);
+}
+
+#[tokio::test]
+async fn prepared_statements_are_planned_once_and_expire_on_close() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let prepared = service
+ .do_action_create_prepared_statement(
+ ActionCreatePreparedStatementRequest {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("statement prepared");
+
+ assert!(
+ !prepared.dataset_schema.is_empty(),
+ "clients need the result schema before executing"
+ );
+
+ let handle = prepared.prepared_statement_handle.clone();
+ service
+ .get_flight_info_prepared_statement(
+ arrow_flight::sql::CommandPreparedStatementQuery {
+ prepared_statement_handle: handle.clone(),
+ },
+ descriptor(),
+ )
+ .await
+ .expect("prepared statement executes");
+
+ service
+ .do_action_close_prepared_statement(
+ arrow_flight::sql::ActionClosePreparedStatementRequest {
+ prepared_statement_handle: handle.clone(),
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("closes");
+
+ let err = service
+ .get_flight_info_prepared_statement(
+ arrow_flight::sql::CommandPreparedStatementQuery {
+ prepared_statement_handle: handle,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("a closed handle must not linger");
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn an_authenticator_makes_tokenless_requests_fail() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend).with_authenticator(Arc::new(DenyAll));
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT 1".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("no anonymous fallback when an authenticator is installed");
+
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+#[tokio::test]
+async fn the_default_service_allows_anonymous_access() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+ assert!(service.allows_anonymous());
+
+ let strict = make_service(Arc::new(StubBackend::default()))
+ .with_authenticator(Arc::new(DenyAll));
+ assert!(!strict.allows_anonymous());
+
+ // And the shipped default really is the permissive one, so the scheduler's
+ // startup warning is not dead code.
+ assert!(AnonymousAuthenticator.allows_anonymous());
+}
+
+#[tokio::test]
+async fn cancelling_a_query_reaches_the_backend() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query submitted")
+ .into_inner();
+
+ service
+ .do_action_cancel_query(
+ arrow_flight::sql::ActionCancelQueryRequest {
+ info: info.encode_to_vec().into(),
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("cancellation accepted");
+
+ assert_eq!(
+ backend.cancelled.get(JOB_ID).map(|v| *v),
+ Some(1),
+ "cancel must be routed to the job the FlightInfo describes"
+ );
+}
+
+#[tokio::test]
+async fn unknown_tokens_are_rejected() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+
+ let mut request = descriptor();
+ request
+ .metadata_mut()
+ .insert("authorization", "Bearer not-a-real-token".parse().unwrap());
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT 1".to_string(),
+ transaction_id: None,
+ },
+ request,
+ )
+ .await
+ .expect_err("a stale token must not silently fall back to a shared session");
+
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+/// The frontend replaces the standalone Flight proxy when mounted, so
+/// Ballista's own client tickets must still be recognised. The fallback is
+/// only reachable if arrow-flight's dispatcher does not mistake a Ballista
+/// action for a Flight SQL command.
+#[tokio::test]
+async fn ballista_client_tickets_are_recognised_by_the_fallback() {
Review Comment:
**Test.** The doc comment says this checks that arrow-flight's dispatcher routes a Ballista ticket to the fallback, but the test calls `do_get_fallback` directly, so the dispatcher never runs. `ballista_tickets_route_to_the_fallback` in `service.rs` covers half of it through `Any::decode`. One test that calls `FlightService::do_get(&service, ...)` with a Ballista ticket would cover both, and the other can go.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#issuecomment-5858146294
Here's where the rest of the review summary landed.
**Missing tests.** The stub backend now returns three partitions, and `select_produces_one_endpoint_per_partition_with_no_location` checks that each endpoint names a different one. `an_authenticator_makes_tokenless_requests_fail` covers `DoGet` and the prepared statement handlers. `EXECUTE` is in `statements_that_cannot_be_distributed_are_refused`. The `PushStaged` test went to #2492 with the fix.
**Duplicate tests.** `the_default_service_allows_anonymous_access` is now one assertion in the authenticator test. The `GetTables` test in `tests/service.rs` is gone, so `GetTables` is covered once in `metadata.rs` and once end to end. `sql_info_is_buildable`, `xdbc_type_info_is_buildable` and `plain_ddl_runs_on_the_scheduler` are removed. `local_results_are_single_use` folded into the session binding test, and the two bearer token tests are one.
**Comments.** All four are fixed. `LocalResult` no longer mentions DML. `query_failed` no longer claims to avoid a 500. The `QueryBackend::session` doc says the frontend caches the context, so a backend may build a fresh one on every call. The notes about the pre-46.0.0 implementation are out of the code and tests.
**Robustness.** The reaper starts on the first request, so the service can be built outside a Tokio runtime. A handshake no longer builds a `SessionContext`. The context is built on first use. I haven't added count caps. Everything in the store expires after the 30 minute idle TTL, and local results are dropped as soon as they're redeemed. The same port also lets any client create scheduler sessions with no cap through `SchedulerGrpc::create_update_session`. So a cap here wouldn't bound memory against a hostile client, which is why the docs say to keep the port on a trusted network. I'm open to adding caps if you think they're worth it for runaway clients.
**Performance.** Agreed, this is a real limit. Endpoints carry no location so clients never need to reach an executor (#1012, #1349), which makes the scheduler the funnel for result bytes. The per-`DoGet` connection is how the existing `BallistaFlightProxyService` works, and this reuses it. Pooling connections in the proxy, or an opt-in to advertise executor locations to clients that can reach them, would both help. I'd like to leave those for a follow-up with a benchmark that uses large results. The TPC-H runner in #2496 times `DoGet` separately, but TPC-H answers are too small to show this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#issuecomment-5799570027
Thanks @andygrove the direction makes a lot of sense to feed downstream BI/reporting/visualize layers using Arrow format, considering the layer already supports the Arrow capable endpoint. Apache Superset https://github.com/apache/superset is one of the good examples.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416#discussion_r4116285230
##########
ballista/scheduler/src/flight_sql.rs:
##########
@@ -0,0 +1,111 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Binds the [`ballista_flight_sql`] frontend to this scheduler.
+//!
+//! The frontend does not know about `SchedulerServer`; this is the only place
+//! the two meet, which is what keeps the Flight SQL code from re-acquiring the
+//! coupling that made the pre-46.0.0 implementation unmaintainable.
+
+use std::sync::Arc;
+
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::serde::protobuf::{JobStatus, SuccessfulJob, job_status};
+use ballista_flight_sql::backend::{QueryBackend, QueryResult};
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+use datafusion_proto::logical_plan::AsLogicalPlan;
+use datafusion_proto::physical_plan::AsExecutionPlan;
+
+use crate::scheduler_server::SchedulerServer;
+
+/// Buffer for the job status stream. Statuses are consumed as fast as they
+/// arrive, so this only has to absorb bursts.
+const STATUS_BUFFER: usize = 16;
+
+#[async_trait::async_trait]
+impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> QueryBackend
+ for SchedulerServer<T, U>
+{
+ async fn session(&self, session_id: &str) -> Result<Arc<SessionContext>> {
+ let config = self.state.session_manager.produce_config();
+ self.state
+ .session_manager
+ .create_or_update_session(session_id, &config)
+ .await
+ }
+
+ async fn close_session(&self, session_id: &str) -> Result<()> {
+ self.state.session_manager.remove_session(session_id).await
+ }
+
+ async fn execute(
+ &self,
+ job_name: &str,
+ ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ ) -> Result<QueryResult> {
+ // The schema handed to the client has to be the one the shuffle files
+ // actually carry, which is the physical plan's, not the logical plan's:
+ // the two disagree about nullability often enough to matter (`version()`
+ // is nullable logically and not-null in the data; a list literal goes the
+ // other way), and a Flight SQL client that is promised one schema and
+ // handed another -- ADBC among them -- rejects the result outright.
+ // Planning twice costs a few milliseconds against a job that is about to
+ // run on the cluster.
+ let schema = ctx.state().create_physical_plan(&plan).await?.schema();
Review Comment:
Yes. `result_schema` now reads the schema from the job's final stage once it succeeds, since that's the plan that wrote the files. It only plans again if the graph is no longer available. I also just fixed the `QueryResult::schema` doc, which still said the schema came from the submitted plan.
##########
ballista/flight-sql/tests/service.rs:
##########
@@ -0,0 +1,593 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Protocol-level tests for the Flight SQL frontend, driven through the
+//! `FlightSqlService` trait against a stub backend.
+//!
+//! These deliberately avoid a live cluster so they run anywhere; the
+//! scheduler's `tests/flight_sql.rs` covers the same surface against real
+//! executors. Everything here is asserted the way a client sees it — tickets
+//! are decoded from the wire bytes rather than through crate internals — so
+//! the tests fail if the wire format changes.
+
+use std::sync::Arc;
+
+use arrow_flight::decode::FlightRecordBatchStream;
+use arrow_flight::error::FlightError;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCreatePreparedStatementRequest, Any, CommandGetCatalogs, CommandGetTables,
+ CommandStatementQuery, TicketStatementQuery,
+};
+use arrow_flight::{Action, FlightDescriptor, FlightInfo, Ticket};
+use async_trait::async_trait;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::serde::decode_protobuf;
+use ballista_core::serde::protobuf::{ExecutorMetadata, PartitionId, PartitionLocation};
+use ballista_core::serde::scheduler::{
+ Action as BallistaAction, ShuffleFileKind, ShuffleLayout,
+};
+use ballista_flight_sql::backend::{QueryBackend, QueryResult};
+use ballista_flight_sql::{
+ ANONYMOUS_SESSION, AnonymousAuthenticator, Authenticator, BallistaFlightSqlService,
+ Identity,
+};
+use dashmap::DashMap;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::empty::EmptyTable;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+use futures::TryStreamExt;
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Status};
+
+/// A backend that never contacts a cluster: it hands out real
+/// `SessionContext`s so planning is genuine, and fabricates the partition
+/// locations a completed job would have produced.
+#[derive(Default)]
+struct StubBackend {
+ sessions: DashMap<String, Arc<SessionContext>>,
+ executed: DashMap<String, usize>,
+ cancelled: DashMap<String, usize>,
+}
+
+const JOB_ID: &str = "job-1";
+
+#[async_trait]
+impl QueryBackend for StubBackend {
+ async fn session(&self, session_id: &str) -> Result<Arc<SessionContext>> {
+ Ok(self
+ .sessions
+ .entry(session_id.to_string())
+ .or_insert_with(|| Arc::new(SessionContext::new()))
+ .clone())
+ }
+
+ async fn close_session(&self, session_id: &str) -> Result<()> {
+ self.sessions.remove(session_id);
+ Ok(())
+ }
+
+ async fn execute(
+ &self,
+ _job_name: &str,
+ _ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ ) -> Result<QueryResult> {
+ *self.executed.entry(JOB_ID.to_string()).or_insert(0) += 1;
+
+ Ok(QueryResult {
+ job_id: JOB_ID.to_string(),
+ schema: Arc::new(plan.schema().as_arrow().clone()),
+ partitions: vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: Some(PartitionId {
+ job_id: JOB_ID.to_string(),
+ stage_id: 3,
+ partition_id: 7,
+ }),
+ executor_meta: Some(ExecutorMetadata {
+ id: "executor-1".to_string(),
+ host: "executor-host".to_string(),
+ port: 50051,
+ grpc_port: 50052,
+ specification: None,
+ os_info: None,
+ }),
+ partition_stats: None,
+ file_id: None,
+ is_sort_shuffle: false,
+ }],
+ })
+ }
+
+ async fn cancel(&self, job_id: &str) -> Result<()> {
+ *self.cancelled.entry(job_id.to_string()).or_insert(0) += 1;
+ Ok(())
+ }
+}
+
+/// Rejects everything, to check that the frontend actually consults the
+/// authenticator and refuses tokenless requests when one is installed.
+struct DenyAll;
+
+#[async_trait]
+impl Authenticator for DenyAll {
+ async fn authenticate(
+ &self,
+ _headers: &MetadataMap,
+ ) -> std::result::Result<Identity, Status> {
+ Err(Status::unauthenticated("nope"))
+ }
+}
+
+fn make_service(backend: Arc<StubBackend>) -> BallistaFlightSqlService<StubBackend> {
+ BallistaFlightSqlService::new(
+ backend,
+ BallistaFlightProxyService::new(4_194_304, 4_194_304, false, None),
+ )
+}
+
+fn descriptor() -> Request<FlightDescriptor> {
+ Request::new(FlightDescriptor::new_cmd(vec![]))
+}
+
+/// `expect_err` needs `Debug` on the success type, and Flight's stream
+/// responses do not have it.
+fn expect_err<T>(result: std::result::Result<T, Status>, msg: &str) -> Status {
+ match result {
+ Ok(_) => panic!("{msg}"),
+ Err(status) => status,
+ }
+}
+
+/// Pulls the ticket out of an endpoint exactly as a Flight client would.
+fn statement_handle(info: &FlightInfo, index: usize) -> Vec<u8> {
+ let ticket = info.endpoint[index].ticket.as_ref().expect("ticket");
+ let any = Any::decode(&*ticket.ticket).expect("ticket is an Any");
+ let statement: TicketStatementQuery =
+ any.unpack().expect("unpackable").expect("statement ticket");
+ statement.statement_handle.to_vec()
+}
+
+fn ticket_for(info: &FlightInfo, index: usize) -> Request<Ticket> {
+ Request::new(info.endpoint[index].ticket.clone().expect("ticket"))
+}
+
+type DoGetStream = <BallistaFlightSqlService<StubBackend> as FlightService>::DoGetStream;
+
+async fn collect(response: tonic::Response<DoGetStream>) -> Vec<RecordBatch> {
+ let stream = response
+ .into_inner()
+ .map_err(|e| FlightError::Tonic(Box::new(e)));
+ FlightRecordBatchStream::new_from_flight_data(stream)
+ .try_collect()
+ .await
+ .expect("results decode")
+}
+
+async fn register_table(backend: &StubBackend, session_id: &str) {
+ let ctx = backend.session(session_id).await.unwrap();
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, true),
+ Field::new("name", DataType::Utf8, true),
+ ]));
+ ctx.register_table("people", Arc::new(EmptyTable::new(schema)))
+ .unwrap();
+}
+
+#[tokio::test]
+async fn select_produces_one_endpoint_per_partition_with_no_location() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query planned and submitted")
+ .into_inner();
+
+ assert_eq!(info.endpoint.len(), 1);
+ assert!(
+ info.endpoint[0].location.is_empty(),
+ "an endpoint with no location tells the client to reuse its existing \
+ connection; advertising the executor is what broke #1012"
+ );
+
+ // The ticket must carry the executor-facing fetch action, so the proxy can
+ // redeem it without any extra server-side lookup.
+ let handle = statement_handle(&info, 0);
+ let action = decode_protobuf(&handle[1..]).expect("handle wraps a Ballista action");
+ match action {
+ BallistaAction::FetchPartition {
+ job_id,
+ stage_id,
+ partition_id,
+ host,
+ port,
+ ..
+ } => {
+ assert_eq!(job_id.to_string(), JOB_ID);
+ assert_eq!(stage_id, 3);
+ assert_eq!(partition_id, 7);
+ assert_eq!(host, "executor-host");
+ assert_eq!(port, 50051);
+ }
+ }
+}
+
+#[tokio::test]
+async fn ddl_runs_on_the_scheduler_and_its_result_is_single_use() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE SCHEMA reporting".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("DDL executes")
+ .into_inner();
+
+ // DDL never reaches the cluster.
+ assert!(backend.executed.is_empty());
+
+ // The schema really was created in the session the client will query.
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ ctx.catalog("datafusion")
+ .unwrap()
+ .schema("reporting")
+ .is_some()
+ );
+
+ let handle = statement_handle(&info, 0);
+ assert_eq!(handle[0], 1, "DDL results use the local-result tag");
+
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.into(),
+ };
+ service
+ .do_get_statement(ticket.clone(), ticket_for(&info, 0))
+ .await
+ .expect("first fetch succeeds");
+
+ let err = expect_err(
+ service.do_get_statement(ticket, ticket_for(&info, 0)).await,
+ "a ticket is redeemable once",
+ );
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn writes_are_refused_rather_than_run_on_the_scheduler() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "INSERT INTO people VALUES (1, 'x')".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("the distributed write path is not implemented");
+
+ assert_eq!(err.code(), tonic::Code::Unimplemented);
+ assert!(backend.executed.is_empty());
+}
+
+/// `CREATE TABLE AS SELECT` arrives as DDL, and DDL runs on the scheduler. It
+/// must be refused rather than quietly executed on a single node.
+#[tokio::test]
+async fn ctas_is_refused_rather_than_run_on_one_node() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let err = expect_err(
+ service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE TABLE big AS SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await,
+ "CTAS would give up distribution",
+ );
+
+ assert_eq!(err.code(), tonic::Code::Unimplemented);
+
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ !ctx.table_exist("big").unwrap(),
+ "the refused statement must not have taken effect"
+ );
+}
+
+#[tokio::test]
+async fn catalog_metadata_reflects_the_session() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let command = CommandGetTables {
+ catalog: None,
+ db_schema_filter_pattern: None,
+ table_name_filter_pattern: None,
+ table_types: vec![],
+ include_schema: true,
+ };
+ let batches = collect(
+ service
+ .do_get_tables(command, Request::new(Ticket::default()))
+ .await
+ .expect("tables listed"),
+ )
+ .await;
+
+ let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(
+ rows, 1,
+ "the table registered in the session must be listed"
+ );
+
+ let batches = collect(
+ service
+ .do_get_catalogs(CommandGetCatalogs {}, Request::new(Ticket::default()))
+ .await
+ .expect("catalogs listed"),
+ )
+ .await;
+ let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(rows, 1);
+}
+
+#[tokio::test]
+async fn prepared_statements_are_planned_once_and_expire_on_close() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let prepared = service
+ .do_action_create_prepared_statement(
+ ActionCreatePreparedStatementRequest {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("statement prepared");
+
+ assert!(
+ !prepared.dataset_schema.is_empty(),
+ "clients need the result schema before executing"
+ );
+
+ let handle = prepared.prepared_statement_handle.clone();
+ service
+ .get_flight_info_prepared_statement(
+ arrow_flight::sql::CommandPreparedStatementQuery {
+ prepared_statement_handle: handle.clone(),
+ },
+ descriptor(),
+ )
+ .await
+ .expect("prepared statement executes");
+
+ service
+ .do_action_close_prepared_statement(
+ arrow_flight::sql::ActionClosePreparedStatementRequest {
+ prepared_statement_handle: handle.clone(),
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("closes");
+
+ let err = service
+ .get_flight_info_prepared_statement(
+ arrow_flight::sql::CommandPreparedStatementQuery {
+ prepared_statement_handle: handle,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("a closed handle must not linger");
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn an_authenticator_makes_tokenless_requests_fail() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend).with_authenticator(Arc::new(DenyAll));
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT 1".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("no anonymous fallback when an authenticator is installed");
+
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+#[tokio::test]
+async fn the_default_service_allows_anonymous_access() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+ assert!(service.allows_anonymous());
+
+ let strict = make_service(Arc::new(StubBackend::default()))
+ .with_authenticator(Arc::new(DenyAll));
+ assert!(!strict.allows_anonymous());
+
+ // And the shipped default really is the permissive one, so the scheduler's
+ // startup warning is not dead code.
+ assert!(AnonymousAuthenticator.allows_anonymous());
+}
+
+#[tokio::test]
+async fn cancelling_a_query_reaches_the_backend() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query submitted")
+ .into_inner();
+
+ service
+ .do_action_cancel_query(
+ arrow_flight::sql::ActionCancelQueryRequest {
+ info: info.encode_to_vec().into(),
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("cancellation accepted");
+
+ assert_eq!(
+ backend.cancelled.get(JOB_ID).map(|v| *v),
+ Some(1),
+ "cancel must be routed to the job the FlightInfo describes"
+ );
+}
+
+#[tokio::test]
+async fn unknown_tokens_are_rejected() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+
+ let mut request = descriptor();
+ request
+ .metadata_mut()
+ .insert("authorization", "Bearer not-a-real-token".parse().unwrap());
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT 1".to_string(),
+ transaction_id: None,
+ },
+ request,
+ )
+ .await
+ .expect_err("a stale token must not silently fall back to a shared session");
+
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+/// The frontend replaces the standalone Flight proxy when mounted, so
+/// Ballista's own client tickets must still be recognised. The fallback is
+/// only reachable if arrow-flight's dispatcher does not mistake a Ballista
+/// action for a Flight SQL command.
+#[tokio::test]
+async fn ballista_client_tickets_are_recognised_by_the_fallback() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+
+ let garbage = Ticket {
+ ticket: vec![0xff, 0xff, 0xff].into(),
+ };
+ let err = expect_err(
+ service
+ .do_get_fallback(Request::new(garbage), Any::default())
+ .await,
+ "an unrecognized ticket is a client error",
+ );
+ assert_eq!(err.code(), tonic::Code::InvalidArgument);
+
+ // A well-formed Ballista ticket gets past validation and is handed to the
+ // proxy, which then fails to dial the (nonexistent) executor. Reaching a
+ // connection error is the assertion: it proves the ticket was accepted.
+ let action = BallistaAction::FetchPartition {
+ job_id: "job".into(),
+ stage_id: 1,
+ partition_id: 0,
+ host: "127.0.0.1".to_string(),
+ port: 1,
+ file_id: None,
+ layout: ShuffleLayout::Passthrough,
+ file_kind: ShuffleFileKind::Data,
+ byte_ranges: vec![],
+ };
+ let encoded: ballista_core::serde::protobuf::Action = action.try_into().unwrap();
+ let err = expect_err(
+ service
+ .do_get_fallback(
+ Request::new(Ticket {
+ ticket: encoded.encode_to_vec().into(),
+ }),
+ Any::default(),
+ )
+ .await,
+ "no executor is listening",
+ );
+ assert_ne!(
+ err.code(),
+ tonic::Code::InvalidArgument,
+ "the ticket itself must be accepted: {err}"
+ );
+}
+
+#[tokio::test]
+async fn sessions_are_isolated_per_token() -> std::result::Result<(), BallistaError> {
Review Comment:
Rewritten to go through the frontend. `sessions_are_isolated_per_token` does two real handshakes over a served channel and creates a view in one session. It then checks that the other session can't see the view, can't redeem its result, and can't run or close its prepared statement.
##########
ballista/flight-sql/tests/service.rs:
##########
@@ -0,0 +1,593 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Protocol-level tests for the Flight SQL frontend, driven through the
+//! `FlightSqlService` trait against a stub backend.
+//!
+//! These deliberately avoid a live cluster so they run anywhere; the
+//! scheduler's `tests/flight_sql.rs` covers the same surface against real
+//! executors. Everything here is asserted the way a client sees it — tickets
+//! are decoded from the wire bytes rather than through crate internals — so
+//! the tests fail if the wire format changes.
+
+use std::sync::Arc;
+
+use arrow_flight::decode::FlightRecordBatchStream;
+use arrow_flight::error::FlightError;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCreatePreparedStatementRequest, Any, CommandGetCatalogs, CommandGetTables,
+ CommandStatementQuery, TicketStatementQuery,
+};
+use arrow_flight::{Action, FlightDescriptor, FlightInfo, Ticket};
+use async_trait::async_trait;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::serde::decode_protobuf;
+use ballista_core::serde::protobuf::{ExecutorMetadata, PartitionId, PartitionLocation};
+use ballista_core::serde::scheduler::{
+ Action as BallistaAction, ShuffleFileKind, ShuffleLayout,
+};
+use ballista_flight_sql::backend::{QueryBackend, QueryResult};
+use ballista_flight_sql::{
+ ANONYMOUS_SESSION, AnonymousAuthenticator, Authenticator, BallistaFlightSqlService,
+ Identity,
+};
+use dashmap::DashMap;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::empty::EmptyTable;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+use futures::TryStreamExt;
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Status};
+
+/// A backend that never contacts a cluster: it hands out real
+/// `SessionContext`s so planning is genuine, and fabricates the partition
+/// locations a completed job would have produced.
+#[derive(Default)]
+struct StubBackend {
+ sessions: DashMap<String, Arc<SessionContext>>,
+ executed: DashMap<String, usize>,
+ cancelled: DashMap<String, usize>,
+}
+
+const JOB_ID: &str = "job-1";
+
+#[async_trait]
+impl QueryBackend for StubBackend {
+ async fn session(&self, session_id: &str) -> Result<Arc<SessionContext>> {
+ Ok(self
+ .sessions
+ .entry(session_id.to_string())
+ .or_insert_with(|| Arc::new(SessionContext::new()))
+ .clone())
+ }
+
+ async fn close_session(&self, session_id: &str) -> Result<()> {
+ self.sessions.remove(session_id);
+ Ok(())
+ }
+
+ async fn execute(
+ &self,
+ _job_name: &str,
+ _ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ ) -> Result<QueryResult> {
+ *self.executed.entry(JOB_ID.to_string()).or_insert(0) += 1;
+
+ Ok(QueryResult {
+ job_id: JOB_ID.to_string(),
+ schema: Arc::new(plan.schema().as_arrow().clone()),
+ partitions: vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: Some(PartitionId {
+ job_id: JOB_ID.to_string(),
+ stage_id: 3,
+ partition_id: 7,
+ }),
+ executor_meta: Some(ExecutorMetadata {
+ id: "executor-1".to_string(),
+ host: "executor-host".to_string(),
+ port: 50051,
+ grpc_port: 50052,
+ specification: None,
+ os_info: None,
+ }),
+ partition_stats: None,
+ file_id: None,
+ is_sort_shuffle: false,
+ }],
+ })
+ }
+
+ async fn cancel(&self, job_id: &str) -> Result<()> {
+ *self.cancelled.entry(job_id.to_string()).or_insert(0) += 1;
+ Ok(())
+ }
+}
+
+/// Rejects everything, to check that the frontend actually consults the
+/// authenticator and refuses tokenless requests when one is installed.
+struct DenyAll;
+
+#[async_trait]
+impl Authenticator for DenyAll {
+ async fn authenticate(
+ &self,
+ _headers: &MetadataMap,
+ ) -> std::result::Result<Identity, Status> {
+ Err(Status::unauthenticated("nope"))
+ }
+}
+
+fn make_service(backend: Arc<StubBackend>) -> BallistaFlightSqlService<StubBackend> {
+ BallistaFlightSqlService::new(
+ backend,
+ BallistaFlightProxyService::new(4_194_304, 4_194_304, false, None),
+ )
+}
+
+fn descriptor() -> Request<FlightDescriptor> {
+ Request::new(FlightDescriptor::new_cmd(vec![]))
+}
+
+/// `expect_err` needs `Debug` on the success type, and Flight's stream
+/// responses do not have it.
+fn expect_err<T>(result: std::result::Result<T, Status>, msg: &str) -> Status {
+ match result {
+ Ok(_) => panic!("{msg}"),
+ Err(status) => status,
+ }
+}
+
+/// Pulls the ticket out of an endpoint exactly as a Flight client would.
+fn statement_handle(info: &FlightInfo, index: usize) -> Vec<u8> {
+ let ticket = info.endpoint[index].ticket.as_ref().expect("ticket");
+ let any = Any::decode(&*ticket.ticket).expect("ticket is an Any");
+ let statement: TicketStatementQuery =
+ any.unpack().expect("unpackable").expect("statement ticket");
+ statement.statement_handle.to_vec()
+}
+
+fn ticket_for(info: &FlightInfo, index: usize) -> Request<Ticket> {
+ Request::new(info.endpoint[index].ticket.clone().expect("ticket"))
+}
+
+type DoGetStream = <BallistaFlightSqlService<StubBackend> as FlightService>::DoGetStream;
+
+async fn collect(response: tonic::Response<DoGetStream>) -> Vec<RecordBatch> {
+ let stream = response
+ .into_inner()
+ .map_err(|e| FlightError::Tonic(Box::new(e)));
+ FlightRecordBatchStream::new_from_flight_data(stream)
+ .try_collect()
+ .await
+ .expect("results decode")
+}
+
+async fn register_table(backend: &StubBackend, session_id: &str) {
+ let ctx = backend.session(session_id).await.unwrap();
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, true),
+ Field::new("name", DataType::Utf8, true),
+ ]));
+ ctx.register_table("people", Arc::new(EmptyTable::new(schema)))
+ .unwrap();
+}
+
+#[tokio::test]
+async fn select_produces_one_endpoint_per_partition_with_no_location() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query planned and submitted")
+ .into_inner();
+
+ assert_eq!(info.endpoint.len(), 1);
+ assert!(
+ info.endpoint[0].location.is_empty(),
+ "an endpoint with no location tells the client to reuse its existing \
+ connection; advertising the executor is what broke #1012"
+ );
+
+ // The ticket must carry the executor-facing fetch action, so the proxy can
+ // redeem it without any extra server-side lookup.
+ let handle = statement_handle(&info, 0);
+ let action = decode_protobuf(&handle[1..]).expect("handle wraps a Ballista action");
+ match action {
+ BallistaAction::FetchPartition {
+ job_id,
+ stage_id,
+ partition_id,
+ host,
+ port,
+ ..
+ } => {
+ assert_eq!(job_id.to_string(), JOB_ID);
+ assert_eq!(stage_id, 3);
+ assert_eq!(partition_id, 7);
+ assert_eq!(host, "executor-host");
+ assert_eq!(port, 50051);
+ }
+ }
+}
+
+#[tokio::test]
+async fn ddl_runs_on_the_scheduler_and_its_result_is_single_use() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE SCHEMA reporting".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("DDL executes")
+ .into_inner();
+
+ // DDL never reaches the cluster.
+ assert!(backend.executed.is_empty());
+
+ // The schema really was created in the session the client will query.
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ ctx.catalog("datafusion")
+ .unwrap()
+ .schema("reporting")
+ .is_some()
+ );
+
+ let handle = statement_handle(&info, 0);
+ assert_eq!(handle[0], 1, "DDL results use the local-result tag");
+
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.into(),
+ };
+ service
+ .do_get_statement(ticket.clone(), ticket_for(&info, 0))
+ .await
+ .expect("first fetch succeeds");
+
+ let err = expect_err(
+ service.do_get_statement(ticket, ticket_for(&info, 0)).await,
+ "a ticket is redeemable once",
+ );
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn writes_are_refused_rather_than_run_on_the_scheduler() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "INSERT INTO people VALUES (1, 'x')".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("the distributed write path is not implemented");
+
+ assert_eq!(err.code(), tonic::Code::Unimplemented);
+ assert!(backend.executed.is_empty());
+}
+
+/// `CREATE TABLE AS SELECT` arrives as DDL, and DDL runs on the scheduler. It
+/// must be refused rather than quietly executed on a single node.
+#[tokio::test]
+async fn ctas_is_refused_rather_than_run_on_one_node() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let err = expect_err(
+ service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE TABLE big AS SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await,
+ "CTAS would give up distribution",
+ );
+
+ assert_eq!(err.code(), tonic::Code::Unimplemented);
+
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ !ctx.table_exist("big").unwrap(),
+ "the refused statement must not have taken effect"
+ );
+}
+
+#[tokio::test]
+async fn catalog_metadata_reflects_the_session() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let command = CommandGetTables {
+ catalog: None,
+ db_schema_filter_pattern: None,
+ table_name_filter_pattern: None,
+ table_types: vec![],
+ include_schema: true,
+ };
+ let batches = collect(
+ service
+ .do_get_tables(command, Request::new(Ticket::default()))
+ .await
+ .expect("tables listed"),
+ )
+ .await;
+
+ let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(
+ rows, 1,
+ "the table registered in the session must be listed"
+ );
+
+ let batches = collect(
+ service
+ .do_get_catalogs(CommandGetCatalogs {}, Request::new(Ticket::default()))
+ .await
+ .expect("catalogs listed"),
+ )
+ .await;
+ let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(rows, 1);
+}
+
+#[tokio::test]
+async fn prepared_statements_are_planned_once_and_expire_on_close() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let prepared = service
+ .do_action_create_prepared_statement(
+ ActionCreatePreparedStatementRequest {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("statement prepared");
+
+ assert!(
+ !prepared.dataset_schema.is_empty(),
+ "clients need the result schema before executing"
+ );
+
+ let handle = prepared.prepared_statement_handle.clone();
+ service
+ .get_flight_info_prepared_statement(
+ arrow_flight::sql::CommandPreparedStatementQuery {
+ prepared_statement_handle: handle.clone(),
+ },
+ descriptor(),
+ )
+ .await
+ .expect("prepared statement executes");
+
+ service
+ .do_action_close_prepared_statement(
+ arrow_flight::sql::ActionClosePreparedStatementRequest {
+ prepared_statement_handle: handle.clone(),
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("closes");
+
+ let err = service
+ .get_flight_info_prepared_statement(
+ arrow_flight::sql::CommandPreparedStatementQuery {
+ prepared_statement_handle: handle,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("a closed handle must not linger");
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn an_authenticator_makes_tokenless_requests_fail() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend).with_authenticator(Arc::new(DenyAll));
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT 1".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect_err("no anonymous fallback when an authenticator is installed");
+
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+#[tokio::test]
+async fn the_default_service_allows_anonymous_access() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+ assert!(service.allows_anonymous());
+
+ let strict = make_service(Arc::new(StubBackend::default()))
+ .with_authenticator(Arc::new(DenyAll));
+ assert!(!strict.allows_anonymous());
+
+ // And the shipped default really is the permissive one, so the scheduler's
+ // startup warning is not dead code.
+ assert!(AnonymousAuthenticator.allows_anonymous());
+}
+
+#[tokio::test]
+async fn cancelling_a_query_reaches_the_backend() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query submitted")
+ .into_inner();
+
+ service
+ .do_action_cancel_query(
+ arrow_flight::sql::ActionCancelQueryRequest {
+ info: info.encode_to_vec().into(),
+ },
+ Request::new(Action::default()),
+ )
+ .await
+ .expect("cancellation accepted");
+
+ assert_eq!(
+ backend.cancelled.get(JOB_ID).map(|v| *v),
+ Some(1),
+ "cancel must be routed to the job the FlightInfo describes"
+ );
+}
+
+#[tokio::test]
+async fn unknown_tokens_are_rejected() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend);
+
+ let mut request = descriptor();
+ request
+ .metadata_mut()
+ .insert("authorization", "Bearer not-a-real-token".parse().unwrap());
+
+ let err = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT 1".to_string(),
+ transaction_id: None,
+ },
+ request,
+ )
+ .await
+ .expect_err("a stale token must not silently fall back to a shared session");
+
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+/// The frontend replaces the standalone Flight proxy when mounted, so
+/// Ballista's own client tickets must still be recognised. The fallback is
+/// only reachable if arrow-flight's dispatcher does not mistake a Ballista
+/// action for a Flight SQL command.
+#[tokio::test]
+async fn ballista_client_tickets_are_recognised_by_the_fallback() {
Review Comment:
Done. `ballista_client_tickets_reach_the_proxy` calls `FlightService::do_get` on the service, so the dispatcher runs, and `ballista_tickets_route_to_the_fallback` is gone.
##########
ballista/flight-sql/tests/service.rs:
##########
@@ -0,0 +1,593 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Protocol-level tests for the Flight SQL frontend, driven through the
+//! `FlightSqlService` trait against a stub backend.
+//!
+//! These deliberately avoid a live cluster so they run anywhere; the
+//! scheduler's `tests/flight_sql.rs` covers the same surface against real
+//! executors. Everything here is asserted the way a client sees it — tickets
+//! are decoded from the wire bytes rather than through crate internals — so
+//! the tests fail if the wire format changes.
+
+use std::sync::Arc;
+
+use arrow_flight::decode::FlightRecordBatchStream;
+use arrow_flight::error::FlightError;
+use arrow_flight::flight_service_server::FlightService;
+use arrow_flight::sql::server::FlightSqlService;
+use arrow_flight::sql::{
+ ActionCreatePreparedStatementRequest, Any, CommandGetCatalogs, CommandGetTables,
+ CommandStatementQuery, TicketStatementQuery,
+};
+use arrow_flight::{Action, FlightDescriptor, FlightInfo, Ticket};
+use async_trait::async_trait;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::flight_proxy_service::BallistaFlightProxyService;
+use ballista_core::serde::decode_protobuf;
+use ballista_core::serde::protobuf::{ExecutorMetadata, PartitionId, PartitionLocation};
+use ballista_core::serde::scheduler::{
+ Action as BallistaAction, ShuffleFileKind, ShuffleLayout,
+};
+use ballista_flight_sql::backend::{QueryBackend, QueryResult};
+use ballista_flight_sql::{
+ ANONYMOUS_SESSION, AnonymousAuthenticator, Authenticator, BallistaFlightSqlService,
+ Identity,
+};
+use dashmap::DashMap;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::empty::EmptyTable;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::SessionContext;
+use futures::TryStreamExt;
+use prost::Message;
+use tonic::metadata::MetadataMap;
+use tonic::{Request, Status};
+
+/// A backend that never contacts a cluster: it hands out real
+/// `SessionContext`s so planning is genuine, and fabricates the partition
+/// locations a completed job would have produced.
+#[derive(Default)]
+struct StubBackend {
+ sessions: DashMap<String, Arc<SessionContext>>,
+ executed: DashMap<String, usize>,
+ cancelled: DashMap<String, usize>,
+}
+
+const JOB_ID: &str = "job-1";
+
+#[async_trait]
+impl QueryBackend for StubBackend {
+ async fn session(&self, session_id: &str) -> Result<Arc<SessionContext>> {
+ Ok(self
+ .sessions
+ .entry(session_id.to_string())
+ .or_insert_with(|| Arc::new(SessionContext::new()))
+ .clone())
+ }
+
+ async fn close_session(&self, session_id: &str) -> Result<()> {
+ self.sessions.remove(session_id);
+ Ok(())
+ }
+
+ async fn execute(
+ &self,
+ _job_name: &str,
+ _ctx: Arc<SessionContext>,
+ plan: LogicalPlan,
+ ) -> Result<QueryResult> {
+ *self.executed.entry(JOB_ID.to_string()).or_insert(0) += 1;
+
+ Ok(QueryResult {
+ job_id: JOB_ID.to_string(),
+ schema: Arc::new(plan.schema().as_arrow().clone()),
+ partitions: vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: Some(PartitionId {
+ job_id: JOB_ID.to_string(),
+ stage_id: 3,
+ partition_id: 7,
+ }),
+ executor_meta: Some(ExecutorMetadata {
+ id: "executor-1".to_string(),
+ host: "executor-host".to_string(),
+ port: 50051,
+ grpc_port: 50052,
+ specification: None,
+ os_info: None,
+ }),
+ partition_stats: None,
+ file_id: None,
+ is_sort_shuffle: false,
+ }],
+ })
+ }
+
+ async fn cancel(&self, job_id: &str) -> Result<()> {
+ *self.cancelled.entry(job_id.to_string()).or_insert(0) += 1;
+ Ok(())
+ }
+}
+
+/// Rejects everything, to check that the frontend actually consults the
+/// authenticator and refuses tokenless requests when one is installed.
+struct DenyAll;
+
+#[async_trait]
+impl Authenticator for DenyAll {
+ async fn authenticate(
+ &self,
+ _headers: &MetadataMap,
+ ) -> std::result::Result<Identity, Status> {
+ Err(Status::unauthenticated("nope"))
+ }
+}
+
+fn make_service(backend: Arc<StubBackend>) -> BallistaFlightSqlService<StubBackend> {
+ BallistaFlightSqlService::new(
+ backend,
+ BallistaFlightProxyService::new(4_194_304, 4_194_304, false, None),
+ )
+}
+
+fn descriptor() -> Request<FlightDescriptor> {
+ Request::new(FlightDescriptor::new_cmd(vec![]))
+}
+
+/// `expect_err` needs `Debug` on the success type, and Flight's stream
+/// responses do not have it.
+fn expect_err<T>(result: std::result::Result<T, Status>, msg: &str) -> Status {
+ match result {
+ Ok(_) => panic!("{msg}"),
+ Err(status) => status,
+ }
+}
+
+/// Pulls the ticket out of an endpoint exactly as a Flight client would.
+fn statement_handle(info: &FlightInfo, index: usize) -> Vec<u8> {
+ let ticket = info.endpoint[index].ticket.as_ref().expect("ticket");
+ let any = Any::decode(&*ticket.ticket).expect("ticket is an Any");
+ let statement: TicketStatementQuery =
+ any.unpack().expect("unpackable").expect("statement ticket");
+ statement.statement_handle.to_vec()
+}
+
+fn ticket_for(info: &FlightInfo, index: usize) -> Request<Ticket> {
+ Request::new(info.endpoint[index].ticket.clone().expect("ticket"))
+}
+
+type DoGetStream = <BallistaFlightSqlService<StubBackend> as FlightService>::DoGetStream;
+
+async fn collect(response: tonic::Response<DoGetStream>) -> Vec<RecordBatch> {
+ let stream = response
+ .into_inner()
+ .map_err(|e| FlightError::Tonic(Box::new(e)));
+ FlightRecordBatchStream::new_from_flight_data(stream)
+ .try_collect()
+ .await
+ .expect("results decode")
+}
+
+async fn register_table(backend: &StubBackend, session_id: &str) {
+ let ctx = backend.session(session_id).await.unwrap();
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, true),
+ Field::new("name", DataType::Utf8, true),
+ ]));
+ ctx.register_table("people", Arc::new(EmptyTable::new(schema)))
+ .unwrap();
+}
+
+#[tokio::test]
+async fn select_produces_one_endpoint_per_partition_with_no_location() {
+ let backend = Arc::new(StubBackend::default());
+ register_table(&backend, ANONYMOUS_SESSION).await;
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "SELECT id FROM people".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("query planned and submitted")
+ .into_inner();
+
+ assert_eq!(info.endpoint.len(), 1);
+ assert!(
+ info.endpoint[0].location.is_empty(),
+ "an endpoint with no location tells the client to reuse its existing \
+ connection; advertising the executor is what broke #1012"
+ );
+
+ // The ticket must carry the executor-facing fetch action, so the proxy can
+ // redeem it without any extra server-side lookup.
+ let handle = statement_handle(&info, 0);
+ let action = decode_protobuf(&handle[1..]).expect("handle wraps a Ballista action");
+ match action {
+ BallistaAction::FetchPartition {
+ job_id,
+ stage_id,
+ partition_id,
+ host,
+ port,
+ ..
+ } => {
+ assert_eq!(job_id.to_string(), JOB_ID);
+ assert_eq!(stage_id, 3);
+ assert_eq!(partition_id, 7);
+ assert_eq!(host, "executor-host");
+ assert_eq!(port, 50051);
+ }
+ }
+}
+
+#[tokio::test]
+async fn ddl_runs_on_the_scheduler_and_its_result_is_single_use() {
+ let backend = Arc::new(StubBackend::default());
+ let service = make_service(backend.clone());
+
+ let info = service
+ .get_flight_info_statement(
+ CommandStatementQuery {
+ query: "CREATE SCHEMA reporting".to_string(),
+ transaction_id: None,
+ },
+ descriptor(),
+ )
+ .await
+ .expect("DDL executes")
+ .into_inner();
+
+ // DDL never reaches the cluster.
+ assert!(backend.executed.is_empty());
+
+ // The schema really was created in the session the client will query.
+ let ctx = backend.session(ANONYMOUS_SESSION).await.unwrap();
+ assert!(
+ ctx.catalog("datafusion")
+ .unwrap()
+ .schema("reporting")
+ .is_some()
+ );
+
+ let handle = statement_handle(&info, 0);
+ assert_eq!(handle[0], 1, "DDL results use the local-result tag");
+
+ let ticket = TicketStatementQuery {
+ statement_handle: handle.into(),
+ };
+ service
+ .do_get_statement(ticket.clone(), ticket_for(&info, 0))
+ .await
+ .expect("first fetch succeeds");
+
+ let err = expect_err(
+ service.do_get_statement(ticket, ticket_for(&info, 0)).await,
+ "a ticket is redeemable once",
+ );
+ assert_eq!(err.code(), tonic::Code::NotFound);
+}
+
+#[tokio::test]
+async fn writes_are_refused_rather_than_run_on_the_scheduler() {
Review Comment:
Done. `statements_that_cannot_be_distributed_are_refused` loops over `INSERT`, CTAS, `COPY`, `EXECUTE` and `EXPLAIN ANALYZE INSERT`. It also checks that nothing reached the backend and that the CTAS table doesn't exist.
##########
ballista/scheduler/tests/flight_sql.rs:
##########
@@ -0,0 +1,339 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! End-to-end coverage for the Arrow Flight SQL frontend against a real
+//! (in-process) cluster: scheduler plus one executor, driven entirely through
+//! the Flight SQL protocol.
+//!
+//! The pre-46.0.0 implementation shipped with no tests at all, and the bugs
+//! that got it removed (#1012, #941, #839, #756) were exactly the ones an
+//! end-to-end test catches: endpoints pointing somewhere the client cannot
+//! reach, and results that fail to decode.
+
+#![cfg(feature = "flight-sql")]
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow_flight::sql::client::FlightSqlServiceClient;
+use arrow_flight::sql::{CommandGetTables, SqlInfo};
+use ballista_core::config::TaskSchedulingPolicy;
+use ballista_core::error::BallistaError;
+use ballista_core::serde::BallistaCodec;
+use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
+use ballista_core::utils::{
+ GrpcClientConfig, create_grpc_client_connection, default_config_producer,
+ default_session_builder,
+};
+use ballista_scheduler::cluster::BallistaCluster;
+use ballista_scheduler::config::SchedulerConfig;
+use ballista_scheduler::metrics::default_metrics_collector;
+use ballista_scheduler::scheduler_process::start_grpc_service_with_listener;
+use ballista_scheduler::scheduler_server::SchedulerServer;
+use datafusion::arrow::array::RecordBatch;
+use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+use futures::TryStreamExt;
+use tonic::transport::Channel;
+
+/// Test errors are only ever printed, and the client, transport, and Ballista
+/// error types all implement `Error` — so one boxed type spares every call site
+/// a `map_err`.
+type Result<T = ()> = std::result::Result<T, Box<dyn std::error::Error>>;
+
+/// Boots a scheduler serving Flight SQL plus one executor, and returns the
+/// scheduler's URL.
+async fn start_cluster() -> Result<String> {
+ let mut config = SchedulerConfig::default()
+ .with_scheduler_policy(TaskSchedulingPolicy::PullStaged);
+ config.flight_sql = true;
+
+ let cluster = BallistaCluster::new_memory(
+ "localhost:50050",
+ Arc::new(default_session_builder),
+ Arc::new(default_config_producer),
+ );
+
+ let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerServer::new(
+ "localhost:50050".to_owned(),
+ cluster,
+ BallistaCodec::default(),
+ Arc::new(config),
+ default_metrics_collector()?,
+ );
+ scheduler.init().await?;
+
+ // Bind first so the test learns the port without racing another process
+ // for it.
+ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+ let addr = listener.local_addr()?;
+ tokio::spawn(start_grpc_service_with_listener(listener, scheduler));
+
+ let url = format!("http://{addr}");
+
+ let scheduler_client = connect_with_retry(&url).await?;
+ ballista_executor::new_standalone_executor(
+ scheduler_client,
+ 2,
+ BallistaCodec::default(),
+ )
+ .await?;
+
+ Ok(url)
+}
+
+async fn connect_with_retry(url: &str) -> Result<SchedulerGrpcClient<Channel>> {
+ for _ in 0..100 {
+ if let Ok(client) = SchedulerGrpcClient::connect(url.to_string()).await {
+ return Ok(client);
+ }
+ tokio::time::sleep(Duration::from_millis(50)).await;
+ }
+ Err(BallistaError::General(format!("scheduler at {url} did not come up")).into())
+}
+
+/// Connects a Flight SQL client, authenticates, and registers a `people` table
+/// for it to query.
+///
+/// The `TempDir` is returned because it owns the CSV the table points at; drop
+/// it and the executors have nothing to read.
+async fn client_with_people(
+ url: &str,
+) -> Result<(FlightSqlServiceClient<Channel>, tempfile::TempDir)> {
+ let channel =
+ create_grpc_client_connection(url.to_string(), &GrpcClientConfig::default())
+ .await?;
+ let mut client = FlightSqlServiceClient::new(channel);
+ client.handshake("anonymous", "").await?;
+
+ let csv = tempfile::tempdir()?;
+ let path = csv.path().join("data.csv");
+ std::fs::write(&path, "id,name\n1,alice\n2,bob\n3,carol\n")?;
+ let path = path.to_string_lossy();
+
+ // DDL runs on the scheduler and registers into the session catalog.
+ let info = client
+ .execute(
+ format!(
+ "CREATE EXTERNAL TABLE people STORED AS CSV LOCATION '{path}' \
+ OPTIONS ('format.has_header' 'true')"
+ ),
+ None,
+ )
+ .await?;
+ collect(&mut client, info).await?;
+
+ Ok((client, csv))
+}
+
+/// Collects every endpoint of a `FlightInfo`, which is what a real client does
+/// and what proves the tickets we hand out are redeemable.
+async fn collect(
+ client: &mut FlightSqlServiceClient<Channel>,
+ info: arrow_flight::FlightInfo,
+) -> Result<Vec<RecordBatch>> {
+ let mut batches = Vec::new();
+ for endpoint in info.endpoint {
+ let ticket = endpoint
+ .ticket
+ .ok_or_else(|| BallistaError::General("endpoint has no ticket".into()))?;
+ let mut collected: Vec<RecordBatch> =
+ client.do_get(ticket).await?.try_collect().await?;
+ batches.append(&mut collected);
+ }
+ Ok(batches)
+}
+
+/// Total rows across every batch a query returned.
+fn row_count(batches: &[RecordBatch]) -> usize {
+ batches.iter().map(|b| b.num_rows()).sum()
+}
+
+/// The headline path: authenticate, register a table with DDL, run a query
+/// that the cluster actually distributes, and read the results back through
+/// the scheduler.
+#[tokio::test]
+async fn flight_sql_runs_a_distributed_query() -> Result {
+ let url = start_cluster().await?;
+ let (mut client, _csv) = client_with_people(&url).await?;
+
+ let info = client
+ .execute(
+ "SELECT name FROM people WHERE id > 1 ORDER BY name".to_string(),
+ None,
+ )
+ .await?;
+
+ // Every endpoint must be redeemable on the connection we already have:
+ // no executor address is allowed to leak to the client.
+ assert!(
+ info.endpoint.iter().all(|e| e.location.is_empty()),
+ "endpoints must not advertise cluster-internal addresses: {:?}",
+ info.endpoint
+ );
+
+ let batches = collect(&mut client, info).await?;
+ assert_eq!(
+ row_count(&batches),
+ 2,
+ "expected bob and carol, got {batches:?}"
+ );
+
+ Ok(())
+}
+
+/// A table registered through DDL must still be there for the next request.
+/// `QueryBackend::session` builds a fresh `SessionContext` every call, so this
+/// only holds because the frontend caches it per session.
+#[tokio::test]
+async fn a_session_keeps_its_catalog_across_requests() -> Result {
Review Comment:
Removed. The distributed query and prepared statement tests already cover it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: Arrow Flight SQL frontend for the scheduler + ADBC support [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #2416:
URL: https://github.com/apache/datafusion-ballista/pull/2416
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org