fix: forward remote MV job cancellation

rust/lancedb/src/job.rs

@@ -3,6 +3,8 @@
 
 //! Handles to operations a server may run asynchronously.
 
+use std::future::Future;
+use std::pin::Pin;
 use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 use arrow_array::RecordBatch;
@@ -250,6 +252,18 @@ where
     pub(crate) fn spawned(task: JoinHandle<Result<T>>) -> Self {
         Self::new_typed(Box::new(SpawnedJob::new(task)))
     }
+
+    pub(crate) fn spawned_with_cancellation<F, Fut>(task: JoinHandle<Result<T>>, cancel: F) -> Self
+    where
+        F: Fn() -> Fut + Send + Sync + 'static,
+        Fut: Future<Output = Result<()>> + Send + 'static,
+    {
+        let cancellation: Cancellation = Arc::new(move || Box::pin(cancel()));
+        Self::new_typed(Box::new(SpawnedJob::new_with_cancellation(
+            task,
+            cancellation,
+        )))
+    }
 }
 
 impl<T> Job<T>
@@ -501,10 +515,28 @@ impl Outcome {
 struct SpawnedJob {
     outcome: watch::Receiver<Option<Outcome>>,
     abort: AbortHandle,
+    cancellation: Option<Cancellation>,
 }
 
+type CancellationFuture = Pin<Box<dyn Future<Output = Result<()>> + Send>>;
+type Cancellation = Arc<dyn Fn() -> CancellationFuture + Send + Sync>;
+
 impl SpawnedJob {
     fn new<T>(task: JoinHandle<Result<T>>) -> Self
+    where
+        T: Serialize + Send + 'static,
+    {
+        Self::new_inner(task, None)
+    }
+
+    fn new_with_cancellation<T>(task: JoinHandle<Result<T>>, cancellation: Cancellation) -> Self
+    where
+        T: Serialize + Send + 'static,
+    {
+        Self::new_inner(task, Some(cancellation))
+    }
+
+    fn new_inner<T>(task: JoinHandle<Result<T>>, cancellation: Option<Cancellation>) -> Self
     where
         T: Serialize + Send + 'static,
     {
@@ -526,7 +558,11 @@ impl SpawnedJob {
             };
             let _ = tx.send(Some(outcome));
         });
-        Self { outcome, abort }
+        Self {
+            outcome,
+            abort,
+            cancellation,
+        }
     }
 }
 
@@ -556,6 +592,19 @@ impl JobHandle for SpawnedJob {
     }
 
     async fn cancel(&self) -> Result<()> {
+        if let Some(cancellation) = &self.cancellation {
+            let mut outcome = self.outcome.clone();
+            tokio::select! {
+                biased;
+                settled = outcome.wait_for(Option::is_some) => {
+                    settled.map_err(|_| Error::Runtime {
+                        message: "job outcome was dropped before cancellation completed".to_string(),
+                    })?;
+                    return Ok(());
+                }
+                result = cancellation() => result?,
+            }
+        }
         self.abort.abort();
         Ok(())
     }

rust/lancedb/src/remote/db.rs

@@ -880,12 +880,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
             quote_sql_identifier(&request.name),
             request.query
         );
-        Ok(Job::spawned(tokio::spawn(async move {
-            let query = client.submit(&statement, &namespace).await?;
-            let mut reader = query.reader().await?;
-            while reader.try_next().await?.is_some() {}
-            Ok(())
-        })))
+        Ok(client.submit_as_job(statement, namespace, || async { Ok(()) }))
     }
 
     async fn drop_materialized_view_async(

rust/lancedb/src/remote/sql.rs

@@ -3,6 +3,7 @@
 
 use std::collections::HashMap;
 use std::fs;
+use std::future::{Future, pending};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex as StdMutex, OnceLock};
 use std::time::{Duration, Instant};
@@ -20,12 +21,14 @@ use arrow_schema::{Schema, SchemaRef};
 use futures::TryStreamExt;
 use http::header::{HeaderMap, HeaderName, HeaderValue};
 use prost::Message;
-use tokio::sync::{Mutex, Notify, OnceCell, mpsc};
+use serde::{Serialize, de::DeserializeOwned};
+use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch};
 use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};
 use uuid::Uuid;
 
 use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream};
 use crate::error::{Error, Result};
