use std::io::{BufReader, BufWriter, Read, Stdout};
use std::path::{Path, PathBuf};
use arrow::array::RecordBatch;
use arrow::datatypes::SchemaRef;
use serde_json::{Map, Value};
use crate::frame::{self, Frame, FrameError};
use crate::msg::{ConnectionField, Request, Response, ResultSetMeta};
use crate::options::{self, OptionField};
use crate::{CAP_VALIDATE, Kind, MAX_VERSION, MIN_VERSION, PluginId};
pub type Error = Box<dyn std::error::Error + Send + Sync>;
pub type Result<T> = std::result::Result<T, Error>;
#[derive(Debug, Clone)]
pub struct About {
pub name: &'static str,
pub version: &'static str,
pub capabilities: &'static [&'static str],
}
impl About {
pub fn new(name: &'static str, version: &'static str) -> About {
let version = if version == "0.0.0" { "unreleased" } else { version };
About {
name,
version,
capabilities: &[],
}
}
pub fn capabilities(mut self, caps: &'static [&'static str]) -> About {
self.capabilities = caps;
self
}
}
pub trait ResultSink {
fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()>;
fn begin(&mut self, schema: SchemaRef) -> Result<()>;
fn batch(&mut self, batch: RecordBatch) -> Result<bool>;
}
pub trait Source {
fn connection_fields(&self) -> Vec<ConnectionField> {
Vec::new()
}
fn identifier_quote(&self) -> Option<&'static str> {
None
}
fn open(&mut self, connection: &Map<String, Value>, read_only: bool) -> Result<()>;
fn execute(&mut self, sql: &str, row_limit: Option<u64>, out: &mut dyn ResultSink) -> Result<()>;
fn check(&mut self, _sql: &str) -> Result<()> {
Err("this source can't check statements".into())
}
fn load(&mut self, _name: &str, _data: &mut ResultSet<'_>) -> Result<Loaded> {
Err("this source can't load rows".into())
}
fn close(&mut self) {}
}
pub struct Loaded {
pub relation: String,
pub rows: u64,
pub warning: Option<String>,
}
pub struct ResultSet<'a> {
pub meta: ResultSetMeta,
pub schema: SchemaRef,
first: Option<Vec<RecordBatch>>,
input: &'a mut Input,
done: bool,
}
impl ResultSet<'_> {
pub fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
loop {
if let Some(first) = self.first.as_mut() {
if !first.is_empty() {
return Ok(Some(first.remove(0)));
}
self.first = None;
}
if self.done {
return Ok(None);
}
match self.input.read()? {
Frame::Arrow(ipc) => {
let (_, batches) = frame::decode_batches(&ipc)?;
self.first = Some(batches);
}
Frame::Json(v) => match serde_json::from_value::<Request>(v)? {
Request::ResultSetEnd {} => {
self.done = true;
return Ok(None);
}
other => return Err(format!("expected result data, got {other:?}").into()),
},
}
}
}
fn drain(&mut self) -> Result<()> {
while self.next_batch()?.is_some() {}
Ok(())
}
}
pub struct WriteRequest {
pub path: String,
pub format: String,
pub options: Map<String, Value>,
pub result_sets: Vec<ResultSetMeta>,
pub template: Option<Value>,
}
pub struct ResultSets<'a> {
metas: std::vec::IntoIter<ResultSetMeta>,
input: &'a mut Input,
warnings: Vec<String>,
}
impl ResultSets<'_> {
pub fn warn(&mut self, message: impl Into<String>) {
self.warnings.push(message.into());
}
pub fn next_set(&mut self) -> Result<Option<ResultSet<'_>>> {
let Some(meta) = self.metas.next() else {
return Ok(None);
};
let ipc = match self.input.read()? {
Frame::Arrow(ipc) => ipc,
Frame::Json(v) => {
return Err(format!(
"expected an Arrow frame starting result set `{}`, got {v}",
meta.name
)
.into());
}
};
let (schema, batches) = frame::decode_batches(&ipc)?;
Ok(Some(ResultSet {
meta,
schema,
first: Some(batches),
input: &mut *self.input,
done: false,
}))
}
}
pub trait Format {
fn options(&self) -> Vec<OptionField> {
Vec::new()
}
fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
Vec::new()
}
fn write(&mut self, req: &WriteRequest, sets: &mut ResultSets<'_>) -> Result<Vec<String>>;
}
pub struct Delivery {
pub files: Vec<DeliveryFile>,
pub connection: Map<String, Value>,
pub options: Map<String, Value>,
}
pub struct DeliveryFile {
pub local: PathBuf,
pub remote: Option<String>,
}
pub trait Destination {
fn connection_fields(&self) -> Vec<ConnectionField> {
Vec::new()
}
fn options(&self) -> Vec<OptionField> {
Vec::new()
}
fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
Vec::new()
}
fn deliver(
&mut self,
_local: &Path,
_remote: Option<&str>,
_connection: &Map<String, Value>,
) -> Result<String> {
Err("this destination doesn't implement `deliver`".into())
}
fn deliver_files(&mut self, d: &Delivery) -> Result<String> {
match d.files.as_slice() {
[f] => self.deliver(&f.local, f.remote.as_deref(), &d.connection),
_ => Err("this destination takes one file per delivery".into()),
}
}
}
pub struct Input {
r: BufReader<Box<dyn Read>>,
}
impl Input {
fn read(&mut self) -> Result<Frame> {
frame::read_frame(&mut self.r).map_err(|e| -> Error {
match e {
FrameError::Eof => "core closed the connection".into(),
e => e.to_string().into(),
}
})
}
}
struct Output {
w: BufWriter<Stdout>,
}
impl Output {
fn send(&mut self, r: &Response) {
if frame::write_json(&mut self.w, r).is_err() {
std::process::exit(1);
}
}
fn batch(&mut self, b: &RecordBatch) -> Result<()> {
let ipc = frame::encode_batch(b)?;
frame::write_arrow(&mut self.w, &ipc).map_err(|_| -> Error { "core closed the connection".into() })
}
}
enum Handler<'a> {
Source(&'a mut dyn Source),
Format(&'a mut dyn Format),
Destination(&'a mut dyn Destination),
}
impl Handler<'_> {
fn kind(&self) -> Kind {
match self {
Handler::Source(_) => Kind::Source,
Handler::Format(_) => Kind::Format,
Handler::Destination(_) => Kind::Destination,
}
}
fn option_fields(&self) -> Vec<OptionField> {
match self {
Handler::Source(_) => Vec::new(),
Handler::Format(f) => f.options(),
Handler::Destination(d) => d.options(),
}
}
fn check(&self, name: &str, o: &Map<String, Value>) -> Vec<String> {
let errs = options::check(self.kind(), name, &self.option_fields(), o);
if !errs.is_empty() {
return errs;
}
match self {
Handler::Source(_) => Vec::new(),
Handler::Format(f) => f.validate(o),
Handler::Destination(d) => d.validate(o),
}
}
fn checked(&self, name: &str, o: &Map<String, Value>) -> Result<()> {
let errs = self.check(name, o);
if errs.is_empty() {
Ok(())
} else {
Err(errs.join("; ").into())
}
}
}
pub fn serve_source(about: About, s: impl Source + 'static) -> ! {
serve_package(vec![Plugin::Source(about, Box::new(s))])
}
pub fn serve_format(about: About, f: impl Format + 'static) -> ! {
serve_package(vec![Plugin::Format(about, Box::new(f))])
}
pub fn serve_destination(about: About, d: impl Destination + 'static) -> ! {
serve_package(vec![Plugin::Destination(about, Box::new(d))])
}
pub enum Plugin {
Source(About, Box<dyn Source>),
Format(About, Box<dyn Format>),
Destination(About, Box<dyn Destination>),
}
impl Plugin {
fn id(&self) -> PluginId {
let (kind, about) = match self {
Plugin::Source(a, _) => (Kind::Source, a),
Plugin::Format(a, _) => (Kind::Format, a),
Plugin::Destination(a, _) => (Kind::Destination, a),
};
PluginId::new(kind, about.name)
}
fn handler(&mut self) -> (&About, Handler<'_>) {
match self {
Plugin::Source(a, s) => (a, Handler::Source(s.as_mut())),
Plugin::Format(a, f) => (a, Handler::Format(f.as_mut())),
Plugin::Destination(a, d) => (a, Handler::Destination(d.as_mut())),
}
}
}
pub fn serve_package(mut plugins: Vec<Plugin>) -> ! {
assert!(!plugins.is_empty(), "a package serves at least one plugin");
let mut input = Input {
r: BufReader::new(Box::new(std::io::stdin())),
};
let mut out = Output {
w: BufWriter::new(std::io::stdout()),
};
let provides: Vec<PluginId> = plugins.iter().map(Plugin::id).collect();
let chosen = loop {
let req = match next_request(&mut input, &mut out) {
Some(r) => r,
None => continue,
};
let Request::Hello {
min_version,
max_version,
plugin,
..
} = req
else {
out.send(&Response::Error {
message: "the first request must be `hello`".into(),
});
continue;
};
let Some(chosen) = negotiate((min_version, max_version), (MIN_VERSION, MAX_VERSION)) else {
out.send(&Response::VersionMismatch {
min_version: MIN_VERSION,
max_version: MAX_VERSION,
});
std::process::exit(1);
};
let i = match &plugin {
None => 0,
Some(want) => match provides.iter().position(|p| p == want) {
Some(i) => i,
None => {
let list: Vec<String> = provides.iter().map(|p| p.to_string()).collect();
out.send(&Response::Error {
message: format!("this executable provides {}, not {want}", list.join(", ")),
});
std::process::exit(1);
}
},
};
let (about, h) = plugins[i].handler();
out.send(&Response::Hello {
protocol_version: chosen,
kind: h.kind(),
name: about.name.to_string(),
version: about.version.to_string(),
capabilities: about
.capabilities
.iter()
.chain(std::iter::once(&CAP_VALIDATE))
.map(|c| c.to_string())
.collect(),
provides: if provides.len() > 1 {
provides.clone()
} else {
Vec::new()
},
});
break i;
};
let (about, h) = plugins[chosen].handler();
serve(about.name, h, input, out)
}
fn next_request(input: &mut Input, out: &mut Output) -> Option<Request> {
let frame = match frame::read_frame(&mut input.r) {
Ok(f) => f,
Err(FrameError::Eof) => std::process::exit(0),
Err(e) => {
eprintln!("{e}");
out.send(&Response::Error {
message: e.to_string(),
});
std::process::exit(2);
}
};
match frame {
Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
Ok(r) => Some(r),
Err(_) => {
let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
out.send(&Response::Error {
message: format!("unsupported request `{t}`"),
});
None
}
},
Frame::Arrow(_) => {
out.send(&Response::Error {
message: "unexpected Arrow frame".into(),
});
None
}
}
}
fn serve(name: &str, mut h: Handler<'_>, mut input: Input, mut out: Output) -> ! {
loop {
let frame = match frame::read_frame(&mut input.r) {
Ok(f) => f,
Err(FrameError::Eof) => {
if let Handler::Source(src) = &mut h {
src.close();
}
std::process::exit(0)
}
Err(e) => {
eprintln!("{e}");
out.send(&Response::Error {
message: e.to_string(),
});
std::process::exit(2);
}
};
let req = match frame {
Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
Ok(r) => r,
Err(_) => {
let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
out.send(&Response::Error {
message: format!("unsupported request `{t}`"),
});
continue;
}
},
Frame::Arrow(_) => {
out.send(&Response::Error {
message: "unexpected Arrow frame".into(),
});
continue;
}
};
if let Request::Hello { .. } = req {
out.send(&Response::Error {
message: "`hello` was already answered".into(),
});
continue;
}
if let Request::Close {} = req {
if let Handler::Source(src) = &mut h {
src.close();
}
out.send(&Response::Ok {});
std::process::exit(0);
}
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
handle(&mut h, name, req, &mut input, &mut out)
}));
match result {
Ok(Ok(())) => {}
Ok(Err(e)) => out.send(&Response::Error {
message: e.to_string(),
}),
Err(p) => {
let msg = p
.downcast_ref::<String>()
.cloned()
.or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()))
.unwrap_or_else(|| "unknown panic".into());
out.send(&Response::Error {
message: format!("plugin panicked: {msg}"),
});
std::process::exit(101);
}
}
}
}
fn handle(h: &mut Handler<'_>, name: &str, req: Request, input: &mut Input, out: &mut Output) -> Result<()> {
if let Request::Describe {} = req {
let connection_fields = match h {
Handler::Source(s) => s.connection_fields(),
Handler::Destination(d) => d.connection_fields(),
Handler::Format(_) => Vec::new(),
};
let identifier_quote = match h {
Handler::Source(s) => s.identifier_quote().map(str::to_string),
_ => None,
};
out.send(&Response::Describe {
connection_fields,
option_fields: h.option_fields(),
identifier_quote,
});
return Ok(());
}
if let Request::Validate { options } = req {
out.send(&Response::Validated {
errors: h.check(name, &options),
});
return Ok(());
}
let checked = match &req {
Request::Write { options, .. } | Request::Deliver { options, .. } => h.checked(name, options),
_ => Ok(()),
};
match (h, req) {
(
Handler::Source(s),
Request::Open {
connection,
read_only,
},
) => {
s.open(&connection, read_only)?;
out.send(&Response::Ok {});
}
(Handler::Source(s), Request::Load { name }) => {
let ipc = match input.read()? {
Frame::Arrow(ipc) => ipc,
Frame::Json(v) => return Err(format!("expected the rows to load, got {v}").into()),
};
let (schema, batches) = frame::decode_batches(&ipc)?;
let mut data = ResultSet {
meta: ResultSetMeta {
name: name.clone(),
query: name.clone(),
result_index: 1,
anchor: None,
header: None,
columns: Default::default(),
},
schema,
first: Some(batches),
input,
done: false,
};
let r = s.load(&name, &mut data);
data.drain()?;
let loaded = r?;
out.send(&Response::Loaded {
relation: loaded.relation,
rows: loaded.rows,
warning: loaded.warning,
});
}
(Handler::Source(s), Request::Check { sql }) => {
s.check(&sql)?;
out.send(&Response::Ok {});
}
(Handler::Source(s), Request::Execute { sql, row_limit }) => {
let mut sink = FrameSink {
out,
row_limit,
rows: 0,
state: SinkState::Idle,
schema: None,
sent_any: false,
};
let r = s.execute(&sql, row_limit, &mut sink);
match (r, sink.state) {
(Err(e), _) => return Err(e),
(Ok(()), SinkState::Idle) => {
return Err("the source produced neither a result nor `no_result`".into());
}
(Ok(()), SinkState::NoResult) => {}
(Ok(()), SinkState::Result) => {
if !sink.sent_any {
let schema = sink.schema.clone().unwrap();
sink.out.batch(&RecordBatch::new_empty(schema))?;
}
let rows = sink.rows;
sink.out.send(&Response::ResultEnd { rows });
}
}
}
(
Handler::Format(f),
Request::Write {
path,
format,
options,
result_sets,
template,
},
) => {
let req = WriteRequest {
path,
format,
options,
result_sets: result_sets.clone(),
template,
};
let mut sets = ResultSets {
metas: result_sets.into_iter(),
input,
warnings: Vec::new(),
};
let written = checked.and_then(|()| f.write(&req, &mut sets));
if let Err(e) = &written {
out.send(&Response::Error {
message: e.to_string(),
});
}
loop {
match sets.input.read()? {
Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("finish") => break,
Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("result_set_end") => {}
Frame::Arrow(_) => {}
other => return Err(format!("expected `finish`, got {other:?}").into()),
}
}
if let Ok(files) = written {
out.send(&Response::Written {
files,
warnings: sets.warnings,
});
}
}
(
Handler::Destination(d),
Request::Deliver {
local_path,
remote_path,
files,
connection,
options,
},
) => {
checked?;
let files = match (local_path, files.is_empty()) {
(Some(local), true) => vec![DeliveryFile {
local: PathBuf::from(local),
remote: remote_path,
}],
(None, false) => files
.into_iter()
.map(|f| DeliveryFile {
local: PathBuf::from(f.local_path),
remote: f.remote_path,
})
.collect(),
_ => return Err("`deliver` needs exactly one of `local_path` or `files`".into()),
};
let location = d.deliver_files(&Delivery {
files,
connection,
options,
})?;
out.send(&Response::Delivered { location });
}
(h, req) => {
let t = serde_json::to_value(&req)
.ok()
.and_then(|v| v.get("type").cloned())
.unwrap_or_default();
return Err(format!("a {} plugin doesn't handle {t} requests", h.kind()).into());
}
}
Ok(())
}
#[derive(PartialEq)]
enum SinkState {
Idle,
NoResult,
Result,
}
struct FrameSink<'a> {
out: &'a mut Output,
row_limit: Option<u64>,
rows: u64,
state: SinkState,
schema: Option<SchemaRef>,
sent_any: bool,
}
impl ResultSink for FrameSink<'_> {
fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()> {
if self.state != SinkState::Idle {
return Err("`no_result` after the result started".into());
}
self.state = SinkState::NoResult;
self.out.send(&Response::NoResult { rows_affected });
Ok(())
}
fn begin(&mut self, schema: SchemaRef) -> Result<()> {
if self.state != SinkState::Idle {
return Err("a statement can only produce one result".into());
}
self.state = SinkState::Result;
self.out.send(&Response::Result {
columns: schema.fields().iter().map(|f| f.name().clone()).collect(),
});
self.schema = Some(schema);
Ok(())
}
fn batch(&mut self, b: RecordBatch) -> Result<bool> {
if self.state != SinkState::Result {
return Err("`batch` before `begin`".into());
}
if self.row_limit.is_some_and(|l| self.rows >= l) {
return Ok(false);
}
let b = match self.row_limit {
Some(l) if self.rows + b.num_rows() as u64 > l => b.slice(0, (l - self.rows) as usize),
_ => b,
};
if b.num_rows() > 0 || !self.sent_any {
self.rows += b.num_rows() as u64;
self.out.batch(&b)?;
self.sent_any = true;
}
Ok(self.row_limit.is_none_or(|l| self.rows < l))
}
}
fn negotiate(core: (u32, u32), ours: (u32, u32)) -> Option<u32> {
let (lo, hi) = (core.0.max(ours.0), core.1.min(ours.1));
(lo <= hi).then_some(hi)
}
pub fn conn_str<'a>(c: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
c.get(key).and_then(Value::as_str).filter(|s| !s.is_empty())
}
pub fn conn_required<'a>(c: &'a Map<String, Value>, key: &str) -> Result<&'a str> {
conn_str(c, key).ok_or_else(|| format!("the profile output needs a `{key}` field").into())
}
pub fn conn_bool(c: &Map<String, Value>, key: &str) -> Option<bool> {
match c.get(key)? {
Value::Bool(b) => Some(*b),
Value::String(s) => Some(matches!(s.as_str(), "true" | "yes" | "1")),
_ => None,
}
}