use std::collections::VecDeque;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::process::{Child, ChildStdin, Command, ExitStatus, Stdio};
use std::sync::mpsc::{Receiver, RecvTimeoutError, channel};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use arrow::array::RecordBatch;
use arrow::datatypes::SchemaRef;
use serde_json::{Map, Value};
use crate::frame::{self, Frame, FrameError};
use crate::msg::{ConnectionField, DeliveryFile, Request, Response, ResultSetMeta};
use crate::options::OptionField;
use crate::{Kind, MAX_VERSION, MIN_VERSION, PluginId, parse_executable_name};
pub type LogSink = Arc<dyn Fn(&str, &str) + Send + Sync>;
pub fn stderr_log() -> LogSink {
Arc::new(|plugin, line| eprintln!("[{plugin}] {line}"))
}
pub const DEFAULT_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);
pub const INTERACTIVE_ENV: &str = "DRE_INTERACTIVE";
const STDERR_TAIL: usize = 20;
#[derive(Debug, Clone, PartialEq)]
pub struct PluginInfo {
pub protocol_version: u32,
pub kind: Kind,
pub name: String,
pub version: String,
pub capabilities: Vec<String>,
pub provides: Vec<PluginId>,
}
#[derive(Debug)]
pub enum HostError {
Spawn {
plugin: String,
error: std::io::Error,
},
Crashed {
plugin: String,
status: Option<ExitStatus>,
stderr: Vec<String>,
},
Malformed {
plugin: String,
message: String,
},
Incompatible {
plugin: String,
core: (u32, u32),
plugin_range: (u32, u32),
},
Timeout {
plugin: String,
waiting_for: &'static str,
},
Unexpected {
plugin: String,
expected: &'static str,
got: String,
},
Plugin {
plugin: String,
message: String,
},
Arrow {
plugin: String,
message: String,
},
}
impl std::fmt::Display for HostError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
HostError::Spawn { plugin, error } => write!(f, "can't start plugin `{plugin}`: {error}"),
HostError::Crashed {
plugin,
status,
stderr,
} => {
match status {
Some(s) => write!(f, "plugin `{plugin}` exited unexpectedly ({s})")?,
None => write!(f, "plugin `{plugin}` closed its output unexpectedly")?,
}
if !stderr.is_empty() {
write!(f, "; its last log lines:\n {}", stderr.join("\n "))?;
}
Ok(())
}
HostError::Malformed { plugin, message } => {
write!(f, "plugin `{plugin}` broke the protocol: {message}")
}
HostError::Incompatible {
plugin,
core,
plugin_range,
} => write!(
f,
"plugin `{plugin}` speaks protocol versions {}..={}, but this DRE core speaks {}..={}; update {}",
plugin_range.0,
plugin_range.1,
core.0,
core.1,
if plugin_range.1 < core.0 {
"the plugin"
} else {
"DRE"
}
),
HostError::Timeout { plugin, waiting_for } => {
write!(
f,
"plugin `{plugin}` didn't answer in time (waiting for {waiting_for})"
)
}
HostError::Unexpected {
plugin,
expected,
got,
} => {
write!(f, "plugin `{plugin}` sent {got} where {expected} was expected")
}
HostError::Plugin { message, .. } => f.write_str(message),
HostError::Arrow { plugin, message } => {
write!(f, "plugin `{plugin}` sent invalid Arrow data: {message}")
}
}
}
}
impl std::error::Error for HostError {}
pub type Result<T> = std::result::Result<T, HostError>;
pub enum Incoming {
Json(Response),
Arrow(Vec<u8>),
}
impl std::fmt::Debug for Incoming {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Incoming::Json(r) => write!(f, "Json({r:?})"),
Incoming::Arrow(b) => write!(f, "Arrow frame ({} bytes)", b.len()),
}
}
}
#[derive(Debug)]
pub enum Execution {
NoResult { rows_affected: Option<u64> },
Result { schema: SchemaRef, rows: u64 },
}
pub struct PluginProcess {
label: String,
path: PathBuf,
child: Child,
stdin: Option<BufWriter<ChildStdin>>,
rx: Receiver<std::result::Result<Frame, FrameError>>,
stderr: Arc<Mutex<VecDeque<String>>>,
stderr_thread: Option<std::thread::JoinHandle<()>>,
info: Option<PluginInfo>,
early: Option<std::result::Result<Frame, FrameError>>,
serve: Option<PluginId>,
}
impl PluginProcess {
pub fn start(path: &Path, log: LogSink) -> Result<PluginProcess> {
Self::start_in(path, log, None)
}
pub fn start_in(path: &Path, log: LogSink, cwd: Option<&Path>) -> Result<PluginProcess> {
Self::start_for(path, None, log, cwd)
}
pub fn start_for(
path: &Path,
plugin: Option<&PluginId>,
log: LogSink,
cwd: Option<&Path>,
) -> Result<PluginProcess> {
let timeout = std::env::var("DRE_PLUGIN_HANDSHAKE_TIMEOUT_MS")
.ok()
.and_then(|v| v.parse().ok())
.map(Duration::from_millis)
.unwrap_or(DEFAULT_HANDSHAKE_TIMEOUT);
let mut p = Self::spawn_in(path, log, &[], cwd)?;
if let Some(id) = plugin {
p.serve = Some(id.clone());
}
p.handshake((MIN_VERSION, MAX_VERSION), timeout)?;
Ok(p)
}
pub fn start_with(
path: &Path,
log: LogSink,
versions: (u32, u32),
timeout: Duration,
) -> Result<PluginProcess> {
let mut p = Self::spawn(path, log)?;
p.handshake(versions, timeout)?;
Ok(p)
}
pub fn spawn(path: &Path, log: LogSink) -> Result<PluginProcess> {
Self::spawn_env(path, log, &[])
}
pub fn spawn_env(path: &Path, log: LogSink, env: &[(&str, &str)]) -> Result<PluginProcess> {
Self::spawn_in(path, log, env, None)
}
pub fn spawn_in(
path: &Path,
log: LogSink,
env: &[(&str, &str)],
cwd: Option<&Path>,
) -> Result<PluginProcess> {
let label = path
.file_name()
.map(|f| f.to_string_lossy().to_string())
.unwrap_or_default();
let mut cmd = Command::new(path);
if let Some(dir) = cwd {
cmd.current_dir(dir);
}
if std::env::var_os(INTERACTIVE_ENV).is_none() {
use std::io::IsTerminal;
let there = std::io::stderr().is_terminal() && std::io::stdin().is_terminal();
cmd.env(INTERACTIVE_ENV, if there { "1" } else { "0" });
}
let mut child = cmd
.envs(env.iter().copied())
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.map_err(|error| HostError::Spawn {
plugin: label.clone(),
error,
})?;
let stdout = child.stdout.take().unwrap();
let stderr_pipe = child.stderr.take().unwrap();
let stdin = child.stdin.take().map(BufWriter::new);
let (tx, rx) = channel();
std::thread::spawn(move || {
let mut r = BufReader::new(stdout);
loop {
let f = frame::read_frame(&mut r);
let stop = f.is_err();
if tx.send(f).is_err() || stop {
break;
}
}
});
let stderr = Arc::new(Mutex::new(VecDeque::new()));
let tail = stderr.clone();
let who = label.clone();
let stderr_thread = std::thread::spawn(move || {
for line in BufReader::new(stderr_pipe).lines() {
let Ok(line) = line else { break };
log(&who, &line);
let mut t = tail.lock().unwrap();
if t.len() == STDERR_TAIL {
t.pop_front();
}
t.push_back(line);
}
});
let serve = parse_executable_name(&label).map(|(k, n)| PluginId::new(k, n));
Ok(PluginProcess {
label,
path: path.to_path_buf(),
child,
stdin,
rx,
stderr,
stderr_thread: Some(stderr_thread),
info: None,
early: None,
serve,
})
}
pub fn ask_for(&mut self, plugin: Option<PluginId>) {
self.serve = plugin;
}
pub fn handshake(&mut self, (min, max): (u32, u32), timeout: Duration) -> Result<()> {
self.send(&Request::Hello {
min_version: min,
max_version: max,
core_version: core_version(),
plugin: self.serve.clone(),
})?;
match self.recv(Some(timeout), "the hello reply")? {
Incoming::Json(Response::Hello {
protocol_version,
kind,
name,
version,
capabilities,
mut provides,
}) => {
if protocol_version < min || protocol_version > max {
return Err(HostError::Incompatible {
plugin: self.label.clone(),
core: (min, max),
plugin_range: (protocol_version, protocol_version),
});
}
if let Some(want) = &self.serve
&& (want.kind != kind || want.name != name)
{
return Err(HostError::Plugin {
plugin: self.label.clone(),
message: format!(
"`{}` was asked for the {want} plugin but serves {kind}/{name}",
self.label
),
});
}
if provides.is_empty() {
provides.push(PluginId::new(kind, name.clone()));
}
self.info = Some(PluginInfo {
protocol_version,
kind,
name,
version,
capabilities,
provides,
});
Ok(())
}
Incoming::Json(Response::VersionMismatch {
min_version,
max_version,
}) => Err(HostError::Incompatible {
plugin: self.label.clone(),
core: (min, max),
plugin_range: (min_version, max_version),
}),
Incoming::Json(Response::Error { message }) => Err(HostError::Plugin {
plugin: self.label.clone(),
message,
}),
other => Err(self.unexpected("a hello reply", &other)),
}
}
pub fn label(&self) -> &str {
&self.label
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn info(&self) -> &PluginInfo {
self.info.as_ref().expect("handshake completed")
}
pub fn has(&self, capability: &str) -> bool {
self.info
.as_ref()
.is_some_and(|i| i.capabilities.iter().any(|c| c == capability))
}
fn unexpected(&self, expected: &'static str, got: &Incoming) -> HostError {
let got = match got {
Incoming::Json(r) => format!("{:?}", r)
.split([' ', '{'])
.next()
.unwrap_or("?")
.to_string(),
Incoming::Arrow(_) => "Arrow data".into(),
};
HostError::Unexpected {
plugin: self.label.clone(),
expected,
got: format!("`{got}`"),
}
}
fn crashed(&mut self) -> HostError {
let mut status = None;
for _ in 0..50 {
if let Ok(Some(s)) = self.child.try_wait() {
status = Some(s);
break;
}
std::thread::sleep(Duration::from_millis(20));
}
std::thread::sleep(Duration::from_millis(50));
let stderr = self.stderr.lock().unwrap().iter().cloned().collect();
HostError::Crashed {
plugin: self.label.clone(),
status,
stderr,
}
}
pub fn send(&mut self, req: &Request) -> Result<()> {
let ok = match self.stdin.as_mut() {
Some(w) => frame::write_json(w, req).is_ok(),
None => false,
};
if ok { Ok(()) } else { Err(self.write_failed()) }
}
pub fn send_batch(&mut self, batch: &RecordBatch) -> Result<()> {
if self.early.is_none()
&& let Ok(f) = self.rx.try_recv()
{
self.early = Some(f);
}
if let Some(Ok(Frame::Json(v))) = &self.early
&& let Ok(Response::Error { message }) = serde_json::from_value::<Response>(v.clone())
{
return Err(HostError::Plugin {
plugin: self.label.clone(),
message,
});
}
let ipc = frame::encode_batch(batch).map_err(|e| HostError::Arrow {
plugin: self.label.clone(),
message: e.to_string(),
})?;
let ok = match self.stdin.as_mut() {
Some(w) => frame::write_arrow(w, &ipc).is_ok(),
None => false,
};
if ok { Ok(()) } else { Err(self.write_failed()) }
}
fn write_failed(&mut self) -> HostError {
match self.rx.recv_timeout(Duration::from_secs(2)) {
Ok(Ok(Frame::Json(v))) => match serde_json::from_value::<Response>(v) {
Ok(Response::Error { message }) => HostError::Plugin {
plugin: self.label.clone(),
message,
},
_ => self.crashed(),
},
_ => self.crashed(),
}
}
pub fn recv(&mut self, timeout: Option<Duration>, waiting_for: &'static str) -> Result<Incoming> {
let got = match timeout {
_ if self.early.is_some() => self.early.take(),
Some(t) => match self.rx.recv_timeout(t) {
Ok(f) => Some(f),
Err(RecvTimeoutError::Timeout) => {
return Err(HostError::Timeout {
plugin: self.label.clone(),
waiting_for,
});
}
Err(RecvTimeoutError::Disconnected) => None,
},
None => self.rx.recv().ok(),
};
match got {
Some(Ok(Frame::Json(v))) => match serde_json::from_value::<Response>(v.clone()) {
Ok(r) => Ok(Incoming::Json(r)),
Err(e) => Err(HostError::Malformed {
plugin: self.label.clone(),
message: format!("unknown message {v}: {e}"),
}),
},
Some(Ok(Frame::Arrow(b))) => Ok(Incoming::Arrow(b)),
Some(Err(FrameError::Malformed(m))) => Err(HostError::Malformed {
plugin: self.label.clone(),
message: m,
}),
Some(Err(_)) | None => Err(self.crashed()),
}
}
pub fn recv_json(&mut self, waiting_for: &'static str) -> Result<Response> {
match self.recv(None, waiting_for)? {
Incoming::Json(Response::Error { message }) => Err(HostError::Plugin {
plugin: self.label.clone(),
message,
}),
Incoming::Json(r) => Ok(r),
other => Err(self.unexpected(waiting_for, &other)),
}
}
fn call_ok(&mut self, req: &Request, what: &'static str) -> Result<()> {
self.send(req)?;
match self.recv_json(what)? {
Response::Ok {} => Ok(()),
other => Err(self.unexpected(what, &Incoming::Json(other))),
}
}
pub fn describe(&mut self) -> Result<Vec<ConnectionField>> {
Ok(self.describe_all()?.0)
}
pub fn describe_all(&mut self) -> Result<(Vec<ConnectionField>, Vec<OptionField>)> {
self.send(&Request::Describe {})?;
match self.recv_json("a describe reply")? {
Response::Describe {
connection_fields,
option_fields,
} => Ok((connection_fields, option_fields)),
other => Err(self.unexpected("a describe reply", &Incoming::Json(other))),
}
}
pub fn validate(&mut self, options: Map<String, Value>) -> Result<Vec<String>> {
self.send(&Request::Validate { options })?;
match self.recv_json("a validate reply")? {
Response::Validated { errors } => Ok(errors),
other => Err(self.unexpected("a validate reply", &Incoming::Json(other))),
}
}
pub fn open(&mut self, connection: Map<String, Value>, read_only: bool) -> Result<()> {
self.call_ok(
&Request::Open {
connection,
read_only,
},
"an open reply",
)
}
pub fn check(&mut self, sql: &str) -> Result<()> {
self.call_ok(&Request::Check { sql: sql.to_string() }, "a check reply")
}
pub fn execute(
&mut self,
sql: &str,
row_limit: Option<u64>,
mut on_batch: impl FnMut(&SchemaRef, RecordBatch) -> std::result::Result<(), String>,
) -> Result<Execution> {
self.send(&Request::Execute {
sql: sql.to_string(),
row_limit,
})?;
match self.recv_json("an execute reply")? {
Response::NoResult { rows_affected } => return Ok(Execution::NoResult { rows_affected }),
Response::Result { .. } => {}
other => return Err(self.unexpected("an execute reply", &Incoming::Json(other))),
}
let mut schema: Option<SchemaRef> = None;
let mut rows = 0u64;
let mut sink_error = None;
loop {
match self.recv(None, "result data")? {
Incoming::Arrow(ipc) => {
let (s, batches) = frame::decode_batches(&ipc).map_err(|e| HostError::Arrow {
plugin: self.label.clone(),
message: e.to_string(),
})?;
let schema = schema.get_or_insert(s);
for b in batches {
rows += b.num_rows() as u64;
if sink_error.is_none()
&& let Err(e) = on_batch(schema, b)
{
sink_error = Some(e);
}
}
}
Incoming::Json(Response::ResultEnd { .. }) => break,
Incoming::Json(Response::Error { message }) => {
return Err(HostError::Plugin {
plugin: self.label.clone(),
message,
});
}
other => return Err(self.unexpected("result data", &other)),
}
}
if let Some(e) = sink_error {
return Err(HostError::Plugin {
plugin: self.label.clone(),
message: e,
});
}
let schema = schema.ok_or_else(|| HostError::Malformed {
plugin: self.label.clone(),
message: "a result ended without any Arrow frame carrying its schema".into(),
})?;
Ok(Execution::Result { schema, rows })
}
pub fn write_begin(
&mut self,
path: &str,
format: &str,
options: Map<String, Value>,
result_sets: Vec<ResultSetMeta>,
template: Option<Value>,
) -> Result<()> {
self.send(&Request::Write {
path: path.to_string(),
format: format.to_string(),
options,
result_sets,
template,
})
}
pub fn write_result_set(
&mut self,
schema: &SchemaRef,
batches: impl IntoIterator<Item = RecordBatch>,
) -> Result<()> {
let mut any = false;
for b in batches {
any = true;
self.send_batch(&b)?;
}
if !any {
self.send_batch(&RecordBatch::new_empty(schema.clone()))?;
}
self.send(&Request::ResultSetEnd {})
}
pub fn load(
&mut self,
name: &str,
schema: &SchemaRef,
batches: impl IntoIterator<Item = RecordBatch>,
) -> Result<(String, Option<String>)> {
self.send(&Request::Load {
name: name.to_string(),
})?;
self.write_result_set(schema, batches)?;
match self.recv_json("a loaded reply")? {
Response::Loaded {
relation, warning, ..
} => Ok((relation, warning)),
other => Err(self.unexpected("a loaded reply", &Incoming::Json(other))),
}
}
pub fn write_finish(&mut self) -> Result<Vec<String>> {
Ok(self.write_finish_with_warnings()?.0)
}
pub fn write_finish_with_warnings(&mut self) -> Result<(Vec<String>, Vec<String>)> {
self.send(&Request::Finish {})?;
match self.recv_json("a written reply")? {
Response::Written { files, warnings } => Ok((files, warnings)),
other => Err(self.unexpected("a written reply", &Incoming::Json(other))),
}
}
pub fn deliver(
&mut self,
local_path: &str,
remote_path: Option<&str>,
connection: Map<String, Value>,
) -> Result<String> {
let file = DeliveryFile {
local_path: local_path.to_string(),
remote_path: remote_path.map(str::to_string),
};
self.deliver_files(&[file], connection, Map::new())
}
pub fn deliver_files(
&mut self,
files: &[DeliveryFile],
connection: Map<String, Value>,
options: Map<String, Value>,
) -> Result<String> {
let req = match files {
[] => {
return Err(HostError::Malformed {
plugin: self.label.clone(),
message: "nothing to deliver".into(),
});
}
[one] => Request::Deliver {
local_path: Some(one.local_path.clone()),
remote_path: one.remote_path.clone(),
files: Vec::new(),
connection,
options,
},
many => {
if !self.has(crate::CAP_MULTI_FILE) {
return Err(HostError::Plugin {
plugin: self.label.clone(),
message: format!(
"{} files in one delivery, but the plugin doesn't advertise `multi_file`",
many.len()
),
});
}
Request::Deliver {
local_path: None,
remote_path: None,
files: many.to_vec(),
connection,
options,
}
}
};
self.send(&req)?;
match self.recv_json("a delivered reply")? {
Response::Delivered { location } => Ok(location),
other => Err(self.unexpected("a delivered reply", &Incoming::Json(other))),
}
}
pub fn close(mut self) -> Result<()> {
let _ = self.send(&Request::Close {});
let _ = self.recv(Some(Duration::from_secs(5)), "the close reply");
self.stdin.take();
for _ in 0..250 {
if let Ok(Some(_)) = self.child.try_wait() {
return Ok(());
}
std::thread::sleep(Duration::from_millis(20));
}
let _ = self.child.kill();
Ok(())
}
pub fn write_raw(&mut self, bytes: &[u8]) -> Result<()> {
let ok = match self.stdin.as_mut() {
Some(w) => w.write_all(bytes).and_then(|_| w.flush()).is_ok(),
None => false,
};
if ok { Ok(()) } else { Err(self.write_failed()) }
}
pub fn wait_exit(&mut self, timeout: Duration) -> Option<ExitStatus> {
self.stdin.take();
let deadline = std::time::Instant::now() + timeout;
while std::time::Instant::now() < deadline {
if let Ok(Some(s)) = self.child.try_wait() {
return Some(s);
}
std::thread::sleep(Duration::from_millis(20));
}
None
}
}
impl PluginProcess {
fn drain_stderr(&mut self) {
if let Some(t) = self.stderr_thread.take() {
let deadline = std::time::Instant::now() + Duration::from_secs(2);
while !t.is_finished() && std::time::Instant::now() < deadline {
std::thread::sleep(Duration::from_millis(5));
}
if t.is_finished() {
let _ = t.join();
}
}
}
}
impl Drop for PluginProcess {
fn drop(&mut self) {
self.stdin.take();
if let Ok(None) = self.child.try_wait() {
for _ in 0..25 {
if let Ok(Some(_)) = self.child.try_wait() {
break;
}
std::thread::sleep(Duration::from_millis(10));
}
let _ = self.child.kill();
let _ = self.child.wait();
}
self.drain_stderr();
}
}
static CORE_VERSION: OnceLock<String> = OnceLock::new();
pub fn set_core_version(version: &str) {
let _ = CORE_VERSION.set(version.to_string());
}
fn core_version() -> String {
CORE_VERSION.get().cloned().unwrap_or_else(|| "unreleased".into())
}