Skip to main content

lance_io/object_store/providers/
local.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{collections::HashMap, sync::Arc};
5
6#[cfg(any(windows, test))]
7use crate::object_store::LocalDirOperations;
8use crate::object_store::{
9    DEFAULT_LOCAL_BLOCK_SIZE, DEFAULT_LOCAL_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore,
10    ObjectStoreParams, ObjectStoreProvider, StorageOptions,
11};
12use lance_core::Error;
13use lance_core::error::Result;
14use object_store::{local::LocalFileSystem, path::Path};
15#[cfg(any(windows, test))]
16use std::io::ErrorKind;
17use url::Url;
18
19#[derive(Default, Debug)]
20pub struct FileStoreProvider;
21
22#[cfg(any(windows, test))]
23#[derive(Debug)]
24struct FileSystemDirOperations {
25    local_file_system: LocalFileSystem,
26}
27
28#[cfg(any(windows, test))]
29#[async_trait::async_trait]
30impl LocalDirOperations for FileSystemDirOperations {
31    async fn remove_dir_all(&self, path: &Path) -> Result<()> {
32        let local_path = self.local_file_system.path_to_filesystem(path)?;
33        let object_store_path = path.to_string();
34        tokio::task::spawn_blocking(move || {
35            std::fs::remove_dir_all(local_path).map_err(|error| match error.kind() {
36                ErrorKind::NotFound => Error::not_found(object_store_path),
37                _ => Error::from(error),
38            })
39        })
40        .await
41        .map_err(|error| Error::io(format!("recursive directory removal task failed: {error}")))?
42    }
43}
44
45#[cfg(windows)]
46mod windows {
47    use std::path::PathBuf;
48
49    use super::*;
50
51    #[derive(Debug)]
52    pub(super) struct UncPath {
53        pub(super) root: PathBuf,
54        pub(super) relative_path: Path,
55        pub(super) store_prefix: String,
56    }
57
58    pub(super) fn extract_unc_path(url: &Url) -> Result<Option<UncPath>> {
59        if url.scheme() != "file" {
60            return Ok(None);
61        }
62
63        let Some(host) = url.host_str().filter(|host| *host != "localhost") else {
64            return Ok(None);
65        };
66        let encoded_path = url.path().strip_prefix('/').unwrap_or(url.path());
67        let (encoded_share, relative_path) =
68            encoded_path.split_once('/').unwrap_or((encoded_path, ""));
69        if encoded_share.is_empty() {
70            return Err(Error::invalid_input(format!(
71                "UNC URL '{}' is missing a share name",
72                url
73            )));
74        }
75
76        let share = Path::from_url_path(encoded_share).map_err(|error| {
77            Error::invalid_input(format!(
78                "Failed to parse share name from UNC URL '{}': {}",
79                url, error
80            ))
81        })?;
82        if share.parts_count() != 1 || share.as_ref().contains('\\') {
83            return Err(Error::invalid_input(format!(
84                "UNC URL '{}' has an invalid share name",
85                url
86            )));
87        }
88
89        Ok(Some(UncPath {
90            root: PathBuf::from(format!(r"\\{}\{}", host, share)),
91            relative_path: Path::from_url_path(relative_path).map_err(|error| {
92                Error::invalid_input(format!(
93                    "Failed to parse path '{}' from UNC URL '{}': {}",
94                    relative_path, url, error
95                ))
96            })?,
97            store_prefix: format!("{}${}/{}", url.scheme(), host, encoded_share),
98        }))
99    }
100}
101
102#[async_trait::async_trait]
103impl ObjectStoreProvider for FileStoreProvider {
104    async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result<ObjectStore> {
105        let block_size = params.block_size.unwrap_or(DEFAULT_LOCAL_BLOCK_SIZE);
106        let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default());
107        let download_retry_count = storage_options.download_retry_count();
108
109        #[cfg(windows)]
110        let (inner, local_dir_operations) = match windows::extract_unc_path(&base_path)? {
111            Some(unc_path) => {
112                let inner = LocalFileSystem::new_with_prefix(unc_path.root)?;
113                let operations = FileSystemDirOperations {
114                    local_file_system: inner.clone(),
115                };
116                (
117                    inner,
118                    Some(Arc::new(operations) as Arc<dyn LocalDirOperations>),
119                )
120            }
121            None => (LocalFileSystem::new(), None),
122        };
123        #[cfg(not(windows))]
124        let inner = LocalFileSystem::new();
125        #[cfg(not(windows))]
126        let local_dir_operations = None;
127
128        Ok(ObjectStore {
129            inner: Arc::new(inner),
130            local_dir_operations,
131            scheme: base_path.scheme().to_owned(),
132            block_size,
133            max_iop_size: *DEFAULT_MAX_IOP_SIZE,
134            use_constant_size_upload_parts: false,
135            list_is_lexically_ordered: false,
136            io_parallelism: DEFAULT_LOCAL_IO_PARALLELISM,
137            download_retry_count,
138            io_tracker: Default::default(),
139            store_prefix: self
140                .calculate_object_store_prefix(&base_path, params.storage_options())?,
141            // Listed in full: reading a directory is one local walk whatever the page size,
142            // so there is no request for a page to be pushed into.
143            paginated_lister: None,
144        })
145    }
146
147    fn extract_path(&self, url: &Url) -> Result<Path> {
148        #[cfg(windows)]
149        if let Some(unc_path) = windows::extract_unc_path(url)? {
150            return Ok(unc_path.relative_path);
151        }
152        if let Ok(file_path) = url.to_file_path()
153            && let Ok(path) = Path::from_absolute_path(&file_path)
154        {
155            return Ok(path);
156        }
157
158        Path::from_url_path(url.path()).map_err(|e| {
159            Error::invalid_input(format!("Failed to parse path '{}': {}", url.path(), e))
160        })
161    }
162
163    fn calculate_object_store_prefix(
164        &self,
165        url: &Url,
166        _storage_options: Option<&HashMap<String, String>>,
167    ) -> Result<String> {
168        #[cfg(windows)]
169        if let Some(unc_path) = windows::extract_unc_path(url)? {
170            return Ok(unc_path.store_prefix);
171        }
172
173        Ok(url.scheme().to_string())
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use std::fs::{create_dir_all, write};
180    use std::path::Path as StdPath;
181
182    use crate::object_store::uri_to_url;
183    #[cfg(unix)]
184    use std::os::unix::fs::symlink;
185    use tempfile::tempdir;
186
187    use super::*;
188
189    fn rooted_local_store(root: &StdPath) -> ObjectStore {
190        let inner = LocalFileSystem::new_with_prefix(root).unwrap();
191        let local_dir_operations = Arc::new(FileSystemDirOperations {
192            local_file_system: inner.clone(),
193        });
194        ObjectStore {
195            inner: Arc::new(inner),
196            local_dir_operations: Some(local_dir_operations),
197            scheme: "file".to_owned(),
198            block_size: DEFAULT_LOCAL_BLOCK_SIZE,
199            max_iop_size: *DEFAULT_MAX_IOP_SIZE,
200            use_constant_size_upload_parts: false,
201            list_is_lexically_ordered: false,
202            io_parallelism: DEFAULT_LOCAL_IO_PARALLELISM,
203            download_retry_count: 0,
204            io_tracker: Default::default(),
205            store_prefix: "file$rooted-test".to_owned(),
206            paginated_lister: None,
207        }
208    }
209
210    #[tokio::test]
211    async fn test_rooted_remove_dir_all_removes_tree() {
212        let sandbox = tempdir().unwrap();
213        let root = sandbox.path().join("share");
214        let dataset = root.join("dataset");
215        create_dir_all(dataset.join("nested")).unwrap();
216        write(dataset.join("nested/data"), "delete").unwrap();
217
218        rooted_local_store(&root)
219            .remove_dir_all(Path::from("dataset"))
220            .await
221            .unwrap();
222
223        assert!(!dataset.exists(), "recursive deletion must remove the tree");
224    }
225
226    #[cfg(unix)]
227    #[tokio::test]
228    async fn test_rooted_remove_dir_all_does_not_follow_directory_symlink() {
229        let sandbox = tempdir().unwrap();
230        let root = sandbox.path().join("share");
231        let dataset = root.join("dataset");
232        let outside = sandbox.path().join("outside");
233        create_dir_all(&dataset).unwrap();
234        create_dir_all(&outside).unwrap();
235        let sentinel = outside.join("sentinel");
236        write(&sentinel, "keep").unwrap();
237        symlink(&outside, dataset.join("link")).unwrap();
238
239        rooted_local_store(&root)
240            .remove_dir_all(Path::from("dataset"))
241            .await
242            .unwrap();
243
244        assert!(
245            sentinel.exists(),
246            "recursive deletion must not follow links"
247        );
248        assert!(!dataset.exists(), "recursive deletion must remove the tree");
249    }
250
251    #[test]
252    fn test_file_store_path() {
253        let provider = FileStoreProvider;
254
255        let cases = [
256            ("file:///", ""),
257            ("file:///usr/local/bin", "usr/local/bin"),
258            ("file-object-store:///path/to/file", "path/to/file"),
259            ("file:///path/to/foo/../bar", "path/to/bar"),
260        ];
261
262        for (uri, expected_path) in cases {
263            let url = uri_to_url(uri).unwrap();
264            let path = provider.extract_path(&url).unwrap();
265            assert_eq!(path.as_ref(), expected_path, "uri: '{}'", uri);
266        }
267    }
268
269    #[test]
270    fn test_calculate_object_store_prefix() {
271        let provider = FileStoreProvider;
272        assert_eq!(
273            "file",
274            provider
275                .calculate_object_store_prefix(&Url::parse("file:///etc").unwrap(), None)
276                .unwrap()
277        );
278    }
279
280    #[test]
281    fn test_calculate_object_store_prefix_for_file_object_store() {
282        let provider = FileStoreProvider;
283        assert_eq!(
284            "file-object-store",
285            provider
286                .calculate_object_store_prefix(
287                    &Url::parse("file-object-store:///etc").unwrap(),
288                    None
289                )
290                .unwrap()
291        );
292    }
293
294    #[test]
295    #[cfg(windows)]
296    fn test_file_store_path_windows() {
297        let provider = FileStoreProvider;
298
299        let cases = [
300            (
301                "C:\\Users\\ADMINI~1\\AppData\\Local\\",
302                "C:/Users/ADMINI~1/AppData/Local",
303            ),
304            (
305                "C:\\Users\\ADMINI~1\\AppData\\Local\\..\\",
306                "C:/Users/ADMINI~1/AppData",
307            ),
308            (
309                "file-object-store:///C:/Users/ADMINI~1/AppData/Local",
310                "C:/Users/ADMINI~1/AppData/Local",
311            ),
312            (
313                "file:///C:/Users/RUNNER~1/AppData/Local/Temp/tmpm49j_w0f",
314                "C:/Users/RUNNER~1/AppData/Local/Temp/tmpm49j_w0f",
315            ),
316            (
317                "file://192.168.0.1/My%20Share/data/my-dataset.lance",
318                "data/my-dataset.lance",
319            ),
320        ];
321
322        for (uri, expected_path) in cases {
323            let url = uri_to_url(uri).unwrap();
324            let path = provider.extract_path(&url).unwrap();
325            assert_eq!(path.as_ref(), expected_path);
326        }
327    }
328
329    #[test]
330    #[cfg(windows)]
331    fn test_unc_share_path() {
332        let url = Url::parse("file://server/My%20Share/data/my-dataset.lance").unwrap();
333        let unc_path = windows::extract_unc_path(&url).unwrap().unwrap();
334
335        assert_eq!(
336            unc_path.root,
337            std::path::PathBuf::from(r"\\server\My Share")
338        );
339        assert_eq!(unc_path.relative_path.as_ref(), "data/my-dataset.lance");
340        assert_eq!(unc_path.store_prefix, "file$server/My%20Share");
341
342        let object_store_url =
343            Url::parse("file-object-store://server/My%20Share/data/my-dataset.lance").unwrap();
344        assert!(
345            windows::extract_unc_path(&object_store_url)
346                .unwrap()
347                .is_none()
348        );
349    }
350}