feat(sf_core): atomically take result streams

sf_core/src/apis/database_driver_v1/result_set.rs

@@ -390,26 +390,11 @@ async fn snapshot_reader_inputs(
 // --- DatabaseDriverV1 impl ---
 
 impl DatabaseDriverV1 {
-    /// Builds a fresh Arrow [`RecordBatchReader`] for this result set, lazily
-    /// from the stored `RowsetData`, so it can be requested multiple times. The
-    /// protobuf layer wraps it in an `FFI_ArrowArrayStream` at the C boundary.
-    ///
-    /// Awaiting this only builds the reader and never blocks. Iterating it,
-    /// however, must happen in a synchronous context: chunked result sets pull
-    /// chunks via a blocking channel receiver, so draining from within an async
-    /// runtime would call `blocking_recv` and panic. Drain after returning from
-    /// `block_on` (keeping the runtime alive), on a dedicated `std::thread`, or
-    /// via `tokio::task::spawn_blocking`.
-    pub async fn result_set_get_stream(
+    // Build a reader from a retained or atomically consumed result set.
+    async fn build_result_set_stream(
         &self,
-        result_handle: Handle,
+        rs_ptr: Arc<Mutex<ResultSet>>,
     ) -> Result<Box<dyn RecordBatchReader + Send>, ApiError> {
-        let rs_ptr = self
-            .results
-            .get_obj(result_handle)
-            .with_context(|| InvalidArgumentSnafu {
-                argument: "result_handle: ResultSet handle not found".to_string(),
-            })?;
         let (data, http_client, prefetch_config, columns) = snapshot_reader_inputs(&rs_ptr).await;
 
         let nullable_flags: Vec<bool> = columns.iter().map(|c| c.nullable).collect();
@@ -418,16 +403,53 @@ impl DatabaseDriverV1 {
         } else {
             Some(nullable_flags.as_slice())
         };
-        let reader = build_reader_from_rowset_data(
+        build_reader_from_rowset_data(
             &data,
             http_client,
             &prefetch_config,
             &self.wrapper_presets,
             flags,
         )
         .await
-        .context(QueryResponseProcessSnafu)?;
-        Ok(reader)
+        .context(QueryResponseProcessSnafu)
+    }
+
+    /// Builds a fresh Arrow [`RecordBatchReader`] while leaving the result-set
+    /// handle registered, so callers may request the stream again.
+    ///
+    /// Awaiting this only builds the reader. Drain chunked readers outside an
+    /// async runtime, on a dedicated thread, or with `spawn_blocking`.
+    pub async fn result_set_get_stream(
+        &self,
+        result_handle: Handle,
+    ) -> Result<Box<dyn RecordBatchReader + Send>, ApiError> {
+        let rs_ptr = self
+            .results
+            .get_obj(result_handle)
+            .with_context(|| InvalidArgumentSnafu {
+                argument: "result_handle: ResultSet handle not found".to_string(),
+            })?;
+        self.build_result_set_stream(rs_ptr).await
+    }
+
+    /// Atomically consumes a result-set handle and builds its Arrow stream.
+    ///
+    /// The handle is deregistered before stream construction, and remains
+    /// released whether construction succeeds or fails. Use this for one-shot
+    /// consumers to avoid pairing [`Self::result_set_get_stream`] with
+    /// [`Self::result_set_release`] and duplicating cleanup logic. As with
+    /// `result_set_get_stream`, drain chunked readers outside an async runtime.
+    pub async fn result_set_take_stream(
+        &self,
+        result_handle: Handle,
+    ) -> Result<Box<dyn RecordBatchReader + Send>, ApiError> {
+        let rs_ptr =
+            self.results
+                .take_obj(result_handle)
+                .with_context(|| InvalidArgumentSnafu {
+                    argument: "result_handle: ResultSet handle not found".to_string(),
+                })?;
+        self.build_result_set_stream(rs_ptr).await
     }
 
     /// Returns chunk metadata (inline data + remote chunk URLs) for this result set.
@@ -674,6 +696,65 @@ mod tests {
         drop(runtime);
     }
 
+    #[test]
+    fn result_set_take_stream_returns_reader_and_consumes_handle() {
+        let driver = DatabaseDriverV1::new();
+        let data: Data = serde_json::from_str(JSON_ROWSET)
+            .expect("fixture must deserialize into query_response::Data");
+        let descriptor = response_to_descriptor(&data, &WrapperPresets::default());
+        let reader_ctx = ReaderContext {
+            http_client: reqwest::Client::new(),
+            prefetch_config: PrefetchConfig::default(),
+        };
+        let handle = driver.create_result_set(descriptor, data.into_rowset_data(), reader_ctx);
+
+        let runtime = tokio::runtime::Runtime::new().expect("failed to build tokio runtime");
+        let reader = runtime
+            .block_on(driver.result_set_take_stream(handle))
+            .expect("taking an inline JSON result stream should succeed");
+
+        assert!(
+            runtime
+                .block_on(driver.result_set_take_stream(handle))
+                .is_err(),
+            "a consumed handle must not be reusable"
+        );
+        assert!(
+            driver.result_set_release(handle).is_err(),
+            "a consumed handle must already be released"
+        );
+        assert_id_name_reader(reader);
+        drop(runtime);
+    }
+
+    #[test]
+    fn result_set_take_stream_consumes_handle_when_stream_build_fails() {
+        let driver = DatabaseDriverV1::new();
+        let data: Data = serde_json::from_str(ARROW_DATA)
+            .expect("fixture must deserialize into query_response::Data");
+        let descriptor = response_to_descriptor(&data, &WrapperPresets::default());
+        let reader_ctx = ReaderContext {
+            http_client: reqwest::Client::new(),
+            prefetch_config: PrefetchConfig::default(),
+        };
+        let malformed = RowsetData::ArrowSingleChunk {
+            chunk_base64: "not valid base64".to_string(),
+        };
+        let handle = driver.create_result_set(descriptor, malformed, reader_ctx);
+
+        let runtime = tokio::runtime::Runtime::new().expect("failed to build tokio runtime");
+        assert!(
+            runtime
+                .block_on(driver.result_set_take_stream(handle))
+                .is_err(),
+            "malformed Arrow data must fail stream construction"
+        );
+        assert!(
+            driver.result_set_release(handle).is_err(),
+            "the handle must remain released after construction fails"
+        );
+    }
+
     #[test]
     fn result_set_get_stream_reads_arrow_rowset() {
         let driver = DatabaseDriverV1::new();

sf_core/src/apis/database_driver_v1/statement.rs

@@ -227,8 +227,7 @@ impl DatabaseDriverV1 {
                 }
                 .fail();
             };
-            let stream = self.result_set_get_stream(rs_info.handle).await?;
-            self.result_set_release(rs_info.handle)?;
+            let stream = self.result_set_take_stream(rs_info.handle).await?;
 
             let stmt_ptr =
                 self.statements

sf_core/src/handle_manager.rs

@@ -80,37 +80,41 @@ impl<T> HandleManager<T> {
         }
     }
 
-    pub fn delete_handle(&self, handle: Handle) -> bool {
-        let span = span!(target: "handle_manager", Level::INFO, "Deleting handle", handle_id = handle.id, handle_magic = handle.magic);
+    /// Atomically deregisters and returns the value for `handle`. Only one caller
+    /// can take a live handle; later lookups and takes fail.
+    pub(crate) fn take_obj(&self, handle: Handle) -> Option<Arc<T>> {
+        let span = span!(target: "handle_manager", Level::INFO, "HandleManager::take_obj", handle_id = handle.id, handle_magic = handle.magic);
         let _enter = span.enter();
         let index = handle.id as usize;
         let mut handles = self.handles.write_recover();
 
         if index >= handles.len() {
-            tracing::error!("Handle index out of bounds, cannot delete handle");
-            return false;
+            tracing::error!("Handle index out of bounds, cannot take object");
+            return None;
         }
 
         let handle_value = &mut handles[index];
-        let magic = handle_value.magic;
-
-        if magic != handle.magic {
-            tracing::error!("Handle magic mismatch, cannot delete handle");
-            return false;
+        if handle_value.magic != handle.magic {
+            tracing::error!("Handle magic mismatch, cannot take object");
+            return None;
         }
 
         match handle_value.value.take() {
-            Some(_) => {
-                tracing::trace!(target: "handle_manager", "Handle deleted successfully");
-                true
+            Some(value) => {
+                tracing::trace!(target: "handle_manager", "Handle taken successfully");
+                Some(value)
             }
             None => {
-                tracing::error!("Handle not found, cannot delete handle");
-                false
+                tracing::error!("Handle not found, cannot take object");
+                None
             }
         }
     }
 
+    pub fn delete_handle(&self, handle: Handle) -> bool {
+        self.take_obj(handle).is_some()
+    }
+
     /// Deregisters and returns every currently-live value matching `pred`,
     /// leaving non-matching handles untouched. Used by session reaping (e.g.
     /// `stream_transfer::reap_connection_streams`) to sweep up a connection's