From 9f2b29a24c17153988e3e858aed12fe4b87d4583 Mon Sep 17 00:00:00 2001
From: Blake Orth <borth@seerai.space>
Date: Thu, 16 Oct 2025 12:42:09 -0600
Subject: [PATCH] Adds instrumentation to delimited LIST operations in CLI  -
 Adds instrumentation around delimited list operations into the   
 instrumented object store  - Adds test cases for the new code

---
 .../src/object_storage/instrumented.rs        | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/datafusion-cli/src/object_storage/instrumented.rs b/datafusion-cli/src/object_storage/instrumented.rs
index 8acece315f76..94445ee64ef4 100644
--- a/datafusion-cli/src/object_storage/instrumented.rs
+++ b/datafusion-cli/src/object_storage/instrumented.rs
@@ -163,6 +163,28 @@ impl InstrumentedObjectStore {
 
         ret
     }
+
+    async fn instrumented_list_with_delimiter(
+        &self,
+        prefix: Option<&Path>,
+    ) -> Result<ListResult> {
+        let timestamp = Utc::now();
+        let start = Instant::now();
+        let ret = self.inner.list_with_delimiter(prefix).await?;
+        let elapsed = start.elapsed();
+
+        self.requests.lock().push(RequestDetails {
+            op: Operation::List,
+            path: prefix.cloned().unwrap_or_else(|| Path::from("")),
+            timestamp,
+            duration: Some(elapsed),
+            size: None,
+            range: None,
+            extra_display: None,
+        });
+
+        Ok(ret)
+    }
 }
 
 impl fmt::Display for InstrumentedObjectStore {
@@ -217,6 +239,10 @@ impl ObjectStore for InstrumentedObjectStore {
     }
 
     async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
+        if self.enabled() {
+            return self.instrumented_list_with_delimiter(prefix).await;
+        }
+
         self.inner.list_with_delimiter(prefix).await
     }
 
@@ -569,6 +595,29 @@ mod tests {
         assert!(request.extra_display.is_none());
     }
 
+    #[tokio::test]
+    async fn instrumented_store_list_with_delimiter() {
+        let (instrumented, path) = setup_test_store().await;
+
+        // By default no requests should be instrumented/stored
+        assert!(instrumented.requests.lock().is_empty());
+        let _ = instrumented.list_with_delimiter(Some(&path)).await.unwrap();
+        assert!(instrumented.requests.lock().is_empty());
+
+        instrumented.set_instrument_mode(InstrumentedObjectStoreMode::Trace);
+        assert!(instrumented.requests.lock().is_empty());
+        let _ = instrumented.list_with_delimiter(Some(&path)).await.unwrap();
+        assert_eq!(instrumented.requests.lock().len(), 1);
+
+        let request = instrumented.take_requests().pop().unwrap();
+        assert_eq!(request.op, Operation::List);
+        assert_eq!(request.path, path);
+        assert!(request.duration.is_some());
+        assert!(request.size.is_none());
+        assert!(request.range.is_none());
+        assert!(request.extra_display.is_none());
+    }
+
     #[test]
     fn request_details() {
         let rd = RequestDetails {