Skip to main content

dre_protocol/
plugin.rs

1//! The plugin SDK: implement one of [`Source`], [`Format`] or [`Destination`] and call the
2//! matching `serve_*` function from `main`. The SDK owns stdin/stdout, the handshake, framing,
3//! request ids, error replies and panics.
4//!
5//! Formats and destinations declare their options (`options()`) and add any rule a declaration
6//! can't express (`validate()`). The SDK checks a config block against both on `validate` and
7//! before every `write` and `deliver`, so plugin code only sees options that passed.
8//!
9//! - **Logging:** use the `log` crate (`log::info!`, `log::warn!`), or [`log`] for a message with
10//!   fields; the SDK sends them to core as `log` messages. The plugin's own records keep their
11//!   level; its dependencies' are shown only with `-v`. Long requests may report [`progress`].
12//! - **Errors:** return a [`PluginError`] to give core the problem's kind and code; any other
13//!   error is sent as a plain message.
14//! - **Cancellation:** stdin is read on a background thread while a request runs. On `cancel`
15//!   (or when core goes away) [`cancelled`] turns true and the hook set with [`on_cancel`] runs,
16//!   e.g. to cancel a query on the server.
17
18use std::io::{BufReader, BufWriter};
19use std::path::{Path, PathBuf};
20use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
21use std::sync::mpsc::{Receiver, Sender};
22use std::sync::{Arc, Mutex, OnceLock};
23
24use arrow::array::RecordBatch;
25use arrow::datatypes::SchemaRef;
26use serde_json::{Map, Value};
27
28use crate::frame::{self, Frame, FrameError};
29use crate::msg::{ConnectionField, Envelope, LogLevel, Request, Response, ResultSetMeta};
30use crate::options::{self, OptionField};
31use crate::{CAP_VALIDATE, Kind, MAX_VERSION, MIN_VERSION, PluginId};
32
33pub type Error = Box<dyn std::error::Error + Send + Sync>;
34pub type Result<T> = std::result::Result<T, Error>;
35
36/// What kind of problem a [`PluginError`] is (the kinds of DRE's error-code registry).
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum ErrorKind {
39    /// The settings are wrong.
40    Config,
41    /// The plugin itself misbehaved.
42    Plugin,
43    /// Refused something unsafe.
44    Refused,
45    /// A connection couldn't be made; trying again may work.
46    Connection,
47    /// Credentials were refused.
48    Auth,
49    /// A query failed on the database.
50    Query,
51    /// A file or message couldn't be delivered.
52    Delivery,
53    /// Something unexpected: a bug.
54    Internal,
55    Cancelled,
56    TimedOut,
57}
58
59impl ErrorKind {
60    pub fn as_str(self) -> &'static str {
61        match self {
62            ErrorKind::Config => "config",
63            ErrorKind::Plugin => "plugin",
64            ErrorKind::Refused => "refused",
65            ErrorKind::Connection => "connection",
66            ErrorKind::Auth => "auth",
67            ErrorKind::Query => "query",
68            ErrorKind::Delivery => "delivery",
69            ErrorKind::Internal => "internal",
70            ErrorKind::Cancelled => "cancelled",
71            ErrorKind::TimedOut => "timed_out",
72        }
73    }
74}
75
76/// An error with its kind and, optionally, a code. The SDK namespaces the code by the plugin
77/// (`host-key-mismatch` is sent as `sftp/host-key-mismatch`).
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct PluginError {
80    pub kind: ErrorKind,
81    pub code: Option<String>,
82    pub message: String,
83}
84
85impl PluginError {
86    pub fn new(kind: ErrorKind, message: impl Into<String>) -> PluginError {
87        PluginError {
88            kind,
89            code: None,
90            message: message.into(),
91        }
92    }
93
94    pub fn code(mut self, code: impl Into<String>) -> PluginError {
95        self.code = Some(code.into());
96        self
97    }
98}
99
100impl std::fmt::Display for PluginError {
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        f.write_str(&self.message)
103    }
104}
105
106impl std::error::Error for PluginError {}
107
108/// The SDK's process-wide state, shared by the request loop, the stdin reader, and logging
109/// from any thread.
110struct Shared {
111    out: Mutex<BufWriter<Box<dyn std::io::Write + Send>>>,
112    /// The protocol version settled on (0 until the handshake).
113    version: AtomicU32,
114    /// The request running now (0: none).
115    current: AtomicU64,
116    /// The last request core cancelled (0: none).
117    cancelled: AtomicU64,
118    hook: Mutex<Option<CancelHook>>,
119    /// The plugin being served, for namespacing its error codes.
120    name: OnceLock<String>,
121}
122
123type CancelHook = Arc<dyn Fn() + Send + Sync>;
124
125fn shared() -> &'static Shared {
126    static SHARED: OnceLock<Shared> = OnceLock::new();
127    SHARED.get_or_init(|| Shared {
128        out: Mutex::new(BufWriter::new(protocol_out())),
129        version: AtomicU32::new(0),
130        current: AtomicU64::new(0),
131        cancelled: AtomicU64::new(0),
132        hook: Mutex::new(None),
133        name: OnceLock::new(),
134    })
135}
136
137/// Where frames go: the process's original stdout, kept for the protocol alone. Stdout itself is
138/// pointed at stderr, so a dependency that prints (`println!`) can't corrupt the frame stream;
139/// core keeps stderr for crash reports.
140fn protocol_out() -> Box<dyn std::io::Write + Send> {
141    #[cfg(unix)]
142    {
143        use std::os::fd::FromRawFd;
144        // SAFETY: plain descriptor calls on fds 1 and 2, made once before any frame is written.
145        unsafe {
146            let fd = libc::dup(1);
147            if fd >= 0 && libc::dup2(2, 1) >= 0 {
148                return Box::new(std::fs::File::from_raw_fd(fd));
149            }
150        }
151    }
152    #[cfg(windows)]
153    {
154        use std::os::windows::io::FromRawHandle;
155        use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
156        use windows_sys::Win32::System::Console::{
157            GetStdHandle, STD_ERROR_HANDLE, STD_OUTPUT_HANDLE, SetStdHandle,
158        };
159        // SAFETY: the process's own standard handles, swapped once before any frame is written;
160        // Rust's stdout looks its handle up on every write, so `println!` then reaches stderr.
161        unsafe {
162            let out = GetStdHandle(STD_OUTPUT_HANDLE);
163            let err = GetStdHandle(STD_ERROR_HANDLE);
164            if !out.is_null() && out != INVALID_HANDLE_VALUE && SetStdHandle(STD_OUTPUT_HANDLE, err) != 0 {
165                return Box::new(std::fs::File::from_raw_handle(out as _));
166            }
167        }
168    }
169    Box::new(std::io::stdout())
170}
171
172/// Whether core has cancelled the request running now. Long loops check it between steps.
173pub fn cancelled() -> bool {
174    let s = shared();
175    let current = s.current.load(Ordering::SeqCst);
176    current != 0 && s.cancelled.load(Ordering::SeqCst) == current
177}
178
179/// Run `hook` when core cancels a request, on the stdin reader's thread, while the request is
180/// still running on the main one (e.g. ask the server to cancel the query). Replaces any hook
181/// set before. The request then fails, and the SDK replies with kind `cancelled`.
182pub fn on_cancel(hook: impl Fn() + Send + Sync + 'static) {
183    *shared().hook.lock().unwrap() = Some(Arc::new(hook));
184}
185
186/// Cancel request `id`: remembered, so a request cancelled before it starts never runs, and
187/// the hook runs if it's running now.
188fn cancel(id: u64) {
189    let s = shared();
190    s.cancelled.store(id, Ordering::SeqCst);
191    if id != 0 && s.current.load(Ordering::SeqCst) == id {
192        let hook = s.hook.lock().unwrap().clone();
193        if let Some(hook) = hook {
194            hook();
195        }
196    }
197}
198
199/// Send a log message to core, with structured `fields` (never secrets).
200pub fn log_fields(level: LogLevel, message: &str, fields: Map<String, Value>) {
201    let s = shared();
202    if s.version.load(Ordering::SeqCst) == 0 {
203        // Protocol 0 has no `log` message: stderr, with the prefixes core 0.3 shows.
204        let prefix = match level {
205            LogLevel::Error | LogLevel::Warn => "warning: ",
206            LogLevel::Info => "info: ",
207            LogLevel::Debug | LogLevel::Trace => "",
208        };
209        eprintln!("{prefix}{message}");
210        return;
211    }
212    send(&Response::Log {
213        level,
214        message: message.to_string(),
215        fields,
216    });
217}
218
219/// Send a log message to core.
220pub fn log(level: LogLevel, message: &str) {
221    log_fields(level, message, Map::new())
222}
223
224/// Tell core how far a long request has got; it shows the latest.
225pub fn progress(message: Option<&str>, done: Option<u64>, total: Option<u64>) {
226    if shared().version.load(Ordering::SeqCst) == 0 {
227        if let Some(m) = message {
228            log(LogLevel::Info, m);
229        }
230        return;
231    }
232    send(&Response::Progress {
233        message: message.map(str::to_string),
234        done,
235        total,
236    });
237}
238
239/// Bridges the `log` crate to `log` messages.
240struct FrameLogger;
241
242impl ::log::Log for FrameLogger {
243    /// The plugin's own records down to debug; a dependency's (a driver, an HTTP client, a
244    /// layout engine) down to info only: their debug output can run to millions of lines.
245    fn enabled(&self, m: &::log::Metadata) -> bool {
246        m.level() <= ::log::Level::Info || (m.level() <= ::log::Level::Debug && m.target().starts_with("dre"))
247    }
248
249    fn log(&self, r: &::log::Record) {
250        if !self.enabled(r.metadata()) {
251            return;
252        }
253        let level = match r.level() {
254            ::log::Level::Error => LogLevel::Error,
255            ::log::Level::Warn => LogLevel::Warn,
256            ::log::Level::Info => LogLevel::Info,
257            ::log::Level::Debug => LogLevel::Debug,
258            ::log::Level::Trace => LogLevel::Trace,
259        };
260        // A dependency's records (a driver, an HTTP client) are for `-v` only.
261        let level = if r.target().starts_with("dre") {
262            level
263        } else {
264            level.max(LogLevel::Debug)
265        };
266        log(level, &r.args().to_string());
267    }
268
269    fn flush(&self) {}
270}
271
272/// Write one reply or streamed message, with the running request's id from protocol 1.
273fn send(r: &Response) {
274    let s = shared();
275    let id = match s.version.load(Ordering::SeqCst) {
276        0 => None,
277        _ => Some(s.current.load(Ordering::SeqCst)).filter(|id| *id != 0),
278    };
279    let mut out = s.out.lock().unwrap();
280    if frame::write_json(&mut *out, &Envelope::new(id, r)).is_err() {
281        // Core has gone; nothing left to talk to.
282        std::process::exit(1);
283    }
284}
285
286/// The `error` reply for a failed request: kind `cancelled` if core cancelled it, else the
287/// kind and code of a [`PluginError`], else the message alone.
288fn error_reply(e: &Error) -> Response {
289    if cancelled() {
290        return Response::Error {
291            message: e.to_string(),
292            kind: Some("cancelled".into()),
293            code: None,
294        };
295    }
296    match e.downcast_ref::<PluginError>() {
297        Some(pe) => Response::Error {
298            message: pe.message.clone(),
299            kind: Some(pe.kind.as_str().into()),
300            code: pe.code.as_ref().map(|c| match shared().name.get() {
301                Some(name) if !c.contains('/') => format!("{name}/{c}"),
302                _ => c.clone(),
303            }),
304        },
305        None => Response::error(e.to_string()),
306    }
307}
308
309/// Identity reported in the handshake.
310#[derive(Debug, Clone)]
311pub struct About {
312    pub name: &'static str,
313    pub version: &'static str,
314    pub capabilities: &'static [&'static str],
315}
316
317impl About {
318    /// `version` "0.0.0" (the crate has no version yet) is reported as "unreleased".
319    pub fn new(name: &'static str, version: &'static str) -> About {
320        let version = if version == "0.0.0" { "unreleased" } else { version };
321        About {
322            name,
323            version,
324            capabilities: &[],
325        }
326    }
327
328    pub fn capabilities(mut self, caps: &'static [&'static str]) -> About {
329        self.capabilities = caps;
330        self
331    }
332}
333
334/// Where a source pushes the outcome of `execute`: either `no_result`, or `begin` followed by
335/// any number of `batch` calls.
336pub trait ResultSink {
337    fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()>;
338    fn begin(&mut self, schema: SchemaRef) -> Result<()>;
339    /// Returns `false` once the row limit is reached: stop producing batches.
340    fn batch(&mut self, batch: RecordBatch) -> Result<bool>;
341}
342
343pub trait Source {
344    fn connection_fields(&self) -> Vec<ConnectionField> {
345        Vec::new()
346    }
347    /// The character the database quotes identifiers with, reported in `describe`: `"` for
348    /// most SQL databases, a backtick for Databricks and MySQL. Core doubles it inside a name.
349    fn identifier_quote(&self) -> Option<&'static str> {
350        None
351    }
352    /// How core should spell a string holding `'` or `\` (an inlined lookup), reported in
353    /// `describe`: `None` for `chr(39)`/`chr(92)`, `Some("backslash")` for a database without
354    /// `chr()`.
355    fn string_literal(&self) -> Option<&'static str> {
356        None
357    }
358    /// Rules the declared [`Source::connection_fields`] can't express (two keys that can't go
359    /// together, a value's form), each a sentence naming the field, never its value. Static:
360    /// no network. Run by `validate_connection` and before `open`.
361    /// A value set through an unset `env_var()` arrives as `null`.
362    fn validate_connection(&self, _connection: &Map<String, Value>) -> Vec<String> {
363        Vec::new()
364    }
365    fn open(&mut self, connection: &Map<String, Value>, read_only: bool) -> Result<()>;
366    /// Open a read-only session for core's `dre query`: also refuse files, programs and the
367    /// network outside the database where the engine can. Default: a read-only [`Source::open`].
368    fn open_sandboxed(&mut self, connection: &Map<String, Value>) -> Result<()> {
369        self.open(connection, true)
370    }
371    /// Run one statement. `row_limit` is a hint; the SDK enforces it either way.
372    fn execute(&mut self, sql: &str, row_limit: Option<u64>, out: &mut dyn ResultSink) -> Result<()>;
373    /// Verify without executing. Only called when the plugin advertises `check`.
374    fn check(&mut self, _sql: &str) -> Result<()> {
375        Err("this source can't check statements".into())
376    }
377    /// Load `data` into a temporary table on the session, named after `name`. Only called when
378    /// the plugin advertises `load`.
379    fn load(&mut self, _name: &str, _data: &mut ResultSet<'_>) -> Result<Loaded> {
380        Err("this source can't load rows".into())
381    }
382    /// End the session cleanly (called on `close` and at end of input, before exiting).
383    fn close(&mut self) {}
384}
385
386/// The outcome of `Source::load`.
387pub struct Loaded {
388    /// How SQL refers to the loaded rows, e.g. the temp table's name.
389    pub relation: String,
390    pub rows: u64,
391    /// Shown to the user, e.g. when the database has no bulk path.
392    pub warning: Option<String>,
393}
394
395/// One incoming result set, streamed from core.
396pub struct ResultSet<'a> {
397    pub meta: ResultSetMeta,
398    pub schema: SchemaRef,
399    first: Option<Vec<RecordBatch>>,
400    input: &'a mut Input,
401    done: bool,
402}
403
404impl ResultSet<'_> {
405    /// The next batch, or `None` at the end of this result set.
406    pub fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
407        loop {
408            if let Some(first) = self.first.as_mut() {
409                if !first.is_empty() {
410                    return Ok(Some(first.remove(0)));
411                }
412                self.first = None;
413            }
414            if self.done {
415                return Ok(None);
416            }
417            match self.input.read()? {
418                Frame::Arrow(ipc) => {
419                    let (_, batches) = frame::decode_batches(&ipc)?;
420                    self.first = Some(batches);
421                }
422                Frame::Json(v) => match serde_json::from_value::<Request>(v)? {
423                    Request::ResultSetEnd {} => {
424                        self.done = true;
425                        return Ok(None);
426                    }
427                    other => return Err(format!("expected result data, got {other:?}").into()),
428                },
429            }
430        }
431    }
432
433    fn drain(&mut self) -> Result<()> {
434        while self.next_batch()?.is_some() {}
435        Ok(())
436    }
437}
438
439/// A format `write` request.
440pub struct WriteRequest {
441    pub path: String,
442    pub format: String,
443    pub options: Map<String, Value>,
444    pub result_sets: Vec<ResultSetMeta>,
445    pub template: Option<Value>,
446}
447
448/// Hands a format plugin its result sets one at a time, in order.
449pub struct ResultSets<'a> {
450    metas: std::vec::IntoIter<ResultSetMeta>,
451    input: &'a mut Input,
452    warnings: Vec<String>,
453}
454
455impl ResultSets<'_> {
456    /// Tell the person something about the output, shown as a warning by core.
457    pub fn warn(&mut self, message: impl Into<String>) {
458        self.warnings.push(message.into());
459    }
460
461    /// The next result set, or `None` after the last. Each must be read (or is drained) before
462    /// the next is requested.
463    pub fn next_set(&mut self) -> Result<Option<ResultSet<'_>>> {
464        let Some(meta) = self.metas.next() else {
465            return Ok(None);
466        };
467        let ipc = match self.input.read()? {
468            Frame::Arrow(ipc) => ipc,
469            Frame::Json(v) => {
470                return Err(format!(
471                    "expected an Arrow frame starting result set `{}`, got {v}",
472                    meta.name
473                )
474                .into());
475            }
476        };
477        let (schema, batches) = frame::decode_batches(&ipc)?;
478        Ok(Some(ResultSet {
479            meta,
480            schema,
481            first: Some(batches),
482            input: &mut *self.input,
483            done: false,
484        }))
485    }
486}
487
488pub trait Format {
489    /// The options this format takes: every key of a report's `output:` block except `format`,
490    /// `destination`, `template` and `extension`, which core owns.
491    fn options(&self) -> Vec<OptionField> {
492        Vec::new()
493    }
494    /// Problems the declared [`Format::options`] can't catch, each a sentence naming the key.
495    /// Only called on options whose declared types already passed.
496    fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
497        Vec::new()
498    }
499    /// Problems with the output's template (`validate`'s `template`, sent for the formats core
500    /// gives a template context): errors, then warnings. Only called when the options passed.
501    fn validate_template(
502        &self,
503        _options: &Map<String, Value>,
504        _template: &Value,
505    ) -> (Vec<String>, Vec<String>) {
506        (Vec::new(), Vec::new())
507    }
508    /// Write every result set to `req.path` (and siblings, if the format needs several files);
509    /// return the files written.
510    fn write(&mut self, req: &WriteRequest, sets: &mut ResultSets<'_>) -> Result<Vec<String>>;
511}
512
513/// A destination `deliver` request.
514pub struct Delivery {
515    /// One file, or every file of an output when the plugin advertises `multi_file`. With a
516    /// `message`: the files to attach to it (often none).
517    pub files: Vec<DeliveryFile>,
518    pub connection: Map<String, Value>,
519    /// The destination entry's plugin options, rendered by core (empty when none).
520    pub options: Map<String, Value>,
521    /// A message to post; only sent to a plugin advertising `message`.
522    pub message: Option<crate::msg::Message>,
523}
524
525/// One file to deliver: the local copy in `target/run/` and its rendered remote path, if any.
526pub struct DeliveryFile {
527    pub local: PathBuf,
528    pub remote: Option<String>,
529}
530
531pub trait Destination {
532    fn connection_fields(&self) -> Vec<ConnectionField> {
533        Vec::new()
534    }
535    /// The options a report's destination entry takes: every key but `profile` and `path`.
536    /// String values may still hold Jinja when checked (see [`options::is_template`]); core
537    /// renders it before `deliver`.
538    fn options(&self) -> Vec<OptionField> {
539        Vec::new()
540    }
541    /// Problems the declared [`Destination::options`] can't catch, each a sentence naming the key.
542    /// Only called on options whose declared types already passed.
543    fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
544        Vec::new()
545    }
546    /// Rules the declared [`Destination::connection_fields`] can't express, each a sentence
547    /// naming the field, never its value. Static: no network. Run by `validate_connection` and
548    /// before every delivery.
549    fn validate_connection(&self, _connection: &Map<String, Value>) -> Vec<String> {
550        Vec::new()
551    }
552    /// Deliver `local` to `remote` (rendered by core); return where it landed. Enough for a
553    /// destination that takes no options; others implement [`Destination::deliver_files`].
554    fn deliver(
555        &mut self,
556        _local: &Path,
557        _remote: Option<&str>,
558        _connection: &Map<String, Value>,
559    ) -> Result<String> {
560        Err("this destination doesn't implement `deliver`".into())
561    }
562    /// The most characters a message may have in the service (after translation), reported in
563    /// `describe` by a plugin advertising `message`.
564    fn message_limit(&self) -> Option<u64> {
565        None
566    }
567    /// Post `d.message` (with `d.files` attached). Called only for a plugin advertising
568    /// `message`, which implements it.
569    fn deliver_message(&mut self, _d: &Delivery, _m: &crate::msg::Message) -> Result<String> {
570        Err("this destination doesn't take messages".into())
571    }
572    /// The whole request, options included (already checked). The default hands a single file
573    /// to [`Destination::deliver`]; several files arrive only with `multi_file` advertised.
574    fn deliver_files(&mut self, d: &Delivery) -> Result<String> {
575        match d.files.as_slice() {
576            [f] => self.deliver(&f.local, f.remote.as_deref(), &d.connection),
577            _ => Err("this destination takes one file per delivery".into()),
578        }
579    }
580}
581
582/// The frames core sends, read on a background thread (see [`read_stdin`]).
583pub struct Input {
584    rx: Receiver<std::result::Result<Frame, FrameError>>,
585}
586
587impl Input {
588    fn next(&mut self) -> std::result::Result<Frame, FrameError> {
589        self.rx.recv().unwrap_or(Err(FrameError::Eof))
590    }
591
592    fn read(&mut self) -> Result<Frame> {
593        self.next().map_err(|e| -> Error {
594            match e {
595                FrameError::Eof => "core closed the connection".into(),
596                e => e.to_string().into(),
597            }
598        })
599    }
600}
601
602/// Read stdin into `tx`, acting on `cancel` at once rather than queueing it behind the request
603/// it cancels. End of input cancels the running request too: core has gone.
604fn read_stdin(tx: Sender<std::result::Result<Frame, FrameError>>) {
605    let mut r = BufReader::new(std::io::stdin());
606    loop {
607        let f = frame::read_frame(&mut r);
608        match &f {
609            Ok(Frame::Json(v)) if v.get("type").and_then(Value::as_str) == Some("cancel") => {
610                cancel(v.get("id").and_then(Value::as_u64).unwrap_or(0));
611                continue;
612            }
613            Err(_) => cancel(shared().current.load(Ordering::SeqCst)),
614            Ok(_) => {}
615        }
616        let stop = f.is_err();
617        if tx.send(f).is_err() || stop {
618            break;
619        }
620    }
621}
622
623/// Writes replies and result data (see [`send`]).
624struct Output;
625
626impl Output {
627    fn send(&mut self, r: &Response) {
628        send(r)
629    }
630
631    fn batch(&mut self, b: &RecordBatch) -> Result<()> {
632        let ipc = frame::encode_batch(b)?;
633        let mut out = shared().out.lock().unwrap();
634        frame::write_arrow(&mut *out, &ipc).map_err(|_| -> Error { "core closed the connection".into() })
635    }
636}
637
638enum Handler<'a> {
639    Source(&'a mut dyn Source),
640    Format(&'a mut dyn Format),
641    Destination(&'a mut dyn Destination),
642}
643
644impl Handler<'_> {
645    fn kind(&self) -> Kind {
646        match self {
647            Handler::Source(_) => Kind::Source,
648            Handler::Format(_) => Kind::Format,
649            Handler::Destination(_) => Kind::Destination,
650        }
651    }
652
653    fn option_fields(&self) -> Vec<OptionField> {
654        match self {
655            Handler::Source(_) => Vec::new(),
656            Handler::Format(f) => f.options(),
657            Handler::Destination(d) => d.options(),
658        }
659    }
660
661    /// Every problem with a config block: the declared checks, then the plugin's own once the
662    /// declared ones pass.
663    fn check(&self, name: &str, o: &Map<String, Value>) -> Vec<String> {
664        let errs = options::check(self.kind(), name, &self.option_fields(), o);
665        if !errs.is_empty() {
666            return errs;
667        }
668        match self {
669            Handler::Source(_) => Vec::new(),
670            Handler::Format(f) => f.validate(o),
671            Handler::Destination(d) => d.validate(o),
672        }
673    }
674
675    /// The config block's problems as one error, for `write` and `deliver`.
676    fn checked(&self, name: &str, o: &Map<String, Value>) -> Result<()> {
677        let errs = self.check(name, o);
678        if errs.is_empty() {
679            Ok(())
680        } else {
681            // Wrong options are the project's settings: kind `config`.
682            Err(PluginError::new(ErrorKind::Config, errs.join("; ")).into())
683        }
684    }
685}
686
687pub fn serve_source(about: About, s: impl Source + 'static) -> ! {
688    serve_package(vec![Plugin::Source(about, Box::new(s))])
689}
690
691pub fn serve_format(about: About, f: impl Format + 'static) -> ! {
692    serve_package(vec![Plugin::Format(about, Box::new(f))])
693}
694
695pub fn serve_destination(about: About, d: impl Destination + 'static) -> ! {
696    serve_package(vec![Plugin::Destination(about, Box::new(d))])
697}
698
699/// One of the plugins a package's executable serves (see [`serve_package`]).
700pub enum Plugin {
701    Source(About, Box<dyn Source>),
702    Format(About, Box<dyn Format>),
703    Destination(About, Box<dyn Destination>),
704}
705
706impl Plugin {
707    fn id(&self) -> PluginId {
708        let (kind, about) = match self {
709            Plugin::Source(a, _) => (Kind::Source, a),
710            Plugin::Format(a, _) => (Kind::Format, a),
711            Plugin::Destination(a, _) => (Kind::Destination, a),
712        };
713        PluginId::new(kind, about.name)
714    }
715
716    fn handler(&mut self) -> (&About, Handler<'_>) {
717        match self {
718            Plugin::Source(a, s) => (a, Handler::Source(s.as_mut())),
719            Plugin::Format(a, f) => (a, Handler::Format(f.as_mut())),
720            Plugin::Destination(a, d) => (a, Handler::Destination(d.as_mut())),
721        }
722    }
723}
724
725/// Serve every plugin of a package from one executable. Core's `hello` names the plugin it
726/// wants; without one, the first is served.
727pub fn serve_package(mut plugins: Vec<Plugin>) -> ! {
728    assert!(!plugins.is_empty(), "a package serves at least one plugin");
729    // Take stdout for frames before anything else can print to it.
730    shared();
731    let (tx, rx) = std::sync::mpsc::channel();
732    std::thread::spawn(move || read_stdin(tx));
733    let mut input = Input { rx };
734    let mut out = Output;
735    if ::log::set_logger(&FrameLogger).is_ok() {
736        ::log::set_max_level(::log::LevelFilter::Debug);
737    }
738    let provides: Vec<PluginId> = plugins.iter().map(Plugin::id).collect();
739    let chosen = loop {
740        let req = match next_request(&mut input, &mut out) {
741            Some(r) => r,
742            None => continue,
743        };
744        let Request::Hello {
745            min_version,
746            max_version,
747            plugin,
748            ..
749        } = req
750        else {
751            out.send(&Response::error("the first request must be `hello`"));
752            continue;
753        };
754        let Some(chosen) = negotiate((min_version, max_version), (MIN_VERSION, MAX_VERSION)) else {
755            out.send(&Response::VersionMismatch {
756                min_version: MIN_VERSION,
757                max_version: MAX_VERSION,
758            });
759            std::process::exit(1);
760        };
761        let i = match &plugin {
762            None => 0,
763            Some(want) => match provides.iter().position(|p| p == want) {
764                Some(i) => i,
765                None => {
766                    let list: Vec<String> = provides.iter().map(|p| p.to_string()).collect();
767                    out.send(&Response::error(format!(
768                        "this executable provides {}, not {want}",
769                        list.join(", ")
770                    )));
771                    std::process::exit(1);
772                }
773            },
774        };
775        let (about, h) = plugins[i].handler();
776        let _ = shared().name.set(about.name.to_string());
777        out.send(&Response::Hello {
778            protocol_version: chosen,
779            kind: h.kind(),
780            name: about.name.to_string(),
781            version: about.version.to_string(),
782            capabilities: about
783                .capabilities
784                .iter()
785                .chain(std::iter::once(&CAP_VALIDATE))
786                .map(|c| c.to_string())
787                .collect(),
788            provides: if provides.len() > 1 {
789                provides.clone()
790            } else {
791                Vec::new()
792            },
793        });
794        shared().version.store(chosen, Ordering::SeqCst);
795        break i;
796    };
797    let (about, h) = plugins[chosen].handler();
798    serve(about.name, h, input, out)
799}
800
801/// The next request, or `None` after answering one that can't be read. Exits at end of input.
802fn next_request(input: &mut Input, out: &mut Output) -> Option<Request> {
803    let frame = match input.next() {
804        Ok(f) => f,
805        Err(FrameError::Eof) => std::process::exit(0),
806        Err(e) => {
807            eprintln!("{e}");
808            out.send(&Response::error(e.to_string()));
809            std::process::exit(2);
810        }
811    };
812    match frame {
813        Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
814            Ok(r) => Some(r),
815            Err(_) => {
816                let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
817                out.send(&Response::error(format!("unsupported request `{t}`")));
818                None
819            }
820        },
821        Frame::Arrow(_) => {
822            out.send(&Response::error("unexpected Arrow frame"));
823            None
824        }
825    }
826}
827
828/// Serve requests after the handshake.
829fn serve(name: &str, mut h: Handler<'_>, mut input: Input, mut out: Output) -> ! {
830    loop {
831        let frame = match input.next() {
832            Ok(f) => f,
833            Err(FrameError::Eof) => {
834                if let Handler::Source(src) = &mut h {
835                    src.close();
836                }
837                std::process::exit(0)
838            }
839            Err(e) => {
840                eprintln!("{e}");
841                out.send(&Response::error(e.to_string()));
842                std::process::exit(2);
843            }
844        };
845        let s = shared();
846        let id = match &frame {
847            Frame::Json(v) => v.get("id").and_then(Value::as_u64).unwrap_or(0),
848            Frame::Arrow(_) => 0,
849        };
850        s.current.store(id, Ordering::SeqCst);
851        let req = match frame {
852            Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
853                Ok(r) => r,
854                Err(_) => {
855                    let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
856                    out.send(&Response::error(format!("unsupported request `{t}`")));
857                    s.current.store(0, Ordering::SeqCst);
858                    continue;
859                }
860            },
861            Frame::Arrow(_) => {
862                out.send(&Response::error("unexpected Arrow frame"));
863                continue;
864            }
865        };
866        if let Request::Hello { .. } = req {
867            out.send(&Response::error("`hello` was already answered"));
868            s.current.store(0, Ordering::SeqCst);
869            continue;
870        }
871        if let Request::Close {} = req {
872            if let Handler::Source(src) = &mut h {
873                src.close();
874            }
875            out.send(&Response::Ok {});
876            std::process::exit(0);
877        }
878        // Cancelled before it started: it never runs.
879        let result = if cancelled() {
880            Ok(Err("cancelled before it started".into()))
881        } else {
882            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
883                handle(&mut h, name, req, &mut input, &mut out)
884            }))
885        };
886        match result {
887            Ok(Ok(())) => {}
888            Ok(Err(e)) => out.send(&error_reply(&e)),
889            Err(p) => {
890                let msg = p
891                    .downcast_ref::<String>()
892                    .cloned()
893                    .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()))
894                    .unwrap_or_else(|| "unknown panic".into());
895                out.send(&Response::Error {
896                    message: format!("plugin panicked: {msg}"),
897                    kind: Some("internal".into()),
898                    code: None,
899                });
900                std::process::exit(101);
901            }
902        }
903        s.current.store(0, Ordering::SeqCst);
904    }
905}
906
907/// The declared fields' generic checks, then the plugin's own rules (when the generic ones
908/// passed), with secret values redacted.
909fn check_connection(
910    h: &Handler<'_>,
911    name: &str,
912    connection: &Map<String, Value>,
913    unresolved: &[String],
914) -> Result<crate::connection::Checked> {
915    let (kind, fields) = match h {
916        Handler::Source(s) => (Kind::Source, s.connection_fields()),
917        Handler::Destination(d) => (Kind::Destination, d.connection_fields()),
918        Handler::Format(_) => return Err("a format has no connection settings".into()),
919    };
920    let mut c = crate::connection::check(kind, name, &fields, connection, unresolved);
921    if c.errors.is_empty() {
922        // An unresolved value arrives as `null`: set, but not known (a rule about two keys
923        // going together checks `contains_key`).
924        c.errors = match h {
925            Handler::Source(s) => s.validate_connection(connection),
926            Handler::Destination(d) => d.validate_connection(connection),
927            Handler::Format(_) => Vec::new(),
928        };
929    }
930    c.errors = crate::connection::redact(c.errors, &fields, connection);
931    c.warnings = crate::connection::redact(c.warnings, &fields, connection);
932    Ok(c)
933}
934
935fn handle(h: &mut Handler<'_>, name: &str, req: Request, input: &mut Input, out: &mut Output) -> Result<()> {
936    if let Request::Describe {} = req {
937        let connection_fields = match h {
938            Handler::Source(s) => s.connection_fields(),
939            Handler::Destination(d) => d.connection_fields(),
940            Handler::Format(_) => Vec::new(),
941        };
942        let (identifier_quote, string_literal) = match h {
943            Handler::Source(s) => (
944                s.identifier_quote().map(str::to_string),
945                s.string_literal().map(str::to_string),
946            ),
947            _ => (None, None),
948        };
949        let message_limit = match h {
950            Handler::Destination(d) => d.message_limit(),
951            _ => None,
952        };
953        out.send(&Response::Describe {
954            connection_fields,
955            option_fields: h.option_fields(),
956            identifier_quote,
957            string_literal,
958            message_limit,
959        });
960        return Ok(());
961    }
962    if let Request::Validate { options, template } = req {
963        let mut errors = h.check(name, &options);
964        let mut warnings = Vec::new();
965        if let (Handler::Format(f), Some(t), true) = (&*h, &template, errors.is_empty()) {
966            let (e, w) = f.validate_template(&options, t);
967            errors.extend(e);
968            warnings.extend(w);
969        }
970        out.send(&Response::Validated { errors, warnings });
971        return Ok(());
972    }
973    if let Request::ValidateConnection {
974        connection,
975        unresolved,
976    } = req
977    {
978        let c = check_connection(h, name, &connection, &unresolved)?;
979        out.send(&Response::Validated {
980            errors: c.errors,
981            warnings: c.warnings,
982        });
983        return Ok(());
984    }
985    // The same checks before connecting: errors stop it, warnings are logged.
986    if let Request::Open { connection, .. } | Request::Deliver { connection, .. } = &req {
987        let c = check_connection(h, name, connection, &[])?;
988        for w in &c.warnings {
989            crate::log::warn!("{w}");
990        }
991        if !c.errors.is_empty() {
992            return Err(PluginError::new(ErrorKind::Config, c.errors.join("; ")).into());
993        }
994    }
995    let checked = match &req {
996        Request::Write { options, .. } | Request::Deliver { options, .. } => h.checked(name, options),
997        _ => Ok(()),
998    };
999    match (h, req) {
1000        (
1001            Handler::Source(s),
1002            Request::Open {
1003                connection,
1004                read_only,
1005                sandbox,
1006            },
1007        ) => {
1008            if sandbox {
1009                s.open_sandboxed(&connection)?;
1010            } else {
1011                s.open(&connection, read_only)?;
1012            }
1013            out.send(&Response::Ok {});
1014        }
1015        (Handler::Source(s), Request::Load { name }) => {
1016            let ipc = match input.read()? {
1017                Frame::Arrow(ipc) => ipc,
1018                Frame::Json(v) => return Err(format!("expected the rows to load, got {v}").into()),
1019            };
1020            let (schema, batches) = frame::decode_batches(&ipc)?;
1021            let mut data = ResultSet {
1022                meta: ResultSetMeta {
1023                    name: name.clone(),
1024                    query: name.clone(),
1025                    result_index: 1,
1026                    anchor: None,
1027                    header: None,
1028                    columns: Default::default(),
1029                    autofit: None,
1030                    style: None,
1031                },
1032                schema,
1033                first: Some(batches),
1034                input,
1035                done: false,
1036            };
1037            let r = s.load(&name, &mut data);
1038            data.drain()?;
1039            let loaded = r?;
1040            out.send(&Response::Loaded {
1041                relation: loaded.relation,
1042                rows: loaded.rows,
1043                warning: loaded.warning,
1044            });
1045        }
1046        (Handler::Source(s), Request::Check { sql }) => {
1047            s.check(&sql)?;
1048            out.send(&Response::Ok {});
1049        }
1050        (Handler::Source(s), Request::Execute { sql, row_limit }) => {
1051            let mut sink = FrameSink {
1052                out,
1053                row_limit,
1054                rows: 0,
1055                state: SinkState::Idle,
1056                schema: None,
1057                sent_any: false,
1058            };
1059            let r = s.execute(&sql, row_limit, &mut sink);
1060            match (r, sink.state) {
1061                (Err(e), _) => return Err(e),
1062                (Ok(()), SinkState::Idle) => {
1063                    return Err("the source produced neither a result nor `no_result`".into());
1064                }
1065                (Ok(()), SinkState::NoResult) => {}
1066                (Ok(()), SinkState::Result) => {
1067                    if !sink.sent_any {
1068                        let schema = sink.schema.clone().unwrap();
1069                        sink.out.batch(&RecordBatch::new_empty(schema))?;
1070                    }
1071                    let rows = sink.rows;
1072                    sink.out.send(&Response::ResultEnd { rows });
1073                }
1074            }
1075        }
1076        (
1077            Handler::Format(f),
1078            Request::Write {
1079                path,
1080                format,
1081                options,
1082                result_sets,
1083                template,
1084            },
1085        ) => {
1086            let req = WriteRequest {
1087                path,
1088                format,
1089                options,
1090                result_sets: result_sets.clone(),
1091                template,
1092            };
1093            let mut sets = ResultSets {
1094                metas: result_sets.into_iter(),
1095                input,
1096                warnings: Vec::new(),
1097            };
1098            let written = checked.and_then(|()| f.write(&req, &mut sets));
1099            // A failed write is reported at once, so core can stop streaming; this is the
1100            // request's one reply.
1101            if let Err(e) = &written {
1102                out.send(&error_reply(e));
1103            }
1104            // Consume whatever the format didn't read (after an error, possibly the rest of a
1105            // result set), so the stream stays in sync.
1106            loop {
1107                match sets.input.read()? {
1108                    Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("finish") => break,
1109                    Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("result_set_end") => {}
1110                    Frame::Arrow(_) => {}
1111                    other => return Err(format!("expected `finish`, got {other:?}").into()),
1112                }
1113            }
1114            if let Ok(files) = written {
1115                out.send(&Response::Written {
1116                    files,
1117                    warnings: sets.warnings,
1118                });
1119            }
1120        }
1121        (
1122            Handler::Destination(d),
1123            Request::Deliver {
1124                local_path,
1125                remote_path,
1126                files,
1127                connection,
1128                options,
1129                message,
1130            },
1131        ) => {
1132            checked?;
1133            let files = match (local_path, files.is_empty(), &message) {
1134                (Some(local), true, None) => vec![DeliveryFile {
1135                    local: PathBuf::from(local),
1136                    remote: remote_path,
1137                }],
1138                (None, false, _) | (None, true, Some(_)) => files
1139                    .into_iter()
1140                    .map(|f| DeliveryFile {
1141                        local: PathBuf::from(f.local_path),
1142                        remote: f.remote_path,
1143                    })
1144                    .collect(),
1145                _ => return Err("`deliver` needs exactly one of `local_path` or `files`".into()),
1146            };
1147            let delivery = Delivery {
1148                files,
1149                connection,
1150                options,
1151                message,
1152            };
1153            crate::delivery::take_attempts();
1154            let location = match &delivery.message {
1155                Some(m) => d.deliver_message(&delivery, m)?,
1156                None => d.deliver_files(&delivery)?,
1157            };
1158            let attempts = Some(crate::delivery::take_attempts()).filter(|n| *n > 1);
1159            out.send(&Response::Delivered { location, attempts });
1160        }
1161        (h, req) => {
1162            let t = serde_json::to_value(&req)
1163                .ok()
1164                .and_then(|v| v.get("type").cloned())
1165                .unwrap_or_default();
1166            return Err(format!("a {} plugin doesn't handle {t} requests", h.kind()).into());
1167        }
1168    }
1169    Ok(())
1170}
1171
1172#[derive(PartialEq)]
1173enum SinkState {
1174    Idle,
1175    NoResult,
1176    Result,
1177}
1178
1179struct FrameSink<'a> {
1180    out: &'a mut Output,
1181    row_limit: Option<u64>,
1182    rows: u64,
1183    state: SinkState,
1184    schema: Option<SchemaRef>,
1185    sent_any: bool,
1186}
1187
1188impl ResultSink for FrameSink<'_> {
1189    fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()> {
1190        if self.state != SinkState::Idle {
1191            return Err("`no_result` after the result started".into());
1192        }
1193        self.state = SinkState::NoResult;
1194        self.out.send(&Response::NoResult { rows_affected });
1195        Ok(())
1196    }
1197
1198    fn begin(&mut self, schema: SchemaRef) -> Result<()> {
1199        if self.state != SinkState::Idle {
1200            return Err("a statement can only produce one result".into());
1201        }
1202        self.state = SinkState::Result;
1203        self.out.send(&Response::Result {
1204            columns: schema.fields().iter().map(|f| f.name().clone()).collect(),
1205        });
1206        self.schema = Some(schema);
1207        Ok(())
1208    }
1209
1210    fn batch(&mut self, b: RecordBatch) -> Result<bool> {
1211        if self.state != SinkState::Result {
1212            return Err("`batch` before `begin`".into());
1213        }
1214        if self.row_limit.is_some_and(|l| self.rows >= l) {
1215            return Ok(false);
1216        }
1217        let b = match self.row_limit {
1218            Some(l) if self.rows + b.num_rows() as u64 > l => b.slice(0, (l - self.rows) as usize),
1219            _ => b,
1220        };
1221        if b.num_rows() > 0 || !self.sent_any {
1222            self.rows += b.num_rows() as u64;
1223            self.out.batch(&b)?;
1224            self.sent_any = true;
1225        }
1226        Ok(self.row_limit.is_none_or(|l| self.rows < l))
1227    }
1228}
1229
1230/// The highest version in both ranges.
1231fn negotiate(core: (u32, u32), ours: (u32, u32)) -> Option<u32> {
1232    let (lo, hi) = (core.0.max(ours.0), core.1.min(ours.1));
1233    (lo <= hi).then_some(hi)
1234}
1235
1236/// Read a string field from a connection map.
1237pub fn conn_str<'a>(c: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
1238    c.get(key).and_then(Value::as_str).filter(|s| !s.is_empty())
1239}
1240
1241/// Read a required string field from a connection map.
1242pub fn conn_required<'a>(c: &'a Map<String, Value>, key: &str) -> Result<&'a str> {
1243    conn_str(c, key).ok_or_else(|| format!("the profile output needs a `{key}` field").into())
1244}
1245
1246/// Read a flag that may be written as a boolean or a string.
1247pub fn conn_bool(c: &Map<String, Value>, key: &str) -> Option<bool> {
1248    match c.get(key)? {
1249        Value::Bool(b) => Some(*b),
1250        Value::String(s) => Some(matches!(s.as_str(), "true" | "yes" | "1")),
1251        _ => None,
1252    }
1253}