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