Skip to main content

comindly_core/
storage.rs

1//! Index storage at any location LanceDB opens: a folder, `s3://`, `gs://` or `az://`.
2//! Credentials, region and endpoint come from the platform's own environment, never from comindly.
3
4use std::sync::Arc;
5
6use anyhow::{Context, Result};
7use object_store::path::Path as ObjPath;
8use object_store::ObjectStore;
9
10/// A LanceDB connection builder for `uri`, on the process's one session.
11pub fn connect(uri: &str) -> lancedb::connection::ConnectBuilder {
12    lancedb::connect(uri).session(session())
13}
14
15/// Caches and object stores shared by every connection, bounded instead of Lance's 7 GB default.
16fn session() -> Arc<lancedb::Session> {
17    static SESSION: std::sync::OnceLock<Arc<lancedb::Session>> = std::sync::OnceLock::new();
18    SESSION
19        .get_or_init(|| {
20            Arc::new(lancedb::Session::new(
21                256 << 20,
22                64 << 20,
23                Arc::new(lancedb::ObjectStoreRegistry::default()),
24            ))
25        })
26        .clone()
27}
28
29/// The object store and path for a remote `uri`, with the same credentials LanceDB uses.
30pub fn open(uri: &str) -> Result<(Arc<dyn ObjectStore>, ObjPath)> {
31    let (store, path) = crate::rt::block_on(lance_io::object_store::ObjectStore::from_uri(uri))
32        .with_context(|| format!("open {uri}"))?;
33    Ok((store.inner.clone(), path))
34}
35
36/// `text` with any `user:password@` removed from URLs in it.
37pub fn without_userinfo(text: &str) -> String {
38    let mut out = String::with_capacity(text.len());
39    let mut rest = text;
40    while let Some(i) = rest.find("://") {
41        let (head, tail) = rest.split_at(i + 3);
42        out.push_str(head);
43        let end = tail
44            .find(|c: char| c == '/' || c.is_whitespace())
45            .unwrap_or(tail.len());
46        match tail[..end].rfind('@') {
47            Some(at) => rest = &tail[at + 1..],
48            None => rest = tail,
49        }
50    }
51    out.push_str(rest);
52    out
53}
54
55/// A short, secret-free explanation of an error, with storage error replies turned into what to do.
56pub fn explain(raw: &str) -> String {
57    if raw.contains("does not support timezone offsets") {
58        return "The AWS SSO token cache has an expiry time the AWS SDK cannot read; sign in again \
59                (`aws sso login`)"
60            .into();
61    }
62    if raw.contains("Failed to get AWS credentials")
63        || raw.contains("CredentialsNotLoaded")
64        || raw.contains("no credentials")
65    {
66        return "No storage credentials: run comindly where the platform's usual credentials are \
67                available (an IAM or workload role, AWS_PROFILE, GOOGLE_APPLICATION_CREDENTIALS, \
68                AZURE_STORAGE_ACCOUNT_NAME …)"
69            .into();
70    }
71    let code = raw
72        .split("<Code>")
73        .nth(1)
74        .and_then(|r| r.split("</Code>").next())
75        .map(str::to_string);
76    let hint = match code.as_deref() {
77        Some("ExpiredToken" | "TokenRefreshRequired") => {
78            Some("the storage credentials have expired; renew them and try again")
79        }
80        Some("AccessDenied" | "AuthorizationFailure" | "AuthorizationPermissionMismatch") => Some(
81            "the storage service denied access; check this machine's permissions on the bucket",
82        ),
83        Some("NoSuchBucket" | "ContainerNotFound") => {
84            Some("the bucket does not exist (check the name and region)")
85        }
86        Some(
87            "InvalidAccessKeyId"
88            | "SignatureDoesNotMatch"
89            | "InvalidToken"
90            | "AuthenticationFailed",
91        ) => Some("the storage service rejected the credentials"),
92        _ => None,
93    };
94    match (hint, code) {
95        (Some(h), Some(c)) => format!("{h} ({c})"),
96        _ => redact(raw),
97    }
98}
99
100/// Removes session tokens and other long credential-like strings.
101fn redact(s: &str) -> String {
102    let mut out = String::with_capacity(s.len());
103    let mut word = String::new();
104    let flush = |w: &mut String, out: &mut String| {
105        let has = |f: fn(&char) -> bool| w.chars().any(|c| f(&c));
106        let path = w.starts_with('/') || w.starts_with('~');
107        let secretish = !path
108            && w.len() >= 60
109            && has(char::is_ascii_uppercase)
110            && has(char::is_ascii_lowercase)
111            && has(char::is_ascii_digit);
112        out.push_str(if secretish { "[redacted]" } else { w.as_str() });
113        w.clear();
114    };
115    for c in s.chars() {
116        if c.is_ascii_alphanumeric() || "+/=_-%".contains(c) {
117            word.push(c);
118        } else {
119            flush(&mut word, &mut out);
120            out.push(c);
121        }
122    }
123    flush(&mut word, &mut out);
124    out
125}
126
127/// An index in object storage (`s3://`, `gs://`, `az://`), not a folder on this machine.
128pub fn is_remote(index_root: &str) -> bool {
129    index_root.contains("://")
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn drops_credentials_from_urls() {
138        assert_eq!(
139            without_userinfo("fatal: https://oauth2:abc123@gitlab.example.com/a/b.git not found"),
140            "fatal: https://gitlab.example.com/a/b.git not found"
141        );
142        assert_eq!(
143            without_userinfo("git@github.com:a/b.git"),
144            "git@github.com:a/b.git"
145        );
146        assert_eq!(
147            without_userinfo("https://github.com/a/b"),
148            "https://github.com/a/b"
149        );
150    }
151
152    #[test]
153    fn explains_s3_replies_without_secrets() {
154        let token = &"ExampleSessionToken0123456789".repeat(3);
155        let raw = format!(
156            "ListRequest: <Error><Code>ExpiredToken</Code><Token-0>{token}</Token-0></Error>"
157        );
158        let e = explain(&raw);
159        assert!(e.contains("(ExpiredToken)") && !e.contains(token), "{e}");
160        let other = explain(&format!("connect failed: token={token} at s3://bucket/x"));
161        assert!(
162            !other.contains(token) && other.contains("s3://bucket/x"),
163            "{other}"
164        );
165        let path = format!("/home/{}", "long-folder-name-".repeat(6));
166        assert_eq!(explain(&path), path);
167        let offset = "Failed to get AWS credentials: InvalidField { field: \"expiresAt\", source: DateTimeParseError { kind: Invalid(\"Smithy does not support timezone offsets in RFC-3339 date times\") } }";
168        assert!(explain(offset).contains("aws sso login"));
169        let chain = "Failed to get AWS credentials: CredentialsNotLoaded(CredentialsNotLoaded { source: Some(ProviderChainError { attempts: [] }) })";
170        assert!(
171            explain(chain).starts_with("No storage credentials"),
172            "{}",
173            explain(chain)
174        );
175        let mixed = "/Users/Someone/Library/Caches/tmp-Ab12Cd34Ef56/index/_graph/symbols.lance";
176        assert_eq!(explain(mixed), mixed);
177    }
178}