polybase-client 0.8.2

//! Bounded synchronous client for the local Polybase HTTP API. No automatic write retry.
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::{Value, json};
use std::{
    io::{Read, Write},
    net::{SocketAddr, TcpStream},
    time::Duration,
};

#[derive(Debug)]
pub enum Error {
    Io(std::io::Error),
    Json(serde_json::Error),
    Protocol(&'static str),
    Api { status: u16, body: Value },
}
impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}
impl std::error::Error for Error {}
impl From<std::io::Error> for Error {
    fn from(e: std::io::Error) -> Self {
        Self::Io(e)
    }
}
impl From<serde_json::Error> for Error {
    fn from(e: serde_json::Error) -> Self {
        Self::Json(e)
    }
}
pub type Result<T> = std::result::Result<T, Error>;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Query {
    pub key_prefix: String,
    pub predicate: String,
    pub value: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub json: Option<Value>,
}
impl Query {
    pub fn all(prefix: &[u8]) -> Self {
        Self {
            key_prefix: hex(prefix),
            predicate: "any".into(),
            value: String::new(),
            json: None,
        }
    }
    pub fn json(prefix: &[u8], predicate: Value) -> Self {
        Self {
            key_prefix: hex(prefix),
            predicate: "json".into(),
            value: String::new(),
            json: Some(predicate),
        }
    }
}
#[derive(Clone, Debug, Deserialize)]
pub struct Row {
    pub key: String,
    pub value: String,
}
#[derive(Clone, Debug, Deserialize)]
pub struct Page {
    pub position: String,
    pub rows: Vec<Row>,
    pub next: Option<String>,
    pub explain: Value,
}
#[derive(Clone, Debug, Deserialize)]
pub struct RevisionRead {
    pub value: Option<String>,
    pub revision: String,
    pub position: String,
}
#[derive(Clone, Debug, Deserialize)]
pub struct Receipt {
    pub position: String,
    pub wrote: bool,
    pub durable: bool,
}
#[derive(Clone, Debug, Deserialize)]
pub struct Session {
    pub id: String,
    pub instance: String,
    pub token: String,
    pub position: String,
    pub rows: Vec<Row>,
}
#[derive(Clone, Debug, Deserialize)]
pub struct Change {
    pub key: String,
    pub before: Option<String>,
    pub after: Option<String>,
}
#[derive(Clone, Debug, Deserialize)]
pub struct Commit {
    pub position: String,
    pub changes: Vec<Change>,
}
#[derive(Clone, Debug, Deserialize)]
pub struct StreamPage {
    pub commits: Vec<Commit>,
    pub token: String,
    pub position: String,
    pub head: String,
    pub floor: String,
}

