Skip to content

Commit 34867f7

Browse files
committed
[adapters] Revert read-after-write semantics in ad hoc queries.
A recent change that enabled ad hoc queries to observe their own changes to input tables (but not views) introduced a regression that prevented bootstrapping state from materialized tables. We rollback that change here. Signed-off-by: Leonid Ryzhyk <ryzhyk@gmail.com>
1 parent 095a91e commit 34867f7

6 files changed

Lines changed: 78 additions & 213 deletions

File tree

crates/adapters/src/adhoc.rs

Lines changed: 6 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -292,25 +292,18 @@ pub(crate) async fn execute_sql(
292292
.await
293293
.ok_or_else(|| PipelineError::Initializing)?,
294294
);
295-
execute_sql_with_state(state, sql, Some(controller)).await
295+
execute_sql_with_state(state, sql).await
296296
}
297297

298298
/// Plan and translate `sql` against `state`, applying `PREPARE`/`EXECUTE`
299299
/// substitution within the scope of a single ad-hoc request.
300300
///
301301
/// Only the final statement returns rows. Earlier statements may be
302302
/// `PREPARE`s or any non-result-producing statement (e.g. `INSERT`),
303-
/// executed for their side effect. After such an intermediate write
304-
/// runs, the per-request snapshot is refreshed so the trailing SELECT
305-
/// sees the just-written rows.
306-
///
307-
/// `controller` is optional so unit tests can drive this function
308-
/// without a running pipeline; in that case the post-write snapshot
309-
/// refresh is skipped.
303+
/// executed for their side effect.
310304
async fn execute_sql_with_state(
311-
mut state: SessionState,
305+
state: SessionState,
312306
sql: &str,
313-
controller: Option<&Controller>,
314307
) -> Result<DataFrame, PipelineError> {
315308
let mut statements = parse_sql_statements(&state, sql)?;
316309
if statements.is_empty() {
@@ -326,13 +319,6 @@ async fn execute_sql_with_state(
326319
let mut prepared: HashMap<String, LogicalPlan> = HashMap::new();
327320
let sql_options = SQLOptions::new().with_allow_ddl(false);
328321

329-
// Subscribe to `step_watcher` *before* any intermediate writes so the
330-
// steps that drain those writes accumulate as unseen changes; calling
331-
// `changed()` after the writes return then completes immediately
332-
// instead of blocking on a step that may never be triggered.
333-
let mut step_watcher = controller.map(|c| c.step_watcher());
334-
335-
let mut intermediate_wrote_data = false;
336322
while statements.len() > 1 {
337323
let stmt = statements.pop_front().unwrap();
338324
let plan = state.statement_to_plan(stmt).await?;
@@ -356,7 +342,6 @@ async fn execute_sql_with_state(
356342
let values = execute_parameters_to_scalars(&parameters)?;
357343
let bound = prepared_plan.replace_params_with_values(&ParamValues::List(values))?;
358344
sql_options.verify_plan(&bound)?;
359-
intermediate_wrote_data |= plan_writes_data(&bound);
360345
drain_intermediate_plan(&state, bound).await?;
361346
}
362347
other if is_result_producing_plan(&other) => {
@@ -374,22 +359,11 @@ async fn execute_sql_with_state(
374359
// UPDATE, DELETE, EXPLAIN, ...). Execute it for its side
375360
// effects and discard the per-statement count row.
376361
sql_options.verify_plan(&other)?;
377-
intermediate_wrote_data |= plan_writes_data(&other);
378362
drain_intermediate_plan(&state, other).await?;
379363
}
380364
}
381365
}
382366

383-
// The snapshot pinned in `state` was captured at request start, before
384-
// any intermediate INSERT ran. Refresh it so the trailing SELECT sees
385-
// the just-written rows.
386-
if intermediate_wrote_data
387-
&& let Some(controller) = controller
388-
&& let Some(watcher) = step_watcher.as_mut()
389-
{
390-
refresh_snapshot_after_writes(controller, watcher, &mut state).await?;
391-
}
392-
393367
let stmt = statements.pop_front().unwrap();
394368
let plan = state.statement_to_plan(stmt).await?;
395369

@@ -447,43 +421,6 @@ async fn drain_intermediate_plan(
447421
Ok(())
448422
}
449423

450-
/// True if executing this plan mutates a table (and therefore needs the
451-
/// post-write snapshot refresh below). Today the only mutating plan our
452-
/// SQL options allow is a `LogicalPlan::Dml`.
453-
fn plan_writes_data(plan: &LogicalPlan) -> bool {
454-
matches!(plan, LogicalPlan::Dml(_))
455-
}
456-
457-
/// Wait for the controller to drain the intermediate writes, then pin the
458-
/// resulting snapshot in `state`. Relies on two controller invariants:
459-
/// `update_snapshot()` runs before each `Idle` notification (so any `Idle`
460-
/// observed after subscribing already reflects our writes), and `watcher`
461-
/// must be subscribed before the writes are submitted or we may sleep
462-
/// forever on an otherwise idle pipeline.
463-
async fn refresh_snapshot_after_writes(
464-
controller: &Controller,
465-
watcher: &mut tokio::sync::watch::Receiver<feldera_types::coordination::StepStatus>,
466-
state: &mut SessionState,
467-
) -> Result<(), PipelineError> {
468-
use feldera_types::coordination::StepAction;
469-
470-
loop {
471-
let status = *watcher.borrow_and_update();
472-
if matches!(status.action, StepAction::Idle) {
473-
break;
474-
}
475-
if watcher.changed().await.is_err() {
476-
break;
477-
}
478-
}
479-
let snapshot = controller
480-
.latest_consistent_snapshot()
481-
.await
482-
.ok_or(PipelineError::Initializing)?;
483-
set_snapshot(state, snapshot);
484-
Ok(())
485-
}
486-
487424
/// Convert `EXECUTE` positional parameters to DataFusion's `ScalarAndMetadata`
488425
/// list, rejecting anything that is not a literal value.
489426
fn execute_parameters_to_scalars(params: &[Expr]) -> Result<Vec<ScalarAndMetadata>, PipelineError> {
@@ -682,7 +619,7 @@ mod tests {
682619

683620
/// A helper that executes a query and returns results.
684621
async fn collect_rows(state: SessionState, sql: &str) -> Vec<RecordBatch> {
685-
execute_sql_with_state(state, sql, None)
622+
execute_sql_with_state(state, sql)
686623
.await
687624
.unwrap()
688625
.collect()
@@ -707,7 +644,7 @@ mod tests {
707644
#[tokio::test]
708645
async fn execute_without_prepare_errors() {
709646
let state = test_state();
710-
let err = execute_sql_with_state(state, "EXECUTE missing(1)", None)
647+
let err = execute_sql_with_state(state, "EXECUTE missing(1)")
711648
.await
712649
.unwrap_err();
713650
assert!(
@@ -730,7 +667,7 @@ mod tests {
730667
#[tokio::test]
731668
async fn intermediate_select_is_rejected() {
732669
let state = test_state();
733-
let err = execute_sql_with_state(state, "SELECT 1; SELECT 2", None)
670+
let err = execute_sql_with_state(state, "SELECT 1; SELECT 2")
734671
.await
735672
.unwrap_err();
736673
let msg = format!("{err}");

crates/adapters/src/controller.rs

Lines changed: 23 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -3023,28 +3023,25 @@ impl CircuitThread {
30233023

30243024
let transaction_state = self.controller.get_transaction_state();
30253025

3026-
// Update `trace_snapshot` to the latest traces *before* signaling
3027-
// `Idle` on the step watcher, so subscribers can rely on
3028-
// `latest_consistent_snapshot()` being current the moment they
3029-
// observe Idle.
3030-
//
3031-
// We also keep this before updating `total_processed_records` so
3032-
// that ad hoc query results always reflect all data that we have
3033-
// reported processing.
3034-
// Keep `trace_snapshots` current on every step regardless of
3035-
// bootstrap or transaction state, so ad-hoc queries can read
3036-
// the freshest data including read-your-own-writes within one
3037-
// request. Connectors with `send_snapshot=true` still defer
3038-
// until bootstrap completes; that gate now lives in
3039-
// `enqueue_latest_snapshot`.
3040-
Span::new("update")
3041-
.with_category("Step")
3042-
.with_tooltip(|| format!("update ad-hoc tables after step {}", self.step))
3043-
.in_scope(|| self.update_snapshot());
3044-
30453026
if transaction_state == TransactionState::None {
30463027
let bootstrapping = self.circuit.bootstrap_in_progress();
30473028

3029+
// Don't update the snapshot until bootstrapping is complete (including the additional post-bootstrap transaction).
3030+
// This guarantees that:
3031+
// 1. Ad hoc queries observe a consistent view of the data.
3032+
// 2. Ad hoc snapshots are up-to-date before the pipeline is marked as running.
3033+
if !bootstrapping && !self.bootstrapping {
3034+
// Update `trace_snapshot` to the latest traces.
3035+
//
3036+
// We do this before updating `total_processed_records` so
3037+
// that ad hoc query results always reflect all data that we have
3038+
// reported processing.
3039+
Span::new("update")
3040+
.with_category("Step")
3041+
.with_tooltip(|| format!("update ad-hoc tables after step {}", self.step))
3042+
.in_scope(|| self.update_snapshot());
3043+
}
3044+
30483045
// If bootstrapping has completed, clear self.bootstrapping, but don't update the status flag
30493046
// until the circuit performs an extra transaction to initialize output snapshots
30503047
// (`StepTrigger::trigger` makes sure to force a step as long as `controller.status.bootstrap_in_progress()`
@@ -3225,37 +3222,18 @@ impl CircuitThread {
32253222

32263223
/// Update `trace_snapshot` with the latest traces so ad hoc queries pick
32273224
/// up the freshest data.
3228-
///
3229-
/// Tables refresh every step so an INSERT inside a transaction is visible
3230-
/// to a subsequent SELECT in the same request. Views compute their
3231-
/// integral via `accumulate_integrate_trace`, which only advances on a
3232-
/// transaction boundary. We use the the previous snapshot instead.
32333225
fn update_snapshot(&mut self) {
3234-
let bootstrapping = self.controller.status.bootstrap_in_progress();
3235-
let in_transaction =
3236-
!bootstrapping && self.controller.get_transaction_state() != TransactionState::None;
3237-
3238-
let previous_views: Option<ConsistentSnapshot> = if in_transaction {
3239-
let snapshots = self.controller.trace_snapshots.blocking_lock();
3240-
snapshots.iter().next_back().map(|(_, snap)| snap.clone())
3241-
} else {
3242-
None
3243-
};
3226+
assert_eq!(
3227+
self.controller.get_transaction_state(),
3228+
TransactionState::None,
3229+
"update_snapshot must be called when no transaction is in progress"
3230+
);
32443231

32453232
// Assemble a new snapshot.
32463233
let mut snapshot = BTreeMap::new();
32473234
for (name, clh) in self.controller.catalog.output_iter() {
32483235
if let Some(ih) = &clh.integrate_handle {
3249-
let is_table = self
3250-
.controller
3251-
.catalog
3252-
.input_collection_handle(name)
3253-
.is_some();
3254-
3255-
let batches = match previous_views.as_ref().and_then(|s| s.get(name)) {
3256-
Some(prev) if !is_table && in_transaction => prev.clone(),
3257-
_ => ih.take_from_all(),
3258-
};
3236+
let batches = ih.take_from_all();
32593237

32603238
// The first index of a materialized view registers its
32613239
// integral under the view name with `alias_as_index =
@@ -6746,17 +6724,6 @@ impl ControllerInner {
67466724
processed_records: Option<ProcessedRecords>,
67476725
step: Option<Step>,
67486726
) -> bool {
6749-
// Defer the initial snapshot delivery until bootstrap completes
6750-
// and no transaction is in flight; this is the gate that
6751-
// guarantees `send_snapshot=true` connectors receive a
6752-
// consistent, fully-formed view. Adhoc queries hit
6753-
// `trace_snapshots` directly and are not subject to this gate.
6754-
if self.status.bootstrap_in_progress()
6755-
|| self.get_transaction_state() != TransactionState::None
6756-
{
6757-
return false;
6758-
}
6759-
67606727
// Look up the most recent cached snapshot for this stream. Return early
67616728
// if no snapshot has been produced yet (pipeline hasn't completed a step).
67626729
// This method is only called from the circuit thread (via `push_output`),
@@ -6769,7 +6736,7 @@ impl ControllerInner {
67696736
};
67706737

67716738
let Some(snapshot_batches) = snapshot_batches else {
6772-
// No cached snapshot yet (pipeline hasn't completed a step).
6739+
// No cached snapshot yet (pipeline hasn't completed a step or bootstrapping is in progress).
67736740
// Leave the state as `Pending` and let the next step retry.
67746741
return false;
67756742
};

crates/adapters/src/controller/test.rs

Lines changed: 15 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -2276,11 +2276,6 @@ fn lir() {
22762276
"stream_id": 3,
22772277
"to": "4"
22782278
},
2279-
{
2280-
"from": "3",
2281-
"stream_id": 3,
2282-
"to": "7"
2283-
},
22842279
{
22852280
"from": "4",
22862281
"stream_id": 4,
@@ -2289,12 +2284,12 @@ fn lir() {
22892284
{
22902285
"from": "4",
22912286
"stream_id": 4,
2292-
"to": "11"
2287+
"to": "7"
22932288
},
22942289
{
22952290
"from": "4",
22962291
"stream_id": 4,
2297-
"to": "13"
2292+
"to": "11"
22982293
},
22992294
{
23002295
"from": "6",
@@ -2316,6 +2311,11 @@ fn lir() {
23162311
"stream_id": 8,
23172312
"to": "9"
23182313
},
2314+
{
2315+
"from": "7",
2316+
"stream_id": 8,
2317+
"to": "12"
2318+
},
23192319
{
23202320
"from": "9",
23212321
"stream_id": 9,
@@ -2325,26 +2325,6 @@ fn lir() {
23252325
"from": "12",
23262326
"stream_id": 10,
23272327
"to": "13"
2328-
},
2329-
{
2330-
"from": "12",
2331-
"stream_id": null,
2332-
"to": "14"
2333-
},
2334-
{
2335-
"from": "13",
2336-
"stream_id": 13,
2337-
"to": "14"
2338-
},
2339-
{
2340-
"from": "13",
2341-
"stream_id": 13,
2342-
"to": "15"
2343-
},
2344-
{
2345-
"from": "15",
2346-
"stream_id": 14,
2347-
"to": "16"
23482328
}
23492329
],
23502330
"nodes": [
@@ -2394,21 +2374,24 @@ fn lir() {
23942374
{
23952375
"id": "6",
23962376
"implements": [
2397-
"input.output"
2377+
"input.output",
2378+
"output"
23982379
],
23992380
"operation": "Z1 (trace)"
24002381
},
24012382
{
24022383
"id": "7",
24032384
"implements": [
2404-
"input.output"
2385+
"input.output",
2386+
"output"
24052387
],
2406-
"operation": "UntimedTraceAppend"
2388+
"operation": "AccumulateUntimedTraceAppend"
24072389
},
24082390
{
24092391
"id": "8",
24102392
"implements": [
2411-
"input.output"
2393+
"input.output",
2394+
"output"
24122395
],
24132396
"operation": "Z1 (trace)"
24142397
},
@@ -2438,31 +2421,10 @@ fn lir() {
24382421
"implements": [
24392422
"output"
24402423
],
2441-
"operation": "Z1 (trace)"
2442-
},
2443-
{
2444-
"id": "13",
2445-
"implements": [
2446-
"output"
2447-
],
2448-
"operation": "AccumulateUntimedTraceAppend"
2449-
},
2450-
{
2451-
"id": "14",
2452-
"implements": [
2453-
"output"
2454-
],
2455-
"operation": "Z1 (trace)"
2456-
},
2457-
{
2458-
"id": "15",
2459-
"implements": [
2460-
"output"
2461-
],
24622424
"operation": "Apply"
24632425
},
24642426
{
2465-
"id": "16",
2427+
"id": "13",
24662428
"implements": [
24672429
"output"
24682430
],

0 commit comments

Comments
 (0)