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#[derive(Clone, Debug, Deserialize)]
94pub struct Change {
95    pub key: String,
96    pub before: Option<String>,
97    pub after: Option<String>,
98}
99#[derive(Clone, Debug, Deserialize)]
100pub struct Commit {
101    pub position: String,
102    pub changes: Vec<Change>,
103}
104#[derive(Clone, Debug, Deserialize)]
105pub struct StreamPage {
106    pub commits: Vec<Commit>,
107    pub token: String,
108    pub position: String,
109    pub head: String,
110    pub floor: String,
111}
112
113#[derive(Clone)]
114pub struct Client {
115    address: SocketAddr,
116}
117impl Client {
118    pub fn new(address: SocketAddr) -> Result<Self> {
119        if !address.is_ipv4() || !address.ip().is_loopback() || address.port() == 0 {
120            return Err(Error::Protocol(
121                "only IPv4 loopback endpoints are supported",
122            ));
123        }
124        Ok(Self { address })
125    }
126    /// A transport error after sending a write leaves its outcome uncertain.
127    pub fn call<T: DeserializeOwned>(&self, route: &str, body: &impl Serialize) -> Result<T> {
128        if !route.bytes().all(|b| b.is_ascii_lowercase() || b == b'-') {
129            return Err(Error::Protocol("invalid route"));
130        }
131        let body = serde_json::to_vec(body)?;
132        if body.len() > 128 * 1024 {
133            return Err(Error::Protocol("request exceeds 128 KiB"));
134        }
135        let mut socket = TcpStream::connect_timeout(&self.address, Duration::from_secs(5))?;
136        socket.set_read_timeout(Some(Duration::from_secs(5)))?;
137        socket.set_write_timeout(Some(Duration::from_secs(5)))?;
138        write!(
139            socket,
140            "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",
141            self.address,
142            self.address,
143            body.len()
144        )?;
145        socket.write_all(&body)?;
146        let mut response = Vec::new();
147        socket
148            .take(8 * 1024 * 1024 + 8193)
149            .read_to_end(&mut response)?;
150        if response.len() > 8 * 1024 * 1024 + 8192 {
151            return Err(Error::Protocol("response exceeds bound"));
152        }
153        let end = response
154            .windows(4)
155            .position(|w| w == b"\r\n\r\n")
156            .ok_or(Error::Protocol("missing HTTP header"))?;
157        let header = std::str::from_utf8(&response[..end])
158            .map_err(|_| Error::Protocol("invalid HTTP header"))?;
159        let status = header
160            .lines()
161            .next()
162            .and_then(|s| s.strip_prefix("HTTP/1.1 "))
163            .and_then(|s| s.split(' ').next())
164            .and_then(|s| s.parse::<u16>().ok())
165            .ok_or(Error::Protocol("invalid HTTP status"))?;
166        let payload = &response[end + 4..];
167        let length = header.lines().find_map(|s| {
168            s.to_ascii_lowercase()
169                .strip_prefix("content-length:")
170                .and_then(|s| s.trim().parse::<usize>().ok())
171        });
172        if length != Some(payload.len()) {
173            return Err(Error::Protocol("incomplete HTTP response"));
174        }
175        if status != 200 {
176            return Err(Error::Api {
177                status,
178                body: serde_json::from_slice(payload)?,
179            });
180        }
181        Ok(serde_json::from_slice(payload)?)
182    }
183    pub fn page(&self, query: &Query, cursor: Option<&str>, limit: usize) -> Result<Page> {
184        self.page_with_scan(query, cursor, limit, false)
185    }
186    pub fn page_with_scan(
187        &self,
188        query: &Query,
189        cursor: Option<&str>,
190        limit: usize,
191        force_scan: bool,
192    ) -> Result<Page> {
193        let mut body = json!({"query":query,"cursor":cursor,"limit":limit});
194        if force_scan {
195            body["force_scan"] = json!(true);
196        }
197        self.call("query-page", &body)
198    }
199    pub fn read_revision(&self, key: &[u8]) -> Result<RevisionRead> {
200        self.call("read-revision", &json!({"key":hex(key)}))
201    }
202    pub fn write_revision(&self, revision: &str, value: Option<&[u8]>) -> Result<Receipt> {
203        self.call(
204            "write-revision",
205            &json!({"revision":revision,"value":value.map(hex)}),
206        )
207    }
208    pub fn subscribe(&self, query: &Query) -> Result<Session> {
209        self.call("query", query)
210    }
211    /// Supply the last *applied* token; persist rows and this token atomically in the application.
212    pub fn poll(&self, session: &Session, applied_token: &str) -> Result<StreamPage> {
213        self.call(
214            "poll",
215            &json!({"id":session.id,"instance":session.instance,"token":applied_token}),
216        )
217    }
218    pub fn unsubscribe(&self, session: &Session, token: &str) -> Result<Value> {
219        self.call(
220            "unsubscribe",
221            &json!({"id":session.id,"instance":session.instance,"token":token}),
222        )
223    }
224    pub fn snapshot_at(&self, position: u64) -> Result<Value> {
225        self.call("snapshot-at", &json!({"position":position.to_string()}))
226    }
227    pub fn backup(&self, directory: &str) -> Result<Value> {
228        self.call("backup", &json!({"directory":directory}))
229    }
230}
231pub fn hex(bytes: &[u8]) -> String {
232    use std::fmt::Write;
233    bytes
234        .iter()
235        .fold(String::with_capacity(bytes.len() * 2), |mut s, b| {
236            write!(s, "{b:02x}").expect("String write");
237            s
238        })
239}
240pub fn unhex(s: &str) -> Result<Vec<u8>> {
241    if !s.len().is_multiple_of(2)
242        || !s
243            .bytes()
244            .all(|c| c.is_ascii_digit() || (b'a'..=b'f').contains(&c))
245    {
246        return Err(Error::Protocol("invalid hex"));
247    }
248    s.as_bytes()
249        .as_chunks::<2>()
250        .0
251        .iter()
252        .map(|b| {
253            u8::from_str_radix(std::str::from_utf8(b).expect("ASCII"), 16)
254                .map_err(|_| Error::Protocol("invalid hex"))
255        })
256        .collect()
257}
polybase_client/lib.rs

polybase_client/
lib.rs
