Skip to main content

geoarrow_flatgeobuf/reader/
object_store.rs

1//! Integration with the [`object_store`] crate.
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use bytes::Bytes;
7use http_range_client::{AsyncHttpRangeClient, Result as HTTPRangeClientResult};
8use object_store::path::Path;
9use object_store::{ObjectStore, ObjectStoreExt};
10
11/// A wrapper around an [`ObjectStore`] that implements the [`AsyncHttpRangeClient`] trait.
12#[derive(Debug, Clone)]
13pub struct ObjectStoreWrapper {
14    store: Arc<dyn ObjectStore>,
15    location: Path,
16}
17
18impl ObjectStoreWrapper {
19    /// Creates a new [`ObjectStoreWrapper`] with the given store and location.
20    pub fn new(store: Arc<dyn ObjectStore>, location: Path) -> Self {
21        Self { store, location }
22    }
23}
24
25#[async_trait]
26impl AsyncHttpRangeClient for ObjectStoreWrapper {
27    /// Send a GET range request
28    async fn get_range(&self, _url: &str, range: &str) -> HTTPRangeClientResult<Bytes> {
29        assert!(range.starts_with("bytes="));
30
31        let split_range = range[6..].split('-').collect::<Vec<_>>();
32        let start_range = split_range[0].parse::<u64>().unwrap();
33
34        // Add one to the range because HTTP range strings are end-inclusive (I think)
35        let end_range = split_range[1].parse::<u64>().unwrap() + 1;
36
37        let bytes = self
38            .store
39            .get_range(&self.location, start_range..end_range)
40            .await
41            .unwrap();
42        Ok(bytes)
43    }
44
45    /// Send a HEAD request and return response header value
46    async fn head_response_header(
47        &self,
48        _url: &str,
49        header: &str,
50    ) -> HTTPRangeClientResult<Option<String>> {
51        // This is a massive hack to align APIs
52        if header == "content-length" {
53            let meta = self.store.head(&self.location).await.unwrap();
54            Ok(Some(meta.size.to_string()))
55        } else {
56            Ok(None)
57        }
58    }
59}