quarb_objstore/
lib.rs

1//! Object-store adapter for the Quarb query engine: Google Cloud
2//! Storage and Amazon S3 buckets as lazy directory trees.
3//!
4//! Object keys with `/` separators span the tree the way a
5//! filesystem does — the adapter lists one "directory" per touch
6//! (delimiter listing, paginated), and an object's content is its
7//! value, fetched on first read and cached. Under composition
8//! (`qua` wraps object stores by default), a bucket of JSON, CSV,
9//! or source files is directly queryable: the object is a leaf,
10//! its parsed content the subtree — grafting is the point of this
11//! adapter.
12//!
13//! **Targets**:
14//! - `gs://BUCKET[/PREFIX]` — GCS, JSON API. Public buckets work
15//!   anonymously; private ones authenticate like the other GCP
16//!   drivers (`QUARB_GCP_TOKEN`, else `gcloud auth
17//!   print-access-token`, `?account=EMAIL` to pick the account —
18//!   set `?auth=1` to force a token for non-public buckets).
19//! - `s3://BUCKET[/PREFIX][?region=R]` — S3, ListObjectsV2.
20//!   **Anonymous only in v1**: public buckets read without
21//!   credentials; SigV4 request signing is a recorded extension,
22//!   so private S3 buckets refuse honestly rather than
23//!   half-work.
24//!
25//! Metadata: `::;size`, `::;updated` on objects; traits
26//! `<object>` / `<prefix>`. Read-only, as always.
27
28use quarb::{AstAdapter, NodeId, Value};
29use std::cell::RefCell;
30
31/// An error connecting to a bucket.
32#[derive(Debug, thiserror::Error)]
33pub enum ObjstoreError {
34    #[error("objstore: {0}")]
35    Http(String),
36    #[error("objstore target: {0} (expected gs://BUCKET[/PREFIX] or s3://BUCKET[/PREFIX])")]
37    Target(String),
38}
39
40enum Backend {
41    Gcs { token: Option<String> },
42    S3 { region: String },
43}
44
45struct Node {
46    /// Full key prefix (dirs end without `/`; root is "").
47    key: String,
48    name: Option<String>,
49    parent: Option<NodeId>,
50    is_object: bool,
51    size: Option<i64>,
52    updated: Option<String>,
53    children: RefCell<Option<Vec<NodeId>>>,
54    content: RefCell<Option<String>>,
55}
56
57/// A bucket (or prefix of one), exposed as an arbor.
58pub struct ObjstoreAdapter {
59    backend: Backend,
60    bucket: String,
61    /// The target's prefix, "" for the whole bucket.
62    base: String,
63    nodes: RefCell<Vec<Node>>,
64}
65
66fn gcp_token(account: Option<&str>) -> Option<String> {
67    if let Ok(t) = std::env::var("QUARB_GCP_TOKEN")
68        && !t.trim().is_empty()
69    {
70        return Some(t.trim().to_string());
71    }
72    let mut cmd = std::process::Command::new("gcloud");
73    cmd.args(["auth", "print-access-token"]);
74    if let Some(a) = account {
75        cmd.arg(a);
76    }
77    let out = cmd.output().ok()?;
78    out.status
79        .success()
80        .then(|| String::from_utf8_lossy(&out.stdout).trim().to_string())
81}
82
83fn urlencode(s: &str) -> String {
84    let mut out = String::new();
85    for b in s.bytes() {
86        match b {
87            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
88                out.push(b as char)
89            }
90            other => out.push_str(&format!("%{other:02X}")),
91        }
92    }
93    out
94}
95
96/// One listing page: (dir prefixes, objects as (key, size, updated)).
97type Page = (Vec<String>, Vec<(String, Option<i64>, Option<String>)>);
98
99impl ObjstoreAdapter {
100    /// Connect to `gs://...` or `s3://...`; one listing probes the
101    /// bucket.
102    pub fn connect(target: &str) -> Result<Self, ObjstoreError> {
103        let (backend, rest) = if let Some(r) = target.strip_prefix("gs://") {
104            ("gs", r)
105        } else if let Some(r) = target.strip_prefix("s3://") {
106            ("s3", r)
107        } else {
108            return Err(ObjstoreError::Target(target.to_string()));
109        };
110        let (path, query) = match rest.split_once('?') {
111            Some((p, q)) => (p, Some(q)),
112            None => (rest, None),
113        };
114        let (bucket, prefix) = match path.split_once('/') {
115            Some((b, p)) => (b.to_string(), p.trim_end_matches('/').to_string()),
116            None => (path.to_string(), String::new()),
117        };
118        if bucket.is_empty() {
119            return Err(ObjstoreError::Target(target.to_string()));
120        }
121        let param = |k: &str| {
122            query.and_then(|q| {
123                q.split('&')
124                    .find_map(|kv| kv.strip_prefix(&format!("{k}=")).map(str::to_string))
125            })
126        };
127        let backend = if backend == "gs" {
128            let token = if param("auth").is_some() || param("account").is_some() {
129                gcp_token(param("account").as_deref())
130            } else {
131                None
132            };
133            Backend::Gcs { token }
134        } else {
135            Backend::S3 {
136                region: param("region").unwrap_or_else(|| "us-east-1".to_string()),
137            }
138        };
139        let adapter = ObjstoreAdapter {
140            backend,
141            bucket,
142            base: prefix.clone(),
143            nodes: RefCell::new(vec![Node {
144                key: prefix,
145                name: None,
146                parent: None,
147                is_object: false,
148                size: None,
149                updated: None,
150                children: RefCell::new(None),
151                content: RefCell::new(None),
152            }]),
153        };
154        adapter
155            .list(&adapter.nodes.borrow()[0].key.clone())
156            .map_err(|e| ObjstoreError::Http(format!("probing the bucket: {e}")))?;
157        Ok(adapter)
158    }
159
160    /// A human-readable locator: the object key below the base.
161    pub fn locator(&self, node: NodeId) -> String {
162        let key = &self.nodes.borrow()[node.0 as usize].key;
163        let rel = key.strip_prefix(&self.base).unwrap_or(key);
164        format!("/{}", rel.trim_start_matches('/'))
165    }
166
167    fn get(&self, url: &str) -> Result<String, String> {
168        let mut req = ureq::get(url);
169        if let Backend::Gcs { token: Some(t) } = &self.backend {
170            req = req.set("Authorization", &format!("Bearer {t}"));
171        }
172        req.call()
173            .map_err(|e| e.to_string())?
174            .into_string()
175            .map_err(|e| e.to_string())
176    }
177
178    /// One delimiter listing under `prefix`, following pages.
179    fn list(&self, prefix: &str) -> Result<Page, String> {
180        let dir = if prefix.is_empty() {
181            String::new()
182        } else {
183            format!("{prefix}/")
184        };
185        let mut prefixes = Vec::new();
186        let mut objects = Vec::new();
187        let mut page: Option<String> = None;
188        loop {
189            match &self.backend {
190                Backend::Gcs { .. } => {
191                    let mut url = format!(
192                        "https://storage.googleapis.com/storage/v1/b/{}/o?delimiter=/&prefix={}",
193                        self.bucket,
194                        urlencode(&dir)
195                    );
196                    if let Some(p) = &page {
197                        url.push_str(&format!("&pageToken={}", urlencode(p)));
198                    }
199                    let resp: serde_json::Value = serde_json::from_str(&self.get(&url)?)
200                        .map_err(|e| format!("listing: {e}"))?;
201                    if let Some(err) = resp.pointer("/error/message").and_then(|v| v.as_str()) {
202                        return Err(err.to_string());
203                    }
204                    if let Some(ps) = resp.pointer("/prefixes").and_then(|v| v.as_array()) {
205                        prefixes.extend(
206                            ps.iter()
207                                .filter_map(|p| p.as_str())
208                                .map(|p| p.trim_end_matches('/').to_string()),
209                        );
210                    }
211                    if let Some(items) = resp.pointer("/items").and_then(|v| v.as_array()) {
212                        for i in items {
213                            let Some(key) = i.pointer("/name").and_then(|v| v.as_str()) else {
214                                continue;
215                            };
216                            if key.ends_with('/') {
217                                continue; // zero-byte "directory" markers
218                            }
219                            objects.push((
220                                key.to_string(),
221                                i.pointer("/size")
222                                    .and_then(|v| v.as_str())
223                                    .and_then(|s| s.parse().ok()),
224                                i.pointer("/updated")
225                                    .and_then(|v| v.as_str())
226                                    .map(str::to_string),
227                            ));
228                        }
229                    }
230                    page = resp
231                        .pointer("/nextPageToken")
232                        .and_then(|v| v.as_str())
233                        .map(str::to_string);
234                }
235                Backend::S3 { region } => {
236                    let host = if region == "us-east-1" {
237                        format!("{}.s3.amazonaws.com", self.bucket)
238                    } else {
239                        format!("{}.s3.{region}.amazonaws.com", self.bucket)
240                    };
241                    let mut url = format!(
242                        "https://{host}/?list-type=2&delimiter=/&prefix={}",
243                        urlencode(&dir)
244                    );
245                    if let Some(p) = &page {
246                        url.push_str(&format!("&continuation-token={}", urlencode(p)));
247                    }
248                    let xml = self.get(&url)?;
249                    let (ps, os, next) = parse_s3_listing(&xml)?;
250                    prefixes.extend(ps);
251                    objects.extend(os);
252                    page = next;
253                }
254            }
255            if page.is_none() {
256                break;
257            }
258        }
259        Ok((prefixes, objects))
260    }
261
262    fn push(&self, node: Node) -> NodeId {
263        let mut nodes = self.nodes.borrow_mut();
264        let id = NodeId(nodes.len() as u64);
265        nodes.push(node);
266        id
267    }
268
269    /// An object's content, fetched once (text, lossily decoded).
270    fn content_of(&self, node: NodeId) -> Option<String> {
271        if let Some(c) = &*self.nodes.borrow()[node.0 as usize].content.borrow() {
272            return Some(c.clone());
273        }
274        let (key, is_object) = {
275            let nodes = self.nodes.borrow();
276            let n = &nodes[node.0 as usize];
277            (n.key.clone(), n.is_object)
278        };
279        if !is_object {
280            return None;
281        }
282        let url = match &self.backend {
283            Backend::Gcs { .. } => format!(
284                "https://storage.googleapis.com/storage/v1/b/{}/o/{}?alt=media",
285                self.bucket,
286                urlencode(&key)
287            ),
288            Backend::S3 { region } => {
289                let host = if region == "us-east-1" {
290                    format!("{}.s3.amazonaws.com", self.bucket)
291                } else {
292                    format!("{}.s3.{region}.amazonaws.com", self.bucket)
293                };
294                format!("https://{host}/{}", urlencode(&key).replace("%2F", "/"))
295            }
296        };
297        let text = self.get(&url).ok()?;
298        *self.nodes.borrow()[node.0 as usize].content.borrow_mut() = Some(text.clone());
299        Some(text)
300    }
301}
302
303/// Parse an S3 ListObjectsV2 response (streamed, no DOM).
304#[allow(clippy::type_complexity)]
305fn parse_s3_listing(
306    xml: &str,
307) -> Result<
308    (
309        Vec<String>,
310        Vec<(String, Option<i64>, Option<String>)>,
311        Option<String>,
312    ),
313    String,
314> {
315    use quick_xml::events::Event;
316    let mut reader = quick_xml::Reader::from_str(xml);
317    let mut prefixes = Vec::new();
318    let mut objects = Vec::new();
319    let mut next = None;
320    let mut path: Vec<String> = Vec::new();
321    let mut cur: (Option<String>, Option<i64>, Option<String>) = (None, None, None);
322    loop {
323        match reader.read_event().map_err(|e| format!("listing: {e}"))? {
324            Event::Start(e) => path.push(String::from_utf8_lossy(e.name().as_ref()).into_owned()),
325            Event::End(e) => {
326                let name = String::from_utf8_lossy(e.name().as_ref()).into_owned();
327                if name == "Contents"
328                    && let (Some(k), size, updated) = std::mem::take(&mut cur)
329                    && !k.ends_with('/')
330                {
331                    objects.push((k, size, updated));
332                }
333                path.pop();
334            }
335            Event::Text(t) => {
336                let text = t.xml_content().map_err(|e| e.to_string())?.into_owned();
337                match path.as_slice() {
338                    [.., a, b] if a == "CommonPrefixes" && b == "Prefix" => {
339                        prefixes.push(text.trim_end_matches('/').to_string());
340                    }
341                    [.., a, b] if a == "Contents" && b == "Key" => cur.0 = Some(text),
342                    [.., a, b] if a == "Contents" && b == "Size" => {
343                        cur.1 = text.parse().ok();
344                    }
345                    [.., a, b] if a == "Contents" && b == "LastModified" => {
346                        cur.2 = Some(text);
347                    }
348                    [.., b] if b == "NextContinuationToken" => next = Some(text),
349                    [.., a, b] if a == "ListBucketResult" && b == "NextContinuationToken" => {
350                        next = Some(text)
351                    }
352                    _ => {}
353                }
354            }
355            Event::Eof => break,
356            _ => {}
357        }
358    }
359    Ok((prefixes, objects, next))
360}
361
362impl AstAdapter for ObjstoreAdapter {
363    fn root(&self) -> NodeId {
364        NodeId(0)
365    }
366
367    fn children(&self, node: NodeId) -> Vec<NodeId> {
368        if let Some(c) = self.nodes.borrow()[node.0 as usize]
369            .children
370            .borrow()
371            .as_ref()
372        {
373            return c.clone();
374        }
375        let (key, is_object) = {
376            let nodes = self.nodes.borrow();
377            let n = &nodes[node.0 as usize];
378            (n.key.clone(), n.is_object)
379        };
380        if is_object {
381            return Vec::new();
382        }
383        let (prefixes, objects) = self.list(&key).unwrap_or_default();
384        let mut ids = Vec::new();
385        for p in prefixes {
386            let name = p.rsplit('/').next().unwrap_or(&p).to_string();
387            ids.push(self.push(Node {
388                key: p,
389                name: Some(name),
390                parent: Some(node),
391                is_object: false,
392                size: None,
393                updated: None,
394                children: RefCell::new(None),
395                content: RefCell::new(None),
396            }));
397        }
398        for (k, size, updated) in objects {
399            let name = k.rsplit('/').next().unwrap_or(&k).to_string();
400            ids.push(self.push(Node {
401                key: k,
402                name: Some(name),
403                parent: Some(node),
404                is_object: true,
405                size,
406                updated,
407                children: RefCell::new(None),
408                content: RefCell::new(None),
409            }));
410        }
411        *self.nodes.borrow()[node.0 as usize].children.borrow_mut() = Some(ids.clone());
412        ids
413    }
414
415    fn name(&self, node: NodeId) -> Option<String> {
416        self.nodes.borrow()[node.0 as usize].name.clone()
417    }
418
419    fn parent(&self, node: NodeId) -> Option<NodeId> {
420        self.nodes.borrow()[node.0 as usize].parent
421    }
422
423    /// `<object>` / `<prefix>`.
424    fn traits(&self, node: NodeId) -> Vec<String> {
425        let nodes = self.nodes.borrow();
426        let n = &nodes[node.0 as usize];
427        if n.parent.is_none() {
428            return Vec::new();
429        }
430        vec![if n.is_object { "object" } else { "prefix" }.to_string()]
431    }
432
433    /// An object's content (fetched on first read, cached).
434    fn default_value(&self, node: NodeId) -> Option<Value> {
435        self.content_of(node).map(Value::Str)
436    }
437
438    /// `::;size`, `::;updated`, `::;key`.
439    fn metadata(&self, node: NodeId, key: &str) -> Option<Value> {
440        let nodes = self.nodes.borrow();
441        let n = &nodes[node.0 as usize];
442        match key {
443            "size" => n.size.map(Value::Int),
444            "updated" => n.updated.clone().map(Value::Str),
445            "key" => Some(Value::Str(n.key.clone())),
446            _ => None,
447        }
448    }
449}
quarb_objstore/lib.rs

quarb_objstore/
lib.rs