+use crate::job::Job;
 use crate::remote::client::{ClientConfig, TlsConfig};
 use crate::remote::retry::ResolvedRetryConfig;
 use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus};
@@ -209,6 +212,42 @@ impl SqlClient {
         .await
     }
 
+    pub(super) fn submit_as_job<T, F, Fut>(
+        &self,
+        statement: String,
+        namespace: Vec<String>,
+        finish: F,
+    ) -> Job<T>
+    where
+        T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static,
+        F: FnOnce() -> Fut + Send + 'static,
+        Fut: Future<Output = Result<T>> + Send + 'static,
+    {
+        let client = self.clone();
+        let (query_tx, query_rx) = watch::channel(None::<Arc<Query>>);
+        let task = tokio::spawn(async move {
+            let query = Arc::new(client.submit(&statement, &namespace).await?);
+            query_tx.send_replace(Some(query.clone()));
+            let mut reader = query.reader().await?;
+            while reader.try_next().await?.is_some() {}
+            finish().await
+        });
+        Job::spawned_with_cancellation(task, move || {
+            let mut query_rx = query_rx.clone();
+            async move {
+                loop {
+                    let query = query_rx.borrow().clone();
+                    if let Some(query) = query {
+                        return query.cancel().await;
+                    }
+                    if query_rx.changed().await.is_err() {
+                        pending::<()>().await;
+                    }
+                }
+            }
+        })
+    }
+
     pub(super) async fn describe(&self, query_id: Uuid) -> Result<QueryDescription> {
         let query = self
             .queries

rust/lancedb/src/remote/sql_test.rs

@@ -21,7 +21,10 @@ use super::*;
 use crate::database::Database;
 use crate::remote::RemoteCatalogOptions;
 use crate::remote::client::HeaderProvider;
+use crate::remote::client::test_utils::client_with_handler;
 use crate::remote::db::RemoteDatabase;
+use crate::remote::table::RemoteTable;
+use crate::table::BaseTable;
 
 #[derive(Debug, Default)]
 struct DelayedHeaderProvider {
@@ -73,6 +76,7 @@ struct TestSqlService {
     first_continuation_count: Arc<AtomicUsize>,
     transient_poll_failures: Arc<AtomicUsize>,
     headers: Arc<std::sync::Mutex<Vec<CapturedHeaders>>>,
+    gate_refresh_release: Arc<Notify>,
     result: RecordBatch,
     large_result: RecordBatch,
     dictionary_result: RecordBatch,
@@ -121,6 +125,7 @@ impl Default for TestSqlService {
             first_continuation_count: Arc::new(AtomicUsize::new(0)),
             transient_poll_failures: Arc::new(AtomicUsize::new(0)),
             headers: Arc::new(std::sync::Mutex::new(Vec::new())),
+            gate_refresh_release: Arc::new(Notify::new()),
             result,
             large_result,
             dictionary_result,
@@ -279,6 +284,7 @@ impl FlightService for TestSqlService {
     ) -> std::result::Result<Response<<Self as FlightService>::DoGetStream>, Status> {
         self.do_get_count.fetch_add(1, Ordering::SeqCst);
         let ticket = request.get_ref().ticket.as_ref();
+        let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\"";
         let empty = ticket == b"SELECT empty";
         let slow = ticket == b"SELECT slow get";
         let large = ticket == b"SELECT large message";
@@ -290,7 +296,11 @@ impl FlightService for TestSqlService {
             self.result.clone()
         };
         let schema = result.schema();
+        let gate_refresh_release = self.gate_refresh_release.clone();
         let input = futures::stream::once(async move {
+            if gate_refresh {
+                gate_refresh_release.notified().await;
+            }
             if slow {
                 tokio::time::sleep(Duration::from_millis(250)).await;
             }
@@ -448,10 +458,6 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() {
 
 #[tokio::test]
 async fn remote_refresh_twice_keeps_latest_selector() {
-    use crate::remote::client::test_utils::client_with_handler;
-    use crate::remote::table::RemoteTable;
-    use crate::table::BaseTable;
-
     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
     let address = listener.local_addr().unwrap();
     let incoming = futures::stream::try_unfold(listener, |listener| async {
@@ -531,6 +537,66 @@ async fn remote_refresh_twice_keeps_latest_selector() {
     server.await.unwrap().unwrap();
 }
 
+#[tokio::test]
+async fn cancelling_remote_refresh_cancels_sql_query() {
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let address = listener.local_addr().unwrap();
+    let incoming = futures::stream::try_unfold(listener, |listener| async {
+        let (socket, _) = listener.accept().await?;
+        Ok::<_, std::io::Error>(Some((socket, listener)))
+    });
+    let service = TestSqlService::default();
+    let cancel_count = service.cancel_count.clone();
+    let do_get_count = service.do_get_count.clone();
+    let gate_refresh_release = service.gate_refresh_release.clone();
+    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
+    let server = tokio::spawn(
+        tonic::transport::Server::builder()
+            .add_service(FlightServiceServer::new(service))
+            .serve_with_incoming_shutdown(incoming, async {
+                let _ = shutdown_rx.await;
+            }),
+    );
+    let sql_client = SqlClient::new(
+        "analytics".into(),
+        None,
+        "test-key".into(),
+        None,
+        Some(format!("grpc://{address}")),
+        ClientConfig::default(),
+    );
+    let client = client_with_handler(|request| -> http::Response<String> {
+        panic!(
+            "cancelled refresh must not read result metadata: {}",
+            request.url().path()
+        );
+    });
+    let table = RemoteTable::new_with_sql_client(
+        client,
+        "cancel_view".into(),
+        vec![],
+        "cancel_view".into(),
+        Default::default(),
+        Some(sql_client),
+    );
+    let job = table.refresh_materialized_view_async(None).await.unwrap();
+    tokio::time::timeout(Duration::from_secs(10), async {
+        while do_get_count.load(Ordering::SeqCst) == 0 {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .unwrap();
+
+    job.cancel().await.unwrap();
+    assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. })));
+    assert_eq!(cancel_count.load(Ordering::SeqCst), 1);
+
+    gate_refresh_release.notify_one();
+    shutdown_tx.send(()).unwrap();
+    server.await.unwrap().unwrap();
+}
+
 #[tokio::test]
 async fn submits_polls_fetches_cancels_and_reuses_client() {
     let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();

rust/lancedb/src/remote/table.rs

@@ -2157,33 +2157,32 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
         let server_version = self.server_version.clone();
         let freshness = self.freshness.clone();
         let freshness_request = self.snapshot_freshness_headers();
-        Ok(Job::spawned(tokio::spawn(async move {
-            let query = sql_client.submit(&statement, &namespace).await?;
-            let mut reader = query.reader().await?;
-            while reader.try_next().await?.is_some() {}
-
-            let table = Self::new_with_sql_client(
-                client,
-                name,
-                namespace,
-                identifier,
-                server_version,
-                Some(sql_client),
-            );
-            table.checkout_latest().await?;
-            let version = table.version().await?;
-            let rows_written =
-                u64::try_from(table.count_rows(None).await?).map_err(|_| Error::Runtime {
-                    message: "materialized-view row count exceeds u64".to_string(),
-                })?;
-            freshness_request.observe_version(&freshness, version);
-            Ok(RefreshMaterializedViewResult {
-                mode: crate::materialized_view::RefreshMode::Rebuild,
-                rows_written,
-                source_version: 0,
-                version,
-            })
-        })))
+        let job_sql_client = sql_client.clone();
+        Ok(
+            sql_client.submit_as_job(statement, namespace.clone(), move || async move {
+                let table = Self::new_with_sql_client(
+                    client,
+                    name,
+                    namespace,
+                    identifier,
+                    server_version,
+                    Some(job_sql_client),
+                );
+                table.checkout_latest().await?;
+                let version = table.version().await?;
+                let rows_written =
+                    u64::try_from(table.count_rows(None).await?).map_err(|_| Error::Runtime {
+                        message: "materialized-view row count exceeds u64".to_string(),
+                    })?;
+                freshness_request.observe_version(&freshness, version);
+                Ok(RefreshMaterializedViewResult {
+                    mode: crate::materialized_view::RefreshMode::Rebuild,
+                    rows_written,
+                    source_version: 0,
+                    version,
+                })
+            }),
+        )
     }
     async fn query_snapshot(&self) -> Result<Arc<dyn BaseTable>> {
         let description = self.describe().await?;