Skip to main content

datui_lib/
source.rs

1//! Input source detection for local paths vs remote URLs (S3, GCS, HTTP/HTTPS).
2
3use std::path::{Path, PathBuf};
4
5#[derive(Clone, Debug, Eq, PartialEq)]
6pub enum InputSource {
7    Local(PathBuf),
8    S3(String),
9    Gcs(String),
10    /// Azure Blob Storage, always as `abfss://container@account.dfs.core.windows.net/path`
11    /// whichever of its forms it was written in.
12    Azure(String),
13    Http(String),
14}
15
16/// Classifies the path as local, S3, GCS, or HTTP/HTTPS using string parsing only (no filesystem calls).
17pub fn input_source(path: &Path) -> InputSource {
18    let s = path.as_os_str().to_string_lossy();
19    if let Some(after_scheme) = s.find("://") {
20        let prefix = s[..after_scheme].to_lowercase();
21        let rest = s[after_scheme + 3..].to_string();
22        if prefix == "s3" || prefix == "s3a" {
23            return InputSource::S3(rest);
24        }
25        if prefix == "gs" || prefix == "gcs" {
26            return InputSource::Gcs(rest);
27        }
28        if let Some((account, container, key)) = azure_parts(&s) {
29            return InputSource::Azure(azure_url(&account, &container, &key));
30        }
31        if prefix == "http" || prefix == "https" {
32            return InputSource::Http(s.to_string());
33        }
34    }
35    InputSource::Local(path.to_path_buf())
36}
37
38/// Whether `path` is a URL datui reads rather than a path on this machine: whatever
39/// [`input_source`] places remotely, and the Azure forms that name no account, which
40/// are expanded when opened.
41pub fn is_remote_url(path: &Path) -> bool {
42    !matches!(input_source(path), InputSource::Local(_))
43        || path
44            .to_string_lossy()
45            .split_once("://")
46            .is_some_and(|(scheme, _)| is_azure_short_scheme(scheme))
47}
48
49/// Whether this build can open `path`: object stores need the `cloud` feature and web
50/// URLs the `http` feature; a local path always opens.
51pub fn opens_in_this_build(path: &Path) -> bool {
52    match input_source(path) {
53        InputSource::Local(_) if is_remote_url(path) => cfg!(feature = "cloud"),
54        InputSource::Local(_) => true,
55        InputSource::S3(_) | InputSource::Gcs(_) | InputSource::Azure(_) => {
56            cfg!(feature = "cloud")
57        }
58        InputSource::Http(_) => cfg!(feature = "http"),
59    }
60}
61
62/// `az`, `adl` and `azure`: Azure schemes that name a container but no account.
63pub(crate) fn is_azure_short_scheme(scheme: &str) -> bool {
64    matches!(scheme.to_ascii_lowercase().as_str(), "az" | "adl" | "azure")
65}
66
67/// The source an `s3://<id>@bucket/key` URL names, and the URL without it.
68///
69/// Only S3 URLs carry a source, and only for S3-compatible servers, whose bucket names
70/// repeat from one endpoint to the next. Bucket names cannot contain `@`, so an `@` in
71/// the first segment is always a source. Everything else comes back unchanged.
72pub fn split_source_id(url: &str) -> (Option<&str>, std::borrow::Cow<'_, str>) {
73    let Some((scheme, rest)) = url.split_once("://") else {
74        return (None, url.into());
75    };
76    if !matches!(scheme.to_ascii_lowercase().as_str(), "s3" | "s3a") {
77        return (None, url.into());
78    }
79    let first = rest.split('/').next().unwrap_or(rest);
80    match first.split_once('@') {
81        Some((id, _)) if !id.is_empty() => {
82            let plain = format!("{scheme}://{}", &rest[id.len() + 1..]);
83            (Some(id), plain.into())
84        }
85        _ => (None, url.into()),
86    }
87}
88
89/// The account, container and path of an Azure Blob Storage URL.
90///
91/// Accepts the forms that name the account: `abfss://` and `abfs://`
92/// (`container@account.dfs.core.windows.net/path`), and `https://` on the blob or dfs
93/// endpoint (`account.blob.core.windows.net/container/path`). `az://container/path`
94/// does not name the account, so it is not one of them. The path comes back without a
95/// leading slash, and a trailing slash is kept, since it is what marks a directory.
96pub fn azure_parts(url: &str) -> Option<(String, String, String)> {
97    let (scheme, rest) = url.split_once("://")?;
98    let scheme = scheme.to_ascii_lowercase();
99    let (host_part, path) = match rest.split_once('/') {
100        Some((host, path)) => (host, path),
101        None => (rest, ""),
102    };
103    let account_of = |host: &str| {
104        let host = host.to_ascii_lowercase();
105        [".dfs.core.windows.net", ".blob.core.windows.net"]
106            .iter()
107            .find_map(|suffix| host.strip_suffix(suffix).map(str::to_string))
108            .filter(|account| !account.is_empty() && !account.contains('.'))
109    };
110    match scheme.as_str() {
111        "abfss" | "abfs" => {
112            let (container, host) = host_part.split_once('@')?;
113            let account = account_of(host)?;
114            (!container.is_empty()).then(|| (account, container.to_string(), path.to_string()))
115        }
116        "https" | "http" => {
117            let account = account_of(host_part)?;
118            let (container, path) = match path.split_once('/') {
119                Some((container, path)) => (container, path),
120                None => (path, ""),
121            };
122            (!container.is_empty()).then(|| (account, container.to_string(), path.to_string()))
123        }
124        _ => None,
125    }
126}
127
128/// The canonical URL for a place in Azure Blob Storage.
129pub fn azure_url(account: &str, container: &str, path: &str) -> String {
130    format!(
131        "abfss://{container}@{account}.dfs.core.windows.net/{}",
132        path.trim_start_matches('/')
133    )
134}
135
136/// A cloud place in the form used for equality comparisons. Equivalent Azure URL
137/// forms become `abfss://`, and a trailing slash does not distinguish a place.
138pub(crate) fn canonical_cloud_place(url: &str) -> String {
139    let canonical = match azure_parts(url) {
140        Some((account, container, path)) => azure_url(&account, &container, &path),
141        None => url.to_string(),
142    };
143    canonical.trim_end_matches('/').to_string()
144}
145
146/// A cloud location whose shape is a prefix or a glob rather than one object.
147pub(crate) fn is_prefix_or_glob(url: &str) -> bool {
148    url.ends_with('/') || url.contains('*')
149}
150
151/// The characters Polars expands a path on (`polars_io::path_utils::has_glob`).
152pub(crate) fn has_glob_chars(path: &Path) -> bool {
153    path.as_os_str().to_string_lossy().contains(['*', '?', '['])
154}
155
156/// Whether a path is a pattern to expand rather than a name: it carries a glob
157/// character and nothing on disk has that name. An existing `d[1].csv` or `a*b.csv`
158/// is that file; read as a glob, `d[1].csv` is `d1.csv` and `x?.csv` is every
159/// two-letter name. This is the `glob` flag for every Polars scan of a local path.
160pub(crate) fn expands_as_glob(path: &Path) -> bool {
161    has_glob_chars(path) && std::fs::symlink_metadata(path).is_err()
162}
163
164/// A path for a Polars reader that always expands globs (its NDJSON scan has no
165/// `glob` flag): an existing name with a glob character comes back escaped, so the
166/// pattern matches that file alone.
167pub(crate) fn polars_literal_path(
168    path: &Path,
169) -> polars::prelude::PolarsResult<polars::prelude::PlRefPath> {
170    if !has_glob_chars(path) || expands_as_glob(path) {
171        return polars::prelude::PlRefPath::try_from_path(path);
172    }
173    // Escaped as Polars will read it: on Windows that text has `/` separators and no
174    // `\\?\` prefix, so neither a separator nor the prefix's `?` is touched.
175    let text = polars::prelude::PlRefPath::try_from_path(path)?;
176    Ok(polars::prelude::PlRefPath::new(
177        escape_glob(text.as_str()).as_str(),
178    ))
179}
180
181/// `text` as a glob that matches only itself: each glob character in brackets.
182pub(crate) fn escape_glob(text: &str) -> String {
183    let mut escaped = String::with_capacity(text.len() + 8);
184    for c in text.chars() {
185        if matches!(c, '*' | '?' | '[' | ']') {
186            escaped.extend(['[', c, ']']);
187        } else {
188            escaped.push(c);
189        }
190    }
191    escaped
192}
193
194/// True when the path names an object-store location datui scans in place, with range
195/// requests, rather than downloads to a temporary file first: Parquet, or a prefix or
196/// glob of it. A downloaded object reaches the schema phase under its display URL, and
197/// this is what keeps it from being treated as a remote scan.
198pub(crate) fn scans_in_place(path: &Path) -> bool {
199    if !matches!(
200        input_source(path),
201        InputSource::S3(_) | InputSource::Gcs(_) | InputSource::Azure(_)
202    ) {
203        return false;
204    }
205    let url = path.to_string_lossy();
206    let (_, ext) = url_path_extension(&url);
207    !cloud_path_should_download(ext.as_deref(), is_prefix_or_glob(&url))
208}
209
210/// Returns the path segment and file extension for URL format inference.
211/// For S3, path part is everything after `://` (bucket/key). For HTTP/HTTPS, path part is the URL path only (host stripped).
212pub(crate) fn url_path_extension(url: &str) -> (String, Option<String>) {
213    let path_part = if let Some(i) = url.find("://") {
214        let scheme = url[..i].to_lowercase();
215        let after = &url[i + 3..];
216        if scheme == "http" || scheme == "https" {
217            after
218                .find('/')
219                .map(|j| after[j + 1..].to_string())
220                .unwrap_or_default()
221        } else {
222            after.to_string()
223        }
224    } else {
225        String::new()
226    };
227    let last_segment = path_part.rsplit('/').next().unwrap_or(&path_part);
228    let ext = std::path::Path::new(last_segment)
229        .extension()
230        .and_then(|e| e.to_str())
231        .map(String::from);
232    (path_part, ext)
233}
234
235/// The extension a downloaded copy of `url` should keep, so it opens as what it is:
236/// `csv.gz` rather than `gz` for a compressed file, since a temporary `.gz` does not
237/// say what is inside it.
238#[cfg(any(feature = "http", feature = "cloud"))]
239pub(crate) fn download_suffix(url: &str) -> Option<String> {
240    let (path_part, ext) = url_path_extension(url);
241    let ext = ext?;
242    const COMPRESSION: [&str; 6] = ["gz", "zst", "bz2", "xz", "lz4", "zip"];
243    if !COMPRESSION.iter().any(|c| ext.eq_ignore_ascii_case(c)) {
244        return Some(ext);
245    }
246    let name = path_part.rsplit('/').next().unwrap_or(&path_part);
247    let stem = &name[..name.len() - ext.len() - 1];
248    match Path::new(stem).extension().and_then(|e| e.to_str()) {
249        Some(inner) => Some(format!("{inner}.{ext}")),
250        None => Some(ext),
251    }
252}
253
254/// For S3/GCS: Polars can only scan Parquet directly. So we pass through only when the path is
255/// Parquet or looks like a directory/glob (no extension, trailing slash, or *). All other paths
256/// (e.g. .csv, .json, .gz, .csv.gz) must be downloaded first.
257/// Returns true when the path should be downloaded to temp instead of passed to Polars.
258pub(crate) fn cloud_path_should_download(ext: Option<&str>, is_glob: bool) -> bool {
259    if is_glob {
260        return false;
261    }
262    match ext {
263        None => false,
264        Some(e) => !e.eq_ignore_ascii_case("parquet"),
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271
272    #[test]
273    fn an_existing_name_is_never_a_glob() {
274        let dir = tempfile::tempdir().unwrap();
275        let file = dir.path().join("d[1].csv");
276        std::fs::write(&file, "a\n1\n").unwrap();
277        assert!(!expands_as_glob(&file));
278        assert!(expands_as_glob(&dir.path().join("d[2].csv")));
279        assert!(expands_as_glob(&dir.path().join("*.csv")));
280        assert!(!expands_as_glob(&dir.path().join("plain.csv")));
281        assert!(!expands_as_glob(dir.path()));
282
283        let escaped = polars_literal_path(&file).unwrap();
284        assert!(escaped.as_str().ends_with("d[[]1[]].csv"), "{escaped:?}");
285        let pattern = dir.path().join("d[2].csv");
286        let kept = polars_literal_path(&pattern).unwrap();
287        // Polars writes a Windows path with forward slashes.
288        assert_eq!(
289            kept.as_str().replace('\\', "/"),
290            pattern.to_str().unwrap().replace('\\', "/")
291        );
292    }
293
294    /// The escaped name reads that one file through the NDJSON scan, which always
295    /// expands, whatever else the name holds: `]` alone, braces (no glob meaning to
296    /// Polars), a backslash (a plain character on Unix), and a directory with `[`.
297    #[test]
298    fn an_escaped_name_reads_that_file_alone() {
299        use polars::prelude::{LazyFileListReader, LazyJsonLineReader};
300        let dir = tempfile::tempdir().unwrap();
301        let nested = dir.path().join("set[1]");
302        std::fs::create_dir(&nested).unwrap();
303        std::fs::create_dir(dir.path().join("set1")).unwrap();
304        let mut names = vec!["d[1]", "h]", "i{j,k}", "[!x]", "set[1]/p[a]"];
305        if cfg!(unix) {
306            names.extend(["a*b", "x?", "e\\f", "g[x]*?"]);
307        }
308        // What the unescaped patterns would also match.
309        for decoy in ["d1", "ha", "set1/pa", "ab", "xy", "gx", "y"] {
310            std::fs::write(dir.path().join(format!("{decoy}.jsonl")), "{\"v\": 0}\n").unwrap();
311        }
312        for name in names {
313            let file = dir.path().join(format!("{name}.jsonl"));
314            std::fs::write(
315                &file,
316                format!("{{\"v\": \"{}\"}}\n", name.replace('\\', "/")),
317            )
318            .unwrap();
319            let lf = LazyJsonLineReader::new(polars_literal_path(&file).unwrap())
320                .finish()
321                .unwrap();
322            let df = lf.collect().unwrap_or_else(|e| panic!("{name}: {e}"));
323            assert_eq!(df.height(), 1, "{name}");
324            let v = df
325                .column("v")
326                .unwrap()
327                .str()
328                .unwrap()
329                .get(0)
330                .map(str::to_string);
331            assert_eq!(
332                v.as_deref(),
333                Some(name.replace('\\', "/").as_str()),
334                "{name}"
335            );
336        }
337    }
338
339    #[test]
340    fn every_url_datui_reads_is_remote() {
341        for url in [
342            "s3://bucket/key.parquet",
343            "s3a://bucket/key.parquet",
344            "gs://bucket/dir/",
345            "gcs://bucket/dir/",
346            "abfss://release@overturemapswestus2.dfs.core.windows.net/2026-09-23.1/",
347            "abfs://container@account.dfs.core.windows.net/x.parquet",
348            "https://account.blob.core.windows.net/container/x.csv",
349            "az://container/x.csv",
350            "adl://container/x.csv",
351            "azure://container/x.csv",
352            "http://example.com/data.csv",
353            "https://example.com/data.csv",
354        ] {
355            assert!(is_remote_url(Path::new(url)), "{url}");
356        }
357        for path in ["/tmp/file.parquet", "relative.csv", ".", "data/2024.csv"] {
358            assert!(!is_remote_url(Path::new(path)), "{path}");
359        }
360    }
361
362    #[test]
363    fn input_source_local_path() {
364        let p = PathBuf::from("/tmp/file.parquet");
365        assert!(matches!(input_source(&p), InputSource::Local(_)));
366        let p = PathBuf::from("relative.csv");
367        assert!(matches!(input_source(&p), InputSource::Local(_)));
368        let p = PathBuf::from(".");
369        assert!(matches!(input_source(&p), InputSource::Local(_)));
370    }
371
372    #[test]
373    fn input_source_s3() {
374        let p = PathBuf::from("s3://bucket/key.parquet");
375        match input_source(&p) {
376            InputSource::S3(rest) => assert_eq!(rest, "bucket/key.parquet"),
377            _ => panic!("expected S3"),
378        }
379        let p = PathBuf::from("S3://my-bucket/path/to/file.csv");
380        match input_source(&p) {
381            InputSource::S3(rest) => assert_eq!(rest, "my-bucket/path/to/file.csv"),
382            _ => panic!("expected S3"),
383        }
384    }
385
386    #[test]
387    fn a_source_is_split_off_s3_urls_only() {
388        assert_eq!(
389            split_source_id("s3://onprem@sales/2024/q3.parquet"),
390            (Some("onprem"), "s3://sales/2024/q3.parquet".into())
391        );
392        assert_eq!(
393            split_source_id("s3://onprem@sales"),
394            (Some("onprem"), "s3://sales".into())
395        );
396        assert_eq!(
397            split_source_id("s3://sales/a@b.parquet"),
398            (None, "s3://sales/a@b.parquet".into())
399        );
400        assert_eq!(
401            split_source_id("gs://bucket/key"),
402            (None, "gs://bucket/key".into())
403        );
404        assert_eq!(
405            split_source_id("https://user@host/file.csv"),
406            (None, "https://user@host/file.csv".into())
407        );
408        assert_eq!(split_source_id("s3://@sales"), (None, "s3://@sales".into()));
409    }
410
411    #[test]
412    fn azure_urls_in_every_form_that_names_the_account_become_one() {
413        let canonical = "abfss://datui-test@datalake001.dfs.core.windows.net/demo/fred/";
414        for url in [
415            canonical,
416            "abfs://datui-test@datalake001.dfs.core.windows.net/demo/fred/",
417            "https://datalake001.blob.core.windows.net/datui-test/demo/fred/",
418            "https://DataLake001.dfs.core.windows.net/datui-test/demo/fred/",
419        ] {
420            assert_eq!(
421                input_source(Path::new(url)),
422                InputSource::Azure(canonical.to_string()),
423                "{url}"
424            );
425        }
426        assert_eq!(
427            azure_parts("abfss://c@acct.dfs.core.windows.net"),
428            Some(("acct".to_string(), "c".to_string(), String::new()))
429        );
430        // No account, or not Azure at all.
431        assert_eq!(azure_parts("az://container/path"), None);
432        assert!(matches!(
433            input_source(Path::new("https://example.com/c/x.csv")),
434            InputSource::Http(_)
435        ));
436        assert!(scans_in_place(Path::new(
437            "abfss://c@acct.dfs.core.windows.net/x.parquet"
438        )));
439    }
440
441    #[test]
442    fn input_source_http() {
443        let p = PathBuf::from("https://example.com/data.parquet");
444        match input_source(&p) {
445            InputSource::Http(u) => assert_eq!(u, "https://example.com/data.parquet"),
446            _ => panic!("expected Http"),
447        }
448        let p = PathBuf::from("http://host/path/file.csv");
449        match input_source(&p) {
450            InputSource::Http(u) => assert_eq!(u, "http://host/path/file.csv"),
451            _ => panic!("expected Http"),
452        }
453    }
454
455    #[test]
456    fn input_source_gcs() {
457        let p = PathBuf::from("gs://my-bucket/path/file.parquet");
458        match input_source(&p) {
459            InputSource::Gcs(rest) => assert_eq!(rest, "my-bucket/path/file.parquet"),
460            _ => panic!("expected Gcs"),
461        }
462        let p = PathBuf::from("gcs://bucket/key.parquet");
463        match input_source(&p) {
464            InputSource::Gcs(rest) => assert_eq!(rest, "bucket/key.parquet"),
465            _ => panic!("expected Gcs"),
466        }
467    }
468
469    #[test]
470    fn input_source_unknown_scheme_stays_local() {
471        let p = PathBuf::from("file:///tmp/foo.parquet");
472        assert!(matches!(input_source(&p), InputSource::Local(_)));
473    }
474
475    #[test]
476    fn url_path_extension_s3() {
477        let (path, ext) = url_path_extension("s3://bucket/key.parquet");
478        assert_eq!(path, "bucket/key.parquet");
479        assert_eq!(ext.as_deref(), Some("parquet"));
480        let (path, ext) = url_path_extension("s3://b/path/to/file.csv");
481        assert_eq!(path, "b/path/to/file.csv");
482        assert_eq!(ext.as_deref(), Some("csv"));
483    }
484
485    #[cfg(any(feature = "http", feature = "cloud"))]
486    #[test]
487    fn a_download_keeps_what_the_compressed_file_holds() {
488        assert_eq!(
489            download_suffix("s3://b/edge/100%.csv.gz").as_deref(),
490            Some("csv.gz")
491        );
492        assert_eq!(
493            download_suffix("https://x.com/a/log.json.zst").as_deref(),
494            Some("json.zst")
495        );
496        assert_eq!(download_suffix("gs://b/data.csv").as_deref(), Some("csv"));
497        assert_eq!(download_suffix("s3://b/archive.gz").as_deref(), Some("gz"));
498        assert_eq!(download_suffix("s3://b/no-extension"), None);
499    }
500
501    #[test]
502    fn url_path_extension_https() {
503        let (path, ext) = url_path_extension("https://example.com/dir/file.parquet");
504        assert_eq!(path, "dir/file.parquet");
505        assert_eq!(ext.as_deref(), Some("parquet"));
506        let (_, ext) = url_path_extension("https://x.com/file.csv.gz");
507        assert_eq!(ext.as_deref(), Some("gz"));
508    }
509
510    #[test]
511    fn only_parquet_prefixes_and_globs_are_scanned_in_place() {
512        assert!(scans_in_place(Path::new("s3://bucket/obj.parquet")));
513        assert!(scans_in_place(Path::new("gs://bucket/prefix/")));
514        assert!(scans_in_place(Path::new("s3://bucket/year=*/*.parquet")));
515        // Downloaded first, then opened as a local file: not a remote scan.
516        assert!(!scans_in_place(Path::new("s3://bucket/data.csv")));
517        assert!(!scans_in_place(Path::new("s3://bucket/data.csv.gz")));
518        assert!(!scans_in_place(Path::new(
519            "https://example.com/data.parquet"
520        )));
521        assert!(!scans_in_place(Path::new("/data/local.parquet")));
522    }
523
524    #[test]
525    fn cloud_path_should_download() {
526        assert!(super::cloud_path_should_download(Some("csv"), false));
527        assert!(super::cloud_path_should_download(Some("gz"), false));
528        assert!(super::cloud_path_should_download(Some("csv.gz"), false));
529        assert!(!super::cloud_path_should_download(Some("parquet"), false));
530        assert!(!super::cloud_path_should_download(None, false));
531        assert!(!super::cloud_path_should_download(Some("csv"), true));
532        assert!(!super::cloud_path_should_download(Some("parquet"), true));
533    }
534}