#[derive(Clone)]
pub struct Client {
    address: SocketAddr,
}
impl Client {
    pub fn new(address: SocketAddr) -> Result<Self> {
        if !address.is_ipv4() || !address.ip().is_loopback() || address.port() == 0 {
            return Err(Error::Protocol(
                "only IPv4 loopback endpoints are supported",
            ));
        }
        Ok(Self { address })
    }
    /// A transport error after sending a write leaves its outcome uncertain.
    pub fn call<T: DeserializeOwned>(&self, route: &str, body: &impl Serialize) -> Result<T> {
        if !route.bytes().all(|b| b.is_ascii_lowercase() || b == b'-') {
            return Err(Error::Protocol("invalid route"));
        }
        let body = serde_json::to_vec(body)?;
        if body.len() > 128 * 1024 {
            return Err(Error::Protocol("request exceeds 128 KiB"));
        }
        let mut socket = TcpStream::connect_timeout(&self.address, Duration::from_secs(5))?;
        socket.set_read_timeout(Some(Duration::from_secs(5)))?;
        socket.set_write_timeout(Some(Duration::from_secs(5)))?;
        write!(
            socket,
            "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",
            self.address,
            self.address,
            body.len()
        )?;
        socket.write_all(&body)?;
        let mut response = Vec::new();
        socket
            .take(8 * 1024 * 1024 + 8193)
            .read_to_end(&mut response)?;
        if response.len() > 8 * 1024 * 1024 + 8192 {
            return Err(Error::Protocol("response exceeds bound"));
        }
        let end = response
            .windows(4)
            .position(|w| w == b"\r\n\r\n")
            .ok_or(Error::Protocol("missing HTTP header"))?;
        let header = std::str::from_utf8(&response[..end])
            .map_err(|_| Error::Protocol("invalid HTTP header"))?;
        let status = header
            .lines()
            .next()
            .and_then(|s| s.strip_prefix("HTTP/1.1 "))
            .and_then(|s| s.split(' ').next())
            .and_then(|s| s.parse::<u16>().ok())
            .ok_or(Error::Protocol("invalid HTTP status"))?;
        let payload = &response[end + 4..];
        let length = header.lines().find_map(|s| {
            s.to_ascii_lowercase()
                .strip_prefix("content-length:")
                .and_then(|s| s.trim().parse::<usize>().ok())
        });
        if length != Some(payload.len()) {
            return Err(Error::Protocol("incomplete HTTP response"));
        }
        if status != 200 {
            return Err(Error::Api {
                status,
                body: serde_json::from_slice(payload)?,
            });
        }
        Ok(serde_json::from_slice(payload)?)
    }
    pub fn page(&self, query: &Query, cursor: Option<&str>, limit: usize) -> Result<Page> {
        self.page_with_scan(query, cursor, limit, false)
    }
    pub fn page_with_scan(
        &self,
        query: &Query,
        cursor: Option<&str>,
        limit: usize,
        force_scan: bool,
    ) -> Result<Page> {
        let mut body = json!({"query":query,"cursor":cursor,"limit":limit});
        if force_scan {
            body["force_scan"] = json!(true);
        }
        self.call("query-page", &body)
    }
    pub fn read_revision(&self, key: &[u8]) -> Result<RevisionRead> {
        self.call("read-revision", &json!({"key":hex(key)}))
    }
    pub fn write_revision(&self, revision: &str, value: Option<&[u8]>) -> Result<Receipt> {
        self.call(
            "write-revision",
            &json!({"revision":revision,"value":value.map(hex)}),
        )
    }
    pub fn subscribe(&self, query: &Query) -> Result<Session> {
        self.call("query", query)
    }
    /// Supply the last *applied* token; persist rows and this token atomically in the application.
    pub fn poll(&self, session: &Session, applied_token: &str) -> Result<StreamPage> {
        self.call(
            "poll",
            &json!({"id":session.id,"instance":session.instance,"token":applied_token}),
        )
    }
    pub fn unsubscribe(&self, session: &Session, token: &str) -> Result<Value> {
        self.call(
            "unsubscribe",
            &json!({"id":session.id,"instance":session.instance,"token":token}),
        )
    }
    pub fn snapshot_at(&self, position: u64) -> Result<Value> {
        self.call("snapshot-at", &json!({"position":position.to_string()}))
    }
    pub fn backup(&self, directory: &str) -> Result<Value> {
        self.call("backup", &json!({"directory":directory}))
    }
}
pub fn hex(bytes: &[u8]) -> String {
    use std::fmt::Write;
    bytes
        .iter()
        .fold(String::with_capacity(bytes.len() * 2), |mut s, b| {
            write!(s, "{b:02x}").expect("String write");
            s
        })
}
pub fn unhex(s: &str) -> Result<Vec<u8>> {
    if !s.len().is_multiple_of(2)
        || !s
            .bytes()
            .all(|c| c.is_ascii_digit() || (b'a'..=b'f').contains(&c))
    {
        return Err(Error::Protocol("invalid hex"));
    }
    s.as_bytes()
        .as_chunks::<2>()
        .0
        .iter()
        .map(|b| {
            u8::from_str_radix(std::str::from_utf8(b).expect("ASCII"), 16)
                .map_err(|_| Error::Protocol("invalid hex"))
        })
        .collect()
}