adhoc: allow non-PREPARE intermediate statements in multi-statement queries

gz · gz · commit 19c5e73617c0 · 2026-05-15T10:41:59.000-07:00
Previously `execute_sql_with_state` rejected any non-PREPARE statement before the final statement, so submitting `INSERT; INSERT; SELECT` or even two INSERTs in a row required splitting them into separate requests. Relax that restriction: an intermediate non-result-producing statement (DML or a DataFusion `Statement` variant such as a prepared EXECUTE) is now executed for its side effect, its row count is discarded, and the final statement still owns the streaming protocol's single result slot. A result-producing statement (SELECT and other relational operators) in a non-final position is still rejected. Will need to add come up with a streaming protocol to support this. Fixes #2487
diff --git a/crates/adapters/src/adhoc.rs b/crates/adapters/src/adhoc.rs
@@ -297,6 +297,10 @@ pub(crate) async fn execute_sql(
 
 /// Plan and translate `sql` against `state`, applying `PREPARE`/`EXECUTE`
 /// substitution within the scope of a single ad-hoc request.
+///
+/// Only the final statement returns rows. Earlier statements may be
+/// `PREPARE`s or any non-result-producing statement (e.g. `INSERT`),
+/// executed for their side effect.
 async fn execute_sql_with_state(
     state: SessionState,
     sql: &str,
@@ -315,9 +319,6 @@ async fn execute_sql_with_state(
     let mut prepared: HashMap<String, LogicalPlan> = HashMap::new();
     let sql_options = SQLOptions::new().with_allow_ddl(false);
 
-    // For now, only the final statement may produce a result set. All
-    // preceding statements must be PREPAREs whose inner plans are stashed
-    // for a later EXECUTE in the same request.
     while statements.len() > 1 {
         let stmt = statements.pop_front().unwrap();
         let plan = state.statement_to_plan(stmt).await?;
@@ -326,14 +327,40 @@ async fn execute_sql_with_state(
                 sql_options.verify_plan(&input)?;
                 prepared.insert(name, (*input).clone());
             }
-            _ => {
+            LogicalPlan::Statement(Statement::Execute(Execute { name, parameters })) => {
+                // `EXECUTE` of a previously-prepared statement, used here
+                // for its side effects (e.g. a prepared INSERT).
+                let prepared_plan =
+                    prepared
+                        .remove(&name)
+                        .ok_or_else(|| PipelineError::AdHocQueryError {
+                            error: format!(
+                                "prepared statement '{name}' is not defined in this request"
+                            ),
+                            df: None,
+                        })?;
+                let values = execute_parameters_to_scalars(&parameters)?;
+                let bound = prepared_plan.replace_params_with_values(&ParamValues::List(values))?;
+                sql_options.verify_plan(&bound)?;
+                drain_intermediate_plan(&state, bound).await?;
+            }
+            other if is_result_producing_plan(&other) => {
                 return Err(PipelineError::AdHocQueryError {
-                    error: "only PREPARE statements may precede the final statement \
-                            in a multi-statement ad-hoc query"
+                    error: "only the final statement in a multi-statement \
+                            ad-hoc query may return a result set; \
+                            move SELECTs to the end or split into \
+                            separate requests"
                         .to_string(),
                     df: None,
                 });
             }
+            other => {
+                // Non-result-producing intermediate statement (INSERT,
+                // UPDATE, DELETE, EXPLAIN, ...). Execute it for its side
+                // effects and discard the per-statement count row.
+                sql_options.verify_plan(&other)?;
+                drain_intermediate_plan(&state, other).await?;
+            }
         }
     }
 
@@ -374,6 +401,26 @@ async fn execute_sql_with_state(
     Ok(DataFrame::new(state, final_plan))
 }
 
+/// True if executing this plan would surface rows to the caller. Used to
+/// reject queries like `SELECT; INSERT` where the early `SELECT` would
+/// otherwise be silently dropped.
+fn is_result_producing_plan(plan: &LogicalPlan) -> bool {
+    !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Statement(_))
+}
+
+/// Execute an intermediate statement for its side effects and drop the
+/// resulting batches. INSERTs produce a one-row count; we keep that
+/// count out of the response stream so only the request's final
+/// statement contributes rows.
+async fn drain_intermediate_plan(
+    state: &SessionState,
+    plan: LogicalPlan,
+) -> Result<(), PipelineError> {
+    let df = DataFrame::new(state.clone(), plan);
+    let _ = df.collect().await?;
+    Ok(())
+}
+
 /// Convert `EXECUTE` positional parameters to DataFusion's `ScalarAndMetadata`
 /// list, rejecting anything that is not a literal value.
 fn execute_parameters_to_scalars(params: &[Expr]) -> Result<Vec<ScalarAndMetadata>, PipelineError> {
@@ -614,15 +661,74 @@ mod tests {
         assert_eq!(total_rows, 0);
     }
 
+    /// An intermediate `SELECT` (or any other result-producing statement)
+    /// must be rejected: only one result set comes back per request, so
+    /// executing the earlier SELECT silently would discard its rows.
     #[tokio::test]
-    async fn non_prepare_intermediate_statement_errors() {
+    async fn intermediate_select_is_rejected() {
         let state = test_state();
         let err = execute_sql_with_state(state, "SELECT 1; SELECT 2")
             .await
             .unwrap_err();
-        assert!(
-            format!("{err:?}").contains("PREPARE"),
-            "unexpected error: {err:?}"
-        );
+        let msg = format!("{err}");
+        assert!(msg.contains("final statement"), "unexpected error: {msg}");
+    }
+
+    /// Multiple `INSERT`s followed by a `SELECT` must execute in order,
+    /// committing each insert's side effect, and only surface the final
+    /// `SELECT`'s rows.
+    #[tokio::test]
+    async fn intermediate_inserts_run_and_final_select_returns_rows() {
+        use datafusion::arrow::array::Int64Array;
+        use datafusion::arrow::datatypes::{DataType, Field, Schema};
+        use datafusion::datasource::MemTable;
+        use std::sync::Arc;
+
+        // Register a writable in-memory table so DML executes for real.
+        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
+        let mem = MemTable::try_new(schema.clone(), vec![vec![]]).unwrap();
+        let ctx = SessionContext::new_with_state(test_state());
+        ctx.register_table("t", Arc::new(mem)).unwrap();
+        let state = ctx.state();
+
+        let batches = collect_rows(
+            state,
+            "INSERT INTO t VALUES (1); INSERT INTO t VALUES (2); \
+             SELECT SUM(x) AS s FROM t",
+        )
+        .await;
+        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+        assert_eq!(total_rows, 1);
+        let col = batches[0]
+            .column(0)
+            .as_any()
+            .downcast_ref::<Int64Array>()
+            .expect("int64 column");
+        assert_eq!(col.value(0), 3);
+    }
+
+    /// A trailing `INSERT` (no final SELECT) must still execute, and
+    /// the final statement's count row is surfaced as today.
+    #[tokio::test]
+    async fn final_insert_returns_count() {
+        use datafusion::arrow::datatypes::{DataType, Field, Schema};
+        use datafusion::datasource::MemTable;
+        use std::sync::Arc;
+
+        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)]));
+        let mem = MemTable::try_new(schema.clone(), vec![vec![]]).unwrap();
+        let ctx = SessionContext::new_with_state(test_state());
+        ctx.register_table("t", Arc::new(mem)).unwrap();
+        let state = ctx.state();
+
+        let batches = collect_rows(
+            state,
+            "INSERT INTO t VALUES (10); INSERT INTO t VALUES (20)",
+        )
+        .await;
+        // The final INSERT yields a single-row count batch; check only
+        // that one row came back.
+        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+        assert_eq!(total_rows, 1);
     }
 }
