polybase_client/
lib.rs

1//! Bounded synchronous client for the local Polybase HTTP API. No automatic write retry.
2use serde::{Deserialize, Serialize, de::DeserializeOwned};
3use serde_json::{Value, json};
4use std::{
5    io::{Read, Write},
6    net::{SocketAddr, TcpStream},
7    time::Duration,
8};
9
10#[derive(Debug)]
11pub enum Error {
12    Io(std::io::Error),
13    Json(serde_json::Error),
14    Protocol(&'static str),
15    Api { status: u16, body: Value },
16}
17impl std::fmt::Display for Error {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        write!(f, "{self:?}")
20    }
21}
22impl std::error::Error for Error {}
23impl From<std::io::Error> for Error {
24    fn from(e: std::io::Error) -> Self {
25        Self::Io(e)
26    }
27}
28impl From<serde_json::Error> for Error {
29    fn from(e: serde_json::Error) -> Self {
30        Self::Json(e)
31    }
32}
33pub type Result<T> = std::result::Result<T, Error>;
34
35#[derive(Clone, Debug, Serialize, Deserialize)]
36pub struct Query {
37    pub key_prefix: String,
38    pub predicate: String,
39    pub value: String,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub json: Option<Value>,
42}
43impl Query {
44    pub fn all(prefix: &[u8]) -> Self {
45        Self {
46            key_prefix: hex(prefix),
47            predicate: "any".into(),
48            value: String::new(),
49            json: None,
50        }
51    }
52    pub fn json(prefix: &[u8], predicate: Value) -> Self {
53        Self {
54            key_prefix: hex(prefix),
55            predicate: "json".into(),
56            value: String::new(),
57            json: Some(predicate),
58        }
59    }
60}
61#[derive(Clone, Debug, Deserialize)]
62pub struct Row {
63    pub key: String,
64    pub value: String,
65}
66#[derive(Clone, Debug, Deserialize)]
67pub struct Page {
68    pub position: String,
69    pub rows: Vec<Row>,
70    pub next: Option<String>,
71    pub explain: Value,
72}
73#[derive(Clone, Debug, Deserialize)]
74pub struct RevisionRead {
75    pub value: Option<String>,
76    pub revision: String,
77    pub position: String,
78}
79#[derive(Clone, Debug, Deserialize)]
80pub struct Receipt {
81    pub position: String,
82    pub wrote: bool,
83    pub durable: bool,
84}
85#[derive(Clone, Debug, Deserialize)]
86pub struct Session {
87    pub id: String,
88    pub instance: String,
89    pub token: String,
90    pub position: String,
91    pub rows: Vec<Row>,
92}
93/// A resumed subscription has an identity but does not return another snapshot.
94#[derive(Clone, Debug, Deserialize)]
95pub struct ResumedSession {
96    pub id: String,
97    pub instance: String,
98}
99
100pub trait SessionIdentity {
101    fn id(&self) -> &str;
102    fn instance(&self) -> &str;
103}
104impl SessionIdentity for Session {
105    fn id(&self) -> &str {
106        &self.id
107    }
108    fn instance(&self) -> &str {
109        &self.instance
110    }
111}
112impl SessionIdentity for ResumedSession {
113    fn id(&self) -> &str {
114        &self.id
115    }
116    fn instance(&self) -> &str {
117        &self.instance
118    }
119}
120
121/// One bounded current-head search; this does not register a subscription.
122#[derive(Clone, Debug, Serialize)]
123#[serde(tag = "kind", rename_all = "lowercase")]
124pub enum SearchQuery {
125    Json {
126        key_prefix: String,
127        path: Vec<String>,
128        value: Value,
129    },
130    Text {
131        key_prefix: String,
132        text: String,
133    },
134}
135
136#[derive(Clone, Debug, Deserialize)]
137pub struct SearchResult {
138    pub position: String,
139    pub indexed: bool,
140    pub rows: Vec<Row>,
141}
142#[derive(Clone, Debug, Deserialize)]
143pub struct Change {
144    pub key: String,
145    pub before: Option<String>,
146    pub after: Option<String>,
147}
148#[derive(Clone, Debug, Deserialize)]
149pub struct Commit {
150    pub position: String,
151    pub changes: Vec<Change>,
152}
153#[derive(Clone, Debug, Deserialize)]
154pub struct StreamPage {
155    pub commits: Vec<Commit>,
156    pub token: String,
157    pub position: String,
158    pub head: String,
159    pub floor: String,
160}
161
162#[derive(Clone)]
163pub struct Client {
164    address: SocketAddr,
165}
166impl Client {
167    pub fn new(address: SocketAddr) -> Result<Self> {
168        if !address.is_ipv4() || !address.ip().is_loopback() || address.port() == 0 {
169            return Err(Error::Protocol(
170                "only IPv4 loopback endpoints are supported",
171            ));
172        }
173        Ok(Self { address })
174    }
175    /// A transport error after sending a write leaves its outcome uncertain.
176    pub fn call<T: DeserializeOwned>(&self, route: &str, body: &impl Serialize) -> Result<T> {
177        if !route.bytes().all(|b| b.is_ascii_lowercase() || b == b'-') {
178            return Err(Error::Protocol("invalid route"));
179        }
180        let body = serde_json::to_vec(body)?;
181        if body.len() > 128 * 1024 {
182            return Err(Error::Protocol("request exceeds 128 KiB"));
183        }
184        let mut socket = TcpStream::connect_timeout(&self.address, Duration::from_secs(5))?;
185        socket.set_read_timeout(Some(Duration::from_secs(5)))?;
186        socket.set_write_timeout(Some(Duration::from_secs(5)))?;
187        write!(
188            socket,
189            "POST /api/{route} HTTP/1.1\r\nHost: {}\r\nOrigin: http://{}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
190            self.address,
191            self.address,
192            body.len()
193        )?;
194        socket.write_all(&body)?;
195        let mut response = Vec::new();
196        socket
197            .take(8 * 1024 * 1024 + 8193)
198            .read_to_end(&mut response)?;
199        if response.len() > 8 * 1024 * 1024 + 8192 {
200            return Err(Error::Protocol("response exceeds bound"));
201        }
202        let end = response
203            .windows(4)
204            .position(|w| w == b"\r\n\r\n")
205            .ok_or(Error::Protocol("missing HTTP header"))?;
206        let header = std::str::from_utf8(&response[..end])
207            .map_err(|_| Error::Protocol("invalid HTTP header"))?;
208        let status = header
209            .lines()
210            .next()
211            .and_then(|s| s.strip_prefix("HTTP/1.1 "))
212            .and_then(|s| s.split(' ').next())
213            .and_then(|s| s.parse::<u16>().ok())
214            .ok_or(Error::Protocol("invalid HTTP status"))?;
215        let payload = &response[end + 4..];
216        let length = header.lines().find_map(|s| {
217            s.to_ascii_lowercase()
218                .strip_prefix("content-length:")
219                .and_then(|s| s.trim().parse::<usize>().ok())
220        });
221        if length != Some(payload.len()) {
222            return Err(Error::Protocol("incomplete HTTP response"));
223        }
224        if status != 200 {
225            return Err(Error::Api {
226                status,
227                body: serde_json::from_slice(payload).unwrap_or_else(|_| {
228                    Value::String(String::from_utf8_lossy(payload).into_owned())
229                }),
230            });
231        }
232        Ok(serde_json::from_slice(payload)?)
233    }
234    pub fn page(&self, query: &Query, cursor: Option<&str>, limit: usize) -> Result<Page> {
235        self.page_with_scan(query, cursor, limit, false)
236    }
237    pub fn page_with_scan(
238        &self,
239        query: &Query,
240        cursor: Option<&str>,
241        limit: usize,
242        force_scan: bool,
243    ) -> Result<Page> {
244        let mut body = json!({"query":query,"cursor":cursor,"limit":limit});
245        if force_scan {
246            body["force_scan"] = json!(true);
247        }
248        self.call("query-page", &body)
249    }
250    pub fn read_revision(&self, key: &[u8]) -> Result<RevisionRead> {
251        self.call("read-revision", &json!({"key":hex(key)}))
252    }
253    pub fn write_revision(&self, revision: &str, value: Option<&[u8]>) -> Result<Receipt> {
254        self.call(
255            "write-revision",
256            &json!({"revision":revision,"value":value.map(hex)}),
257        )
258    }
259    pub fn subscribe(&self, query: &Query) -> Result<Session> {
260        self.call("query", query)
261    }
262    pub fn resume(&self, token: &str) -> Result<ResumedSession> {
263        self.call("resume", &json!({"token":token}))
264    }
265    /// Supply the last *applied* token; persist rows and this token atomically in the application.
266    pub fn poll(&self, session: &impl SessionIdentity, applied_token: &str) -> Result<StreamPage> {
267        self.call(
268            "poll",
269            &json!({"id":session.id(),"instance":session.instance(),"token":applied_token}),
270        )
271    }
272    pub fn unsubscribe(&self, session: &impl SessionIdentity, token: &str) -> Result<Value> {
273        self.call(
274            "unsubscribe",
275            &json!({"id":session.id(),"instance":session.instance(),"token":token}),
276        )
277    }
278    pub fn snapshot_at(&self, position: u64) -> Result<Value> {
279        self.call("snapshot-at", &json!({"position":position.to_string()}))
280    }
281    pub fn backup(&self, directory: &str) -> Result<Value> {
282        self.call("backup", &json!({"directory":directory}))
283    }
284    pub fn verify_backup(&self, directory: &str) -> Result<Value> {
285        self.call("verify-backup", &json!({"directory":directory}))
286    }
287    /// This route takes prefix bytes as a JSON array, unlike hex key_prefix fields.
288    pub fn create_index(&self, name: &str, prefix: &[u8], path: &[&str]) -> Result<Value> {
289        self.call(
290            "index-create",
291            &json!({"name":name,"prefix":prefix,"path":path}),
292        )
293    }
294    pub fn indexes(&self) -> Result<Value> {
295        self.call("index-list", &json!({}))
296    }
297    pub fn search_index(&self, name: &str, value: Value) -> Result<SearchResult> {
298        self.call("index-search", &json!({"name":name,"value":value}))
299    }
300    pub fn drop_index(&self, name: &str) -> Result<Value> {
301        self.call("index-drop", &json!({"name":name}))
302    }
303    pub fn compaction_advice(&self, retain_commits: u64) -> Result<Value> {
304        self.call(
305            "compaction-advice",
306            &json!({"retain_commits":retain_commits.to_string()}),
307        )
308    }
309    pub fn search(&self, query: &SearchQuery) -> Result<SearchResult> {
310        self.call("search", query)
311    }
312}
313pub fn hex(bytes: &[u8]) -> String {
314    use std::fmt::Write;
315    bytes
316        .iter()
317        .fold(String::with_capacity(bytes.len() * 2), |mut s, b| {
318            write!(s, "{b:02x}").expect("String write");
319            s
320        })
321}
322pub fn unhex(s: &str) -> Result<Vec<u8>> {
323    if !s.len().is_multiple_of(2)
324        || !s
325            .bytes()
326            .all(|c| c.is_ascii_digit() || (b'a'..=b'f').contains(&c))
327    {
328        return Err(Error::Protocol("invalid hex"));
329    }
330    s.as_bytes()
331        .as_chunks::<2>()
332        .0
333        .iter()
334        .map(|b| {
335            u8::from_str_radix(std::str::from_utf8(b).expect("ASCII"), 16)
336                .map_err(|_| Error::Protocol("invalid hex"))
337        })
338        .collect()
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344    use std::{
345        io::{BufRead, BufReader},
346        net::TcpListener,
347        thread,
348    };
349
350    fn response(status: u16, body: &[u8]) -> Result<Value> {
351        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
352        let address = listener.local_addr().unwrap();
353        let body = body.to_vec();
354        let worker = thread::spawn(move || {
355            let (mut socket, _) = listener.accept().unwrap();
356            socket
357                .set_read_timeout(Some(Duration::from_secs(5)))
358                .unwrap();
359            let mut reader = BufReader::new(&mut socket);
360            let mut length = 0;
361            loop {
362                let mut line = String::new();
363                assert!(reader.read_line(&mut line).unwrap() > 0);
364                if line == "\r\n" {
365                    break;
366                }
367                if let Some(value) = line.to_ascii_lowercase().strip_prefix("content-length:") {
368                    length = value.trim().parse::<usize>().unwrap();
369                }
370            }
371            assert!(length <= 128 * 1024);
372            reader.read_exact(&mut vec![0; length]).unwrap();
373            write!(
374                reader.get_mut(),
375                "HTTP/1.1 {status} Test\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
376                body.len()
377            )
378            .unwrap();
379            reader.get_mut().write_all(&body).unwrap();
380        });
381        let result = Client::new(address)
382            .unwrap()
383            .call("read", &json!({"key":""}));
384        worker.join().unwrap();
385        result
386    }
387
388    #[test]
389    fn non_json_errors_keep_status_and_text_while_success_requires_json() {
390        for (status, bytes, text) in [
391            (500, b"upstream panic".as_slice(), "upstream panic"),
392            (503, b"".as_slice(), ""),
393            (502, b"bad \xff".as_slice(), "bad \u{fffd}"),
394        ] {
395            assert!(
396                matches!(response(status, bytes), Err(Error::Api { status: actual, body }) if actual == status && body == json!(text))
397            );
398        }
399        assert!(
400            matches!(response(409, br#"{"error":"conflict"}"#), Err(Error::Api { status: 409, body }) if body == json!({"error":"conflict"}))
401        );
402        assert!(matches!(
403            response(200, b"invalid success"),
404            Err(Error::Json(_))
405        ));
406    }
407}
polybase_client/lib.rs

polybase_client/
lib.rs
