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    /// A hint added to the "unknown option" error for `key`, for an option this format doesn't
509    /// take but a sibling format does (`password` on `html`).
510    fn unknown_option_hint(&self, _key: &str) -> Option<String> {
511        None
512    }
513    /// Write every result set to `req.path` (and siblings, if the format needs several files);
514    /// return the files written.
515    fn write(&mut self, req: &WriteRequest, sets: &mut ResultSets<'_>) -> Result<Vec<String>>;
516}
517
518/// A destination `deliver` request.
519pub struct Delivery {
520    /// One file, or every file of an output when the plugin advertises `multi_file`. With a
521    /// `message`: the files to attach to it (often none).
522    pub files: Vec<DeliveryFile>,
523    pub connection: Map<String, Value>,
524    /// The destination entry's plugin options, rendered by core (empty when none).
525    pub options: Map<String, Value>,
526    /// A message to post; only sent to a plugin advertising `message`.
527    pub message: Option<crate::msg::Message>,
528}
529
530/// One file to deliver: the local copy in `target/run/` and its rendered remote path, if any.
531pub struct DeliveryFile {
532    pub local: PathBuf,
533    pub remote: Option<String>,
534}
535
536pub trait Destination {
537    fn connection_fields(&self) -> Vec<ConnectionField> {
538        Vec::new()
539    }
540    /// The options a report's destination entry takes: every key but `profile` and `path`.
541    /// String values may still hold Jinja when checked (see [`options::is_template`]); core
542    /// renders it before `deliver`.
543    fn options(&self) -> Vec<OptionField> {
544        Vec::new()
545    }
546    /// Problems the declared [`Destination::options`] can't catch, each a sentence naming the key.
547    /// Only called on options whose declared types already passed.
548    fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
549        Vec::new()
550    }
551    /// Rules the declared [`Destination::connection_fields`] can't express, each a sentence
552    /// naming the field, never its value. Static: no network. Run by `validate_connection` and
553    /// before every delivery.
554    fn validate_connection(&self, _connection: &Map<String, Value>) -> Vec<String> {
555        Vec::new()
556    }
557    /// Deliver `local` to `remote` (rendered by core); return where it landed. Enough for a
558    /// destination that takes no options; others implement [`Destination::deliver_files`].
559    fn deliver(
560        &mut self,
561        _local: &Path,
562        _remote: Option<&str>,
563        _connection: &Map<String, Value>,
564    ) -> Result<String> {
565        Err("this destination doesn't implement `deliver`".into())
566    }
567    /// The most characters a message may have in the service (after translation), reported in
568    /// `describe` by a plugin advertising `message`.
569    fn message_limit(&self) -> Option<u64> {
570        None
571    }
572    /// Post `d.message` (with `d.files` attached). Called only for a plugin advertising
573    /// `message`, which implements it.
574    fn deliver_message(&mut self, _d: &Delivery, _m: &crate::msg::Message) -> Result<String> {
575        Err("this destination doesn't take messages".into())
576    }
577    /// The whole request, options included (already checked). The default hands a single file
578    /// to [`Destination::deliver`]; several files arrive only with `multi_file` advertised.
579    fn deliver_files(&mut self, d: &Delivery) -> Result<String> {
580        match d.files.as_slice() {
581            [f] => self.deliver(&f.local, f.remote.as_deref(), &d.connection),
582            _ => Err("this destination takes one file per delivery".into()),
583        }
584    }
585}
586
587/// The frames core sends, read on a background thread (see [`read_stdin`]).
588pub struct Input {
589    rx: Receiver<std::result::Result<Frame, FrameError>>,
590}
591
592impl Input {
593    fn next(&mut self) -> std::result::Result<Frame, FrameError> {
594        self.rx.recv().unwrap_or(Err(FrameError::Eof))
595    }
596
597    fn read(&mut self) -> Result<Frame> {
598        self.next().map_err(|e| -> Error {
599            match e {
600                FrameError::Eof => "core closed the connection".into(),
601                e => e.to_string().into(),
602            }
603        })
604    }
605}
606
607/// Read stdin into `tx`, acting on `cancel` at once rather than queueing it behind the request
608/// it cancels. End of input cancels the running request too: core has gone.
609fn read_stdin(tx: Sender<std::result::Result<Frame, FrameError>>) {
610    let mut r = BufReader::new(std::io::stdin());
611    loop {
612        let f = frame::read_frame(&mut r);
613        match &f {
614            Ok(Frame::Json(v)) if v.get("type").and_then(Value::as_str) == Some("cancel") => {
615                cancel(v.get("id").and_then(Value::as_u64).unwrap_or(0));
616                continue;
617            }
618            Err(_) => cancel(shared().current.load(Ordering::SeqCst)),
619            Ok(_) => {}
620        }
621        let stop = f.is_err();
622        if tx.send(f).is_err() || stop {
623            break;
624        }
625    }
626}
627
628/// Writes replies and result data (see [`send`]).
629struct Output;
630
631impl Output {
632    fn send(&mut self, r: &Response) {
633        send(r)
634    }
635
636    fn batch(&mut self, b: &RecordBatch) -> Result<()> {
637        let ipc = frame::encode_batch(b)?;
638        let mut out = shared().out.lock().unwrap();
639        frame::write_arrow(&mut *out, &ipc).map_err(|_| -> Error { "core closed the connection".into() })
640    }
641}
642
643enum Handler<'a> {
644    Source(&'a mut dyn Source),
645    Format(&'a mut dyn Format),
646    Destination(&'a mut dyn Destination),
647}
648
649impl Handler<'_> {
650    fn kind(&self) -> Kind {
651        match self {
652            Handler::Source(_) => Kind::Source,
653            Handler::Format(_) => Kind::Format,
654            Handler::Destination(_) => Kind::Destination,
655        }
656    }
657
658    fn option_fields(&self) -> Vec<OptionField> {
659        match self {
660            Handler::Source(_) => Vec::new(),
661            Handler::Format(f) => f.options(),
662            Handler::Destination(d) => d.options(),
663        }
664    }
665
666    /// Every problem with a config block: the declared checks, then the plugin's own once the
667    /// declared ones pass.
668    fn check(&self, name: &str, o: &Map<String, Value>) -> Vec<String> {
669        let mut errs = options::check(self.kind(), name, &self.option_fields(), o);
670        if let Handler::Format(f) = self {
671            for e in &mut errs {
672                let hint = o
673                    .keys()
674                    .find(|k| e.starts_with(&format!("unknown option `{k}`")))
675                    .and_then(|k| f.unknown_option_hint(k));
676                if let Some(h) = hint {
677                    e.push_str("; ");
678                    e.push_str(&h);
679                }
680            }
681        }
682        if !errs.is_empty() {
683            return errs;
684        }
685        match self {
686            Handler::Source(_) => Vec::new(),
687            Handler::Format(f) => f.validate(o),
688            Handler::Destination(d) => d.validate(o),
689        }
690    }
691
692    /// The config block's problems as one error, for `write` and `deliver`.
693    fn checked(&self, name: &str, o: &Map<String, Value>) -> Result<()> {
694        let errs = self.check(name, o);
695        if errs.is_empty() {
696            Ok(())
697        } else {
698            // Wrong options are the project's settings: kind `config`.
699            Err(PluginError::new(ErrorKind::Config, errs.join("; ")).into())
700        }
701    }
702}
703
704pub fn serve_source(about: About, s: impl Source + 'static) -> ! {
705    serve_package(vec![Plugin::Source(about, Box::new(s))])
706}
707
708pub fn serve_format(about: About, f: impl Format + 'static) -> ! {
709    serve_package(vec![Plugin::Format(about, Box::new(f))])
710}
711
712pub fn serve_destination(about: About, d: impl Destination + 'static) -> ! {
713    serve_package(vec![Plugin::Destination(about, Box::new(d))])
714}
715
716/// One of the plugins a package's executable serves (see [`serve_package`]).
717pub enum Plugin {
718    Source(About, Box<dyn Source>),
719    Format(About, Box<dyn Format>),
720    Destination(About, Box<dyn Destination>),
721}
722
723impl Plugin {
724    fn id(&self) -> PluginId {
725        let (kind, about) = match self {
726            Plugin::Source(a, _) => (Kind::Source, a),
727            Plugin::Format(a, _) => (Kind::Format, a),
728            Plugin::Destination(a, _) => (Kind::Destination, a),
729        };
730        PluginId::new(kind, about.name)
731    }
732
733    fn handler(&mut self) -> (&About, Handler<'_>) {
734        match self {
735            Plugin::Source(a, s) => (a, Handler::Source(s.as_mut())),
736            Plugin::Format(a, f) => (a, Handler::Format(f.as_mut())),
737            Plugin::Destination(a, d) => (a, Handler::Destination(d.as_mut())),
738        }
739    }
740}
741
742/// Serve every plugin of a package from one executable. Core's `hello` names the plugin it
743/// wants; without one, the first is served.
744pub fn serve_package(mut plugins: Vec<Plugin>) -> ! {
745    assert!(!plugins.is_empty(), "a package serves at least one plugin");
746    // Take stdout for frames before anything else can print to it.
747    shared();
748    let (tx, rx) = std::sync::mpsc::channel();
749    std::thread::spawn(move || read_stdin(tx));
750    let mut input = Input { rx };
751    let mut out = Output;
752    if ::log::set_logger(&FrameLogger).is_ok() {
753        ::log::set_max_level(::log::LevelFilter::Debug);
754    }
755    let provides: Vec<PluginId> = plugins.iter().map(Plugin::id).collect();
756    let chosen = loop {
757        let req = match next_request(&mut input, &mut out) {
758            Some(r) => r,
759            None => continue,
760        };
761        let Request::Hello {
762            min_version,
763            max_version,
764            plugin,
765            ..
766        } = req
767        else {
768            out.send(&Response::error("the first request must be `hello`"));
769            continue;
770        };
771        let Some(chosen) = negotiate((min_version, max_version), (MIN_VERSION, MAX_VERSION)) else {
772            out.send(&Response::VersionMismatch {
773                min_version: MIN_VERSION,
774                max_version: MAX_VERSION,
775            });
776            std::process::exit(1);
777        };
778        let i = match &plugin {
779            None => 0,
780            Some(want) => match provides.iter().position(|p| p == want) {
781                Some(i) => i,
782                None => {
783                    let list: Vec<String> = provides.iter().map(|p| p.to_string()).collect();
784                    out.send(&Response::error(format!(
785                        "this executable provides {}, not {want}",
786                        list.join(", ")
787                    )));
788                    std::process::exit(1);
789                }
790            },
791        };
792        let (about, h) = plugins[i].handler();
793        let _ = shared().name.set(about.name.to_string());
794        out.send(&Response::Hello {
795            protocol_version: chosen,
796            kind: h.kind(),
797            name: about.name.to_string(),
798            version: about.version.to_string(),
799            capabilities: about
800                .capabilities
801                .iter()
802                .chain(std::iter::once(&CAP_VALIDATE))
803                .map(|c| c.to_string())
804                .collect(),
805            provides: if provides.len() > 1 {
806                provides.clone()
807            } else {
808                Vec::new()
809            },
810        });
811        shared().version.store(chosen, Ordering::SeqCst);
812        break i;
813    };
814    let (about, h) = plugins[chosen].handler();
815    serve(about.name, h, input, out)
816}
817
818/// The next request, or `None` after answering one that can't be read. Exits at end of input.
819fn next_request(input: &mut Input, out: &mut Output) -> Option<Request> {
820    let frame = match input.next() {
821        Ok(f) => f,
822        Err(FrameError::Eof) => std::process::exit(0),
823        Err(e) => {
824            eprintln!("{e}");
825            out.send(&Response::error(e.to_string()));
826            std::process::exit(2);
827        }
828    };
829    match frame {
830        Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
831            Ok(r) => Some(r),
832            Err(_) => {
833                let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
834                out.send(&Response::error(format!("unsupported request `{t}`")));
835                None
836            }
837        },
838        Frame::Arrow(_) => {
839            out.send(&Response::error("unexpected Arrow frame"));
840            None
841        }
842    }
843}
844
845/// Serve requests after the handshake.
846fn serve(name: &str, mut h: Handler<'_>, mut input: Input, mut out: Output) -> ! {
847    loop {
848        let frame = match input.next() {
849            Ok(f) => f,
850            Err(FrameError::Eof) => {
851                if let Handler::Source(src) = &mut h {
852                    src.close();
853                }
854                std::process::exit(0)
855            }
856            Err(e) => {
857                eprintln!("{e}");
858                out.send(&Response::error(e.to_string()));
859                std::process::exit(2);
860            }
861        };
862        let s = shared();
863        let id = match &frame {
864            Frame::Json(v) => v.get("id").and_then(Value::as_u64).unwrap_or(0),
865            Frame::Arrow(_) => 0,
866        };
867        s.current.store(id, Ordering::SeqCst);
868        let req = match frame {
869            Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
870                Ok(r) => r,
871                Err(_) => {
872                    let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
873                    out.send(&Response::error(format!("unsupported request `{t}`")));
874                    s.current.store(0, Ordering::SeqCst);
875                    continue;
876                }
877            },
878            Frame::Arrow(_) => {
879                out.send(&Response::error("unexpected Arrow frame"));
880                continue;
881            }
882        };
883        if let Request::Hello { .. } = req {
884            out.send(&Response::error("`hello` was already answered"));
885            s.current.store(0, Ordering::SeqCst);
886            continue;
887        }
888        if let Request::Close {} = req {
889            if let Handler::Source(src) = &mut h {
890                src.close();
891            }
892            out.send(&Response::Ok {});
893            std::process::exit(0);
894        }
895        // Cancelled before it started: it never runs.
896        let result = if cancelled() {
897            Ok(Err("cancelled before it started".into()))
898        } else {
899            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
900                handle(&mut h, name, req, &mut input, &mut out)
901            }))
902        };
903        match result {
904            Ok(Ok(())) => {}
905            Ok(Err(e)) => out.send(&error_reply(&e)),
906            Err(p) => {
907                let msg = p
908                    .downcast_ref::<String>()
909                    .cloned()
910                    .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()))
911                    .unwrap_or_else(|| "unknown panic".into());
912                out.send(&Response::Error {
913                    message: format!("plugin panicked: {msg}"),
914                    kind: Some("internal".into()),
915                    code: None,
916                });
917                std::process::exit(101);
918            }
919        }
920        s.current.store(0, Ordering::SeqCst);
921    }
922}
923
924/// The declared fields' generic checks, then the plugin's own rules (when the generic ones
925/// passed), with secret values redacted.
926fn check_connection(
927    h: &Handler<'_>,
928    name: &str,
929    connection: &Map<String, Value>,
930    unresolved: &[String],
931) -> Result<crate::connection::Checked> {
932    let (kind, fields) = match h {
933        Handler::Source(s) => (Kind::Source, s.connection_fields()),
934        Handler::Destination(d) => (Kind::Destination, d.connection_fields()),
935        Handler::Format(_) => return Err("a format has no connection settings".into()),
936    };
937    let mut c = crate::connection::check(kind, name, &fields, connection, unresolved);
938    if c.errors.is_empty() {
939        // An unresolved value arrives as `null`: set, but not known (a rule about two keys
940        // going together checks `contains_key`).
941        c.errors = match h {
942            Handler::Source(s) => s.validate_connection(connection),
943            Handler::Destination(d) => d.validate_connection(connection),
944            Handler::Format(_) => Vec::new(),
945        };
946    }
947    c.errors = crate::connection::redact(c.errors, &fields, connection);
948    c.warnings = crate::connection::redact(c.warnings, &fields, connection);
949    Ok(c)
950}
951
952fn handle(h: &mut Handler<'_>, name: &str, req: Request, input: &mut Input, out: &mut Output) -> Result<()> {
953    if let Request::Describe {} = req {
954        let connection_fields = match h {
955            Handler::Source(s) => s.connection_fields(),
956            Handler::Destination(d) => d.connection_fields(),
957            Handler::Format(_) => Vec::new(),
958        };
959        let (identifier_quote, string_literal) = match h {
960            Handler::Source(s) => (
961                s.identifier_quote().map(str::to_string),
962                s.string_literal().map(str::to_string),
963            ),
964            _ => (None, None),
965        };
966        let message_limit = match h {
967            Handler::Destination(d) => d.message_limit(),
968            _ => None,
969        };
970        out.send(&Response::Describe {
971            connection_fields,
972            option_fields: h.option_fields(),
973            identifier_quote,
974            string_literal,
975            message_limit,
976        });
977        return Ok(());
978    }
979    if let Request::Validate { options, template } = req {
980        let mut errors = h.check(name, &options);
981        let mut warnings = Vec::new();
982        if let (Handler::Format(f), Some(t), true) = (&*h, &template, errors.is_empty()) {
983            let (e, w) = f.validate_template(&options, t);
984            errors.extend(e);
985            warnings.extend(w);
986        }
987        out.send(&Response::Validated { errors, warnings });
988        return Ok(());
989    }
990    if let Request::ValidateConnection {
991        connection,
992        unresolved,
993    } = req
994    {
995        let c = check_connection(h, name, &connection, &unresolved)?;
996        out.send(&Response::Validated {
997            errors: c.errors,
998            warnings: c.warnings,
999        });
1000        return Ok(());
1001    }
1002    // The same checks before connecting: errors stop it, warnings are logged.
1003    if let Request::Open { connection, .. } | Request::Deliver { connection, .. } = &req {
1004        let c = check_connection(h, name, connection, &[])?;
1005        for w in &c.warnings {
1006            crate::log::warn!("{w}");
1007        }
1008        if !c.errors.is_empty() {
1009            return Err(PluginError::new(ErrorKind::Config, c.errors.join("; ")).into());
1010        }
1011    }
1012    let checked = match &req {
1013        Request::Write { options, .. } | Request::Deliver { options, .. } => h.checked(name, options),
1014        _ => Ok(()),
1015    };
1016    match (h, req) {
1017        (
1018            Handler::Source(s),
1019            Request::Open {
1020                connection,
1021                read_only,
1022                sandbox,
1023            },
1024        ) => {
1025            if sandbox {
1026                s.open_sandboxed(&connection)?;
1027            } else {
1028                s.open(&connection, read_only)?;
1029            }
1030            out.send(&Response::Ok {});
1031        }
1032        (Handler::Source(s), Request::Load { name }) => {
1033            let ipc = match input.read()? {
1034                Frame::Arrow(ipc) => ipc,
1035                Frame::Json(v) => return Err(format!("expected the rows to load, got {v}").into()),
1036            };
1037            let (schema, batches) = frame::decode_batches(&ipc)?;
1038            let mut data = ResultSet {
1039                meta: ResultSetMeta {
1040                    name: name.clone(),
1041                    query: name.clone(),
1042                    result_index: 1,
1043                    anchor: None,
1044                    header: None,
1045                    columns: Default::default(),
1046                    autofit: None,
1047                    style: None,
1048                },
1049                schema,
1050                first: Some(batches),
1051                input,
1052                done: false,
1053            };
1054            let r = s.load(&name, &mut data);
1055            data.drain()?;
1056            let loaded = r?;
1057            out.send(&Response::Loaded {
1058                relation: loaded.relation,
1059                rows: loaded.rows,
1060                warning: loaded.warning,
1061            });
1062        }
1063        (Handler::Source(s), Request::Check { sql }) => {
1064            s.check(&sql)?;
1065            out.send(&Response::Ok {});
1066        }
1067        (Handler::Source(s), Request::Execute { sql, row_limit }) => {
1068            let mut sink = FrameSink {
1069                out,
1070                row_limit,
1071                rows: 0,
1072                state: SinkState::Idle,
1073                schema: None,
1074                sent_any: false,
1075            };
1076            let r = s.execute(&sql, row_limit, &mut sink);
1077            match (r, sink.state) {
1078                (Err(e), _) => return Err(e),
1079                (Ok(()), SinkState::Idle) => {
1080                    return Err("the source produced neither a result nor `no_result`".into());
1081                }
1082                (Ok(()), SinkState::NoResult) => {}
1083                (Ok(()), SinkState::Result) => {
1084                    if !sink.sent_any {
1085                        let schema = sink.schema.clone().unwrap();
1086                        sink.out.batch(&RecordBatch::new_empty(schema))?;
1087                    }
1088                    let rows = sink.rows;
1089                    sink.out.send(&Response::ResultEnd { rows });
1090                }
1091            }
1092        }
1093        (
1094            Handler::Format(f),
1095            Request::Write {
1096                path,
1097                format,
1098                options,
1099                result_sets,
1100                template,
1101            },
1102        ) => {
1103            let req = WriteRequest {
1104                path,
1105                format,
1106                options,
1107                result_sets: result_sets.clone(),
1108                template,
1109            };
1110            let mut sets = ResultSets {
1111                metas: result_sets.into_iter(),
1112                input,
1113                warnings: Vec::new(),
1114            };
1115            let written = checked.and_then(|()| f.write(&req, &mut sets));
1116            // A failed write is reported at once, so core can stop streaming; this is the
1117            // request's one reply.
1118            if let Err(e) = &written {
1119                out.send(&error_reply(e));
1120            }
1121            // Consume whatever the format didn't read (after an error, possibly the rest of a
1122            // result set), so the stream stays in sync.
1123            loop {
1124                match sets.input.read()? {
1125                    Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("finish") => break,
1126                    Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("result_set_end") => {}
1127                    Frame::Arrow(_) => {}
1128                    other => return Err(format!("expected `finish`, got {other:?}").into()),
1129                }
1130            }
1131            if let Ok(files) = written {
1132                out.send(&Response::Written {
1133                    files,
1134                    warnings: sets.warnings,
1135                });
1136            }
1137        }
1138        (
1139            Handler::Destination(d),
1140            Request::Deliver {
1141                local_path,
1142                remote_path,
1143                files,
1144                connection,
1145                options,
1146                message,
1147            },
1148        ) => {
1149            checked?;
1150            let files = match (local_path, files.is_empty(), &message) {
1151                (Some(local), true, None) => vec![DeliveryFile {
1152                    local: PathBuf::from(local),
1153                    remote: remote_path,
1154                }],
1155                (None, false, _) | (None, true, Some(_)) => files
1156                    .into_iter()
1157                    .map(|f| DeliveryFile {
1158                        local: PathBuf::from(f.local_path),
1159                        remote: f.remote_path,
1160                    })
1161                    .collect(),
1162                _ => return Err("`deliver` needs exactly one of `local_path` or `files`".into()),
1163            };
1164            let delivery = Delivery {
1165                files,
1166                connection,
1167                options,
1168                message,
1169            };
1170            crate::delivery::take_attempts();
1171            let location = match &delivery.message {
1172                Some(m) => d.deliver_message(&delivery, m)?,
1173                None => d.deliver_files(&delivery)?,
1174            };
1175            let attempts = Some(crate::delivery::take_attempts()).filter(|n| *n > 1);
1176            out.send(&Response::Delivered { location, attempts });
1177        }
1178        (h, req) => {
1179            let t = serde_json::to_value(&req)
1180                .ok()
1181                .and_then(|v| v.get("type").cloned())
1182                .unwrap_or_default();
1183            return Err(format!("a {} plugin doesn't handle {t} requests", h.kind()).into());
1184        }
1185    }
1186    Ok(())
1187}
1188
1189#[derive(PartialEq)]
1190enum SinkState {
1191    Idle,
1192    NoResult,
1193    Result,
1194}
1195
1196struct FrameSink<'a> {
1197    out: &'a mut Output,
1198    row_limit: Option<u64>,
1199    rows: u64,
1200    state: SinkState,
1201    schema: Option<SchemaRef>,
1202    sent_any: bool,
1203}
1204
1205impl ResultSink for FrameSink<'_> {
1206    fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()> {
1207        if self.state != SinkState::Idle {
1208            return Err("`no_result` after the result started".into());
1209        }
1210        self.state = SinkState::NoResult;
1211        self.out.send(&Response::NoResult { rows_affected });
1212        Ok(())
1213    }
1214
1215    fn begin(&mut self, schema: SchemaRef) -> Result<()> {
1216        if self.state != SinkState::Idle {
1217            return Err("a statement can only produce one result".into());
1218        }
1219        self.state = SinkState::Result;
1220        self.out.send(&Response::Result {
1221            columns: schema.fields().iter().map(|f| f.name().clone()).collect(),
1222        });
1223        self.schema = Some(schema);
1224        Ok(())
1225    }
1226
1227    fn batch(&mut self, b: RecordBatch) -> Result<bool> {
1228        if self.state != SinkState::Result {
1229            return Err("`batch` before `begin`".into());
1230        }
1231        if self.row_limit.is_some_and(|l| self.rows >= l) {
1232            return Ok(false);
1233        }
1234        let b = match self.row_limit {
1235            Some(l) if self.rows + b.num_rows() as u64 > l => b.slice(0, (l - self.rows) as usize),
1236            _ => b,
1237        };
1238        if b.num_rows() > 0 || !self.sent_any {
1239            self.rows += b.num_rows() as u64;
1240            self.out.batch(&b)?;
1241            self.sent_any = true;
1242        }
1243        Ok(self.row_limit.is_none_or(|l| self.rows < l))
1244    }
1245}
1246
1247/// The highest version in both ranges.
1248fn negotiate(core: (u32, u32), ours: (u32, u32)) -> Option<u32> {
1249    let (lo, hi) = (core.0.max(ours.0), core.1.min(ours.1));
1250    (lo <= hi).then_some(hi)
1251}
1252
1253/// Read a string field from a connection map.
1254pub fn conn_str<'a>(c: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
1255    c.get(key).and_then(Value::as_str).filter(|s| !s.is_empty())
1256}
1257
1258/// Read a required string field from a connection map.
1259pub fn conn_required<'a>(c: &'a Map<String, Value>, key: &str) -> Result<&'a str> {
1260    conn_str(c, key).ok_or_else(|| format!("the profile output needs a `{key}` field").into())
1261}
1262
1263/// Read a flag that may be written as a boolean or a string.
1264pub fn conn_bool(c: &Map<String, Value>, key: &str) -> Option<bool> {
1265    match c.get(key)? {
1266        Value::Bool(b) => Some(*b),
1267        Value::String(s) => Some(matches!(s.as_str(), "true" | "yes" | "1")),
1268        _ => None,
1269    }
1270}