diff --git a/crates/fda/test.bash b/crates/fda/test.bash
@@ -131,6 +131,11 @@ fail_on_success fda query p1 "SELECT 1/0"
 fail_on_success fda --format arrow_ipc query p1 "SELECT 1/0"
 fail_on_success fda --format json query p1 "SELECT 1/0"
 
+# Intermediate `SELECT`s still error today; we plan to support
+# `select; select; ...` once a streaming protocol can frame multiple
+# result sets back to the caller.
+fail_on_success fda query p1 "SELECT 1; SELECT 2"
+
 # Transaction tests
 echo "Testing transaction commands..."
 fail_on_success fda commit-transaction p1
diff --git a/python/tests/runtime/test_adhoc_queries.py b/python/tests/runtime/test_adhoc_queries.py
@@ -235,6 +235,20 @@ def norm(rows):
         assert prepared_rows and prepared_rows[0].get("c") == 1
         assert self._count("SELECT COUNT(*) AS c FROM t1") == 7
 
+        # Multiple non-PREPARE statements in a single request: earlier
+        # INSERTs run for their side effect; the trailing SELECT returns
+        # its rows.
+        chain_rows = list(
+            self.pipeline.query(
+                "INSERT INTO t1 VALUES "
+                "(200,'2020-02-02','11111111-1111-1111-1111-111111111111');"
+                "INSERT INTO t1 VALUES "
+                "(201,'2020-02-02','22222222-2222-2222-2222-222222222222');"
+                "SELECT COUNT(*) AS c FROM t1 WHERE id BETWEEN 200 AND 201"
+            )
+        )
+        assert chain_rows and chain_rows[0].get("c") == 2
+
         # Non-materialized table via its materialized view
         assert self._count("SELECT COUNT(*) AS c FROM view_of_not_materialized") == 0
         ins_nm = list(
