fix: prevent despecialization of object store methods (#3506)

Adding enforcement that wrappers implement all ObjectStore methods. If
they don't, then the specialized implementations in the base
`ObjectStore` could be bypasses, causing key optimizations to be missed
out on.

I didn't see any real cases here, but adding this to be cautious.
`delete_stream()` a common one we might worry about it, but it seems
like tracing store already had this.

Author: wjones127

rust/lance-io/src/object_store/tracing.rs

@@ -55,6 +55,7 @@ impl std::fmt::Display for TracedObjectStore {
 }
 
 #[async_trait::async_trait]
+#[deny(clippy::missing_trait_methods)]
 impl object_store::ObjectStore for TracedObjectStore {
     #[instrument(level = "debug", skip(self, bytes))]
     async fn put(&self, location: &Path, bytes: PutPayload) -> OSResult<PutResult> {
@@ -71,6 +72,17 @@ impl object_store::ObjectStore for TracedObjectStore {
         self.target.put_opts(location, bytes, opts).await
     }
 
+    async fn put_multipart(
+        &self,
+        location: &Path,
+    ) -> OSResult<Box<dyn object_store::MultipartUpload>> {
+        let upload = self.target.put_multipart(location).await?;
+        Ok(Box::new(TracedMultipartUpload {
+            target: upload,
+            write_span: debug_span!("put_multipart"),
+        }))
+    }
+
     async fn put_multipart_opts(
         &self,
         location: &Path,
@@ -83,6 +95,11 @@ impl object_store::ObjectStore for TracedObjectStore {
         }))
     }
 
+    #[instrument(level = "debug", skip(self, location))]
+    async fn get(&self, location: &Path) -> OSResult<GetResult> {
+        self.target.get(location).await
+    }
+
     #[instrument(level = "debug", skip(self, options))]
     async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
         self.target.get_opts(location, options).await
@@ -121,6 +138,15 @@ impl object_store::ObjectStore for TracedObjectStore {
         self.target.list(prefix)
     }
 
+    #[instrument(level = "debug", skip(self))]
+    fn list_with_offset(
+        &self,
+        prefix: Option<&Path>,
+        offset: &Path,
+    ) -> BoxStream<'_, OSResult<ObjectMeta>> {
+        self.target.list_with_offset(prefix, offset)
+    }
+
     #[instrument(level = "debug", skip(self))]
     async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
         self.target.list_with_delimiter(prefix).await
@@ -136,6 +162,11 @@ impl object_store::ObjectStore for TracedObjectStore {
         self.target.rename(from, to).await
     }
 
+    #[instrument(level = "debug", skip(self))]
+    async fn rename_if_not_exists(&self, from: &Path, to: &Path) -> OSResult<()> {
+        self.target.rename_if_not_exists(from, to).await
+    }
+
     #[instrument(level = "debug", skip(self))]
     async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> OSResult<()> {
         self.target.copy_if_not_exists(from, to).await

rust/lance/src/utils/test.rs

@@ -325,7 +325,13 @@ impl IoTrackingStore {
 }
 
 #[async_trait::async_trait]
+#[deny(clippy::missing_trait_methods)]
 impl ObjectStore for IoTrackingStore {
+    async fn put(&self, location: &Path, bytes: PutPayload) -> OSResult<PutResult> {
+        self.record_write(bytes.content_length() as u64);
+        self.target.put(location, bytes).await
+    }
+
     async fn put_opts(
         &self,
         location: &Path,
@@ -336,6 +342,14 @@ impl ObjectStore for IoTrackingStore {
         self.target.put_opts(location, bytes, opts).await
     }
 
+    async fn put_multipart(&self, location: &Path) -> OSResult<Box<dyn MultipartUpload>> {
+        let target = self.target.put_multipart(location).await?;
+        Ok(Box::new(IoTrackingMultipartUpload {
+            target,
+            stats: self.stats.clone(),
+        }))
+    }
+
     async fn put_multipart_opts(
         &self,
         location: &Path,
@@ -348,6 +362,15 @@ impl ObjectStore for IoTrackingStore {
         }))
     }
 
+    async fn get(&self, location: &Path) -> OSResult<GetResult> {
+        let result = self.target.get(location).await;
+        if let Ok(result) = &result {
+            let num_bytes = result.range.end - result.range.start;
+            self.record_read(num_bytes as u64);
+        }
+        result
+    }
+
     async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
         let result = self.target.get_opts(location, options).await;
         if let Ok(result) = &result {
@@ -394,6 +417,15 @@ impl ObjectStore for IoTrackingStore {
         self.target.list(prefix)
     }
 
+    fn list_with_offset(
+        &self,
+        prefix: Option<&Path>,
+        offset: &Path,
+    ) -> BoxStream<'_, OSResult<ObjectMeta>> {
+        self.record_read(0);
+        self.target.list_with_offset(prefix, offset)
+    }
+
     async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
         self.record_read(0);
         self.target.list_with_delimiter(prefix).await
@@ -409,6 +441,11 @@ impl ObjectStore for IoTrackingStore {
         self.target.rename(from, to).await
     }
 
+    async fn rename_if_not_exists(&self, from: &Path, to: &Path) -> OSResult<()> {
+        self.record_write(0);
+        self.target.rename_if_not_exists(from, to).await
+    }
+
     async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> OSResult<()> {
         self.record_write(0);
         self.target.copy_if_not_exists(from, to).await