feat: robustify pre-signed URL checks

Author: roeap

kernel/src/engine/default/filesystem.rs

@@ -7,6 +7,7 @@ use object_store::path::Path;
 use object_store::{DynObjectStore, ObjectStore};
 use url::Url;
 
+use super::UrlExt;
 use crate::engine::default::executor::TaskExecutor;
 use crate::{DeltaResult, Error, FileMeta, FileSlice, FileSystemClient};
 
@@ -131,18 +132,15 @@ impl<E: TaskExecutor> FileSystemClient for ObjectStoreFileSystemClient<E> {
                     };
                     let store = store.clone();
                     async move {
-                        match url.scheme() {
-                            "http" | "https" => {
-                                // have to annotate type here or rustc can't figure it out
-                                Ok::<bytes::Bytes, Error>(reqwest::get(url).await?.bytes().await?)
-                            }
-                            _ => {
-                                if let Some(rng) = range {
-                                    Ok(store.get_range(&path, rng).await?)
-                                } else {
-                                    let result = store.get(&path).await?;
-                                    Ok(result.bytes().await?)
-                                }
+                        if url.is_presigned() {
+                            // have to annotate type here or rustc can't figure it out
+                            Ok::<bytes::Bytes, Error>(reqwest::get(url).await?.bytes().await?)
+                        } else {
+                            if let Some(rng) = range {
+                                Ok(store.get_range(&path, rng).await?)
+                            } else {
+                                let result = store.get(&path).await?;
+                                Ok(result.bytes().await?)
                             }
                         }
                     }

kernel/src/engine/default/mod.rs

@@ -158,6 +158,16 @@ impl<E: TaskExecutor> Engine for DefaultEngine<E> {
     }
 }
 
+pub(self) trait UrlExt {
+    fn is_presigned(&self) -> bool;
+}
+
+impl UrlExt for Url {
+    fn is_presigned(&self) -> bool {
+        matches!(self.scheme(), "http" | "https") && self.query().is_some()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::executor::tokio::TokioBackgroundExecutor;

kernel/src/engine/default/parquet.rs

@@ -4,23 +4,24 @@ use std::collections::HashMap;
 use std::ops::Range;
 use std::sync::Arc;
 
-use crate::arrow::array::builder::{MapBuilder, MapFieldNames, StringBuilder};
-use crate::arrow::array::{BooleanArray, Int64Array, RecordBatch, StringArray};
-use crate::parquet::arrow::arrow_reader::{
-    ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder,
-};
-use crate::parquet::arrow::arrow_writer::ArrowWriter;
-use crate::parquet::arrow::async_reader::{ParquetObjectReader, ParquetRecordBatchStreamBuilder};
 use futures::StreamExt;
 use object_store::path::Path;
 use object_store::DynObjectStore;
 use uuid::Uuid;
 
 use super::file_stream::{FileOpenFuture, FileOpener, FileStream};
+use super::UrlExt;
+use crate::arrow::array::builder::{MapBuilder, MapFieldNames, StringBuilder};
+use crate::arrow::array::{BooleanArray, Int64Array, RecordBatch, StringArray};
 use crate::engine::arrow_data::ArrowEngineData;
 use crate::engine::arrow_utils::{fixup_parquet_read, generate_mask, get_requested_indices};
 use crate::engine::default::executor::TaskExecutor;
 use crate::engine::parquet_row_group_skipping::ParquetRowGroupSkipping;
+use crate::parquet::arrow::arrow_reader::{
+    ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder,
+};
+use crate::parquet::arrow::arrow_writer::ArrowWriter;
+use crate::parquet::arrow::async_reader::{ParquetObjectReader, ParquetRecordBatchStreamBuilder};
 use crate::schema::SchemaRef;
 use crate::{
     DeltaResult, EngineData, Error, ExpressionRef, FileDataReadResultIterator, FileMeta,
@@ -191,18 +192,19 @@ impl<E: TaskExecutor> ParquetHandler for DefaultParquetHandler<E> {
         //   -> reqwest to get data
         //   -> parse to parquet
         // SAFETY: we did is_empty check above, this is ok.
-        let file_opener: Box<dyn FileOpener> = match files[0].location.scheme() {
-            "http" | "https" => Box::new(PresignedUrlOpener::new(
+        let file_opener: Box<dyn FileOpener> = if files[0].location.is_presigned() {
+            Box::new(PresignedUrlOpener::new(
                 1024,
                 physical_schema.clone(),
                 predicate,
-            )),
-            _ => Box::new(ParquetOpener::new(
+            ))
+        } else {
+            Box::new(ParquetOpener::new(
                 1024,
                 physical_schema.clone(),
                 predicate,
                 self.store.clone(),
-            )),
+            ))
         };
         FileStream::new_async_read_iterator(
             self.task_executor.clone(),