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//! error replies and panics; plugins log with `eprintln!`.
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
9use std::io::{BufReader, BufWriter, Read, Stdout};
10use std::path::{Path, PathBuf};
11
12use arrow::array::RecordBatch;
13use arrow::datatypes::SchemaRef;
14use serde_json::{Map, Value};
15
16use crate::frame::{self, Frame, FrameError};
17use crate::msg::{ConnectionField, Request, Response, ResultSetMeta};
18use crate::options::{self, OptionField};
19use crate::{CAP_VALIDATE, Kind, MAX_VERSION, MIN_VERSION, PluginId};
20
21pub type Error = Box<dyn std::error::Error + Send + Sync>;
22pub type Result<T> = std::result::Result<T, Error>;
23
24/// Identity reported in the handshake.
25#[derive(Debug, Clone)]
26pub struct About {
27    pub name: &'static str,
28    pub version: &'static str,
29    pub capabilities: &'static [&'static str],
30}
31
32impl About {
33    /// `version` "0.0.0" (the crate has no version yet) is reported as "unreleased".
34    pub fn new(name: &'static str, version: &'static str) -> About {
35        let version = if version == "0.0.0" { "unreleased" } else { version };
36        About {
37            name,
38            version,
39            capabilities: &[],
40        }
41    }
42
43    pub fn capabilities(mut self, caps: &'static [&'static str]) -> About {
44        self.capabilities = caps;
45        self
46    }
47}
48
49/// Where a source pushes the outcome of `execute`: either `no_result`, or `begin` followed by
50/// any number of `batch` calls.
51pub trait ResultSink {
52    fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()>;
53    fn begin(&mut self, schema: SchemaRef) -> Result<()>;
54    /// Returns `false` once the row limit is reached: stop producing batches.
55    fn batch(&mut self, batch: RecordBatch) -> Result<bool>;
56}
57
58pub trait Source {
59    fn connection_fields(&self) -> Vec<ConnectionField> {
60        Vec::new()
61    }
62    /// The character the database quotes identifiers with, reported in `describe`: `"` for
63    /// most SQL databases, a backtick for Databricks and MySQL. Core doubles it inside a name.
64    fn identifier_quote(&self) -> Option<&'static str> {
65        None
66    }
67    fn open(&mut self, connection: &Map<String, Value>, read_only: bool) -> Result<()>;
68    /// Run one statement. `row_limit` is a hint; the SDK enforces it either way.
69    fn execute(&mut self, sql: &str, row_limit: Option<u64>, out: &mut dyn ResultSink) -> Result<()>;
70    /// Verify without executing. Only called when the plugin advertises `check`.
71    fn check(&mut self, _sql: &str) -> Result<()> {
72        Err("this source can't check statements".into())
73    }
74    /// Load `data` into a temporary table on the session, named after `name`. Only called when
75    /// the plugin advertises `load`.
76    fn load(&mut self, _name: &str, _data: &mut ResultSet<'_>) -> Result<Loaded> {
77        Err("this source can't load rows".into())
78    }
79    /// End the session cleanly (called on `close` and at end of input, before exiting).
80    fn close(&mut self) {}
81}
82
83/// The outcome of `Source::load`.
84pub struct Loaded {
85    /// How SQL refers to the loaded rows, e.g. the temp table's name.
86    pub relation: String,
87    pub rows: u64,
88    /// Shown to the user, e.g. when the database has no bulk path.
89    pub warning: Option<String>,
90}
91
92/// One incoming result set, streamed from core.
93pub struct ResultSet<'a> {
94    pub meta: ResultSetMeta,
95    pub schema: SchemaRef,
96    first: Option<Vec<RecordBatch>>,
97    input: &'a mut Input,
98    done: bool,
99}
100
101impl ResultSet<'_> {
102    /// The next batch, or `None` at the end of this result set.
103    pub fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
104        loop {
105            if let Some(first) = self.first.as_mut() {
106                if !first.is_empty() {
107                    return Ok(Some(first.remove(0)));
108                }
109                self.first = None;
110            }
111            if self.done {
112                return Ok(None);
113            }
114            match self.input.read()? {
115                Frame::Arrow(ipc) => {
116                    let (_, batches) = frame::decode_batches(&ipc)?;
117                    self.first = Some(batches);
118                }
119                Frame::Json(v) => match serde_json::from_value::<Request>(v)? {
120                    Request::ResultSetEnd {} => {
121                        self.done = true;
122                        return Ok(None);
123                    }
124                    other => return Err(format!("expected result data, got {other:?}").into()),
125                },
126            }
127        }
128    }
129
130    fn drain(&mut self) -> Result<()> {
131        while self.next_batch()?.is_some() {}
132        Ok(())
133    }
134}
135
136/// A format `write` request.
137pub struct WriteRequest {
138    pub path: String,
139    pub format: String,
140    pub options: Map<String, Value>,
141    pub result_sets: Vec<ResultSetMeta>,
142    pub template: Option<Value>,
143}
144
145/// Hands a format plugin its result sets one at a time, in order.
146pub struct ResultSets<'a> {
147    metas: std::vec::IntoIter<ResultSetMeta>,
148    input: &'a mut Input,
149    warnings: Vec<String>,
150}
151
152impl ResultSets<'_> {
153    /// Tell the person something about the output, shown as a warning by core.
154    pub fn warn(&mut self, message: impl Into<String>) {
155        self.warnings.push(message.into());
156    }
157
158    /// The next result set, or `None` after the last. Each must be read (or is drained) before
159    /// the next is requested.
160    pub fn next_set(&mut self) -> Result<Option<ResultSet<'_>>> {
161        let Some(meta) = self.metas.next() else {
162            return Ok(None);
163        };
164        let ipc = match self.input.read()? {
165            Frame::Arrow(ipc) => ipc,
166            Frame::Json(v) => {
167                return Err(format!(
168                    "expected an Arrow frame starting result set `{}`, got {v}",
169                    meta.name
170                )
171                .into());
172            }
173        };
174        let (schema, batches) = frame::decode_batches(&ipc)?;
175        Ok(Some(ResultSet {
176            meta,
177            schema,
178            first: Some(batches),
179            input: &mut *self.input,
180            done: false,
181        }))
182    }
183}
184
185pub trait Format {
186    /// The options this format takes: every key of a report's `output:` block except `format`,
187    /// `destination`, `template` and `extension`, which core owns.
188    fn options(&self) -> Vec<OptionField> {
189        Vec::new()
190    }
191    /// Problems the declared [`Format::options`] can't catch, each a sentence naming the key.
192    /// Only called on options whose declared types already passed.
193    fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
194        Vec::new()
195    }
196    /// Write every result set to `req.path` (and siblings, if the format needs several files);
197    /// return the files written.
198    fn write(&mut self, req: &WriteRequest, sets: &mut ResultSets<'_>) -> Result<Vec<String>>;
199}
200
201/// A destination `deliver` request.
202pub struct Delivery {
203    /// One file, or every file of an output when the plugin advertises `multi_file`. With a
204    /// `message`: the files to attach to it (often none).
205    pub files: Vec<DeliveryFile>,
206    pub connection: Map<String, Value>,
207    /// The destination entry's plugin options, rendered by core (empty when none).
208    pub options: Map<String, Value>,
209    /// A message to post; only sent to a plugin advertising `message`.
210    pub message: Option<crate::msg::Message>,
211}
212
213/// One file to deliver: the local copy in `target/run/` and its rendered remote path, if any.
214pub struct DeliveryFile {
215    pub local: PathBuf,
216    pub remote: Option<String>,
217}
218
219pub trait Destination {
220    fn connection_fields(&self) -> Vec<ConnectionField> {
221        Vec::new()
222    }
223    /// The options a report's destination entry takes: every key but `profile` and `path`.
224    /// String values may still hold Jinja when checked (see [`options::is_template`]); core
225    /// renders it before `deliver`.
226    fn options(&self) -> Vec<OptionField> {
227        Vec::new()
228    }
229    /// Problems the declared [`Destination::options`] can't catch, each a sentence naming the key.
230    /// Only called on options whose declared types already passed.
231    fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
232        Vec::new()
233    }
234    /// Deliver `local` to `remote` (rendered by core); return where it landed. Enough for a
235    /// destination that takes no options; others implement [`Destination::deliver_files`].
236    fn deliver(
237        &mut self,
238        _local: &Path,
239        _remote: Option<&str>,
240        _connection: &Map<String, Value>,
241    ) -> Result<String> {
242        Err("this destination doesn't implement `deliver`".into())
243    }
244    /// The most characters a message may have in the service (after translation), reported in
245    /// `describe` by a plugin advertising `message`.
246    fn message_limit(&self) -> Option<u64> {
247        None
248    }
249    /// Post `d.message` (with `d.files` attached). Called only for a plugin advertising
250    /// `message`, which implements it.
251    fn deliver_message(&mut self, _d: &Delivery, _m: &crate::msg::Message) -> Result<String> {
252        Err("this destination doesn't take messages".into())
253    }
254    /// The whole request, options included (already checked). The default hands a single file
255    /// to [`Destination::deliver`]; several files arrive only with `multi_file` advertised.
256    fn deliver_files(&mut self, d: &Delivery) -> Result<String> {
257        match d.files.as_slice() {
258            [f] => self.deliver(&f.local, f.remote.as_deref(), &d.connection),
259            _ => Err("this destination takes one file per delivery".into()),
260        }
261    }
262}
263
264pub struct Input {
265    r: BufReader<Box<dyn Read>>,
266}
267
268impl Input {
269    fn read(&mut self) -> Result<Frame> {
270        frame::read_frame(&mut self.r).map_err(|e| -> Error {
271            match e {
272                FrameError::Eof => "core closed the connection".into(),
273                e => e.to_string().into(),
274            }
275        })
276    }
277}
278
279struct Output {
280    w: BufWriter<Stdout>,
281}
282
283impl Output {
284    fn send(&mut self, r: &Response) {
285        if frame::write_json(&mut self.w, r).is_err() {
286            // Core has gone; nothing left to talk to.
287            std::process::exit(1);
288        }
289    }
290
291    fn batch(&mut self, b: &RecordBatch) -> Result<()> {
292        let ipc = frame::encode_batch(b)?;
293        frame::write_arrow(&mut self.w, &ipc).map_err(|_| -> Error { "core closed the connection".into() })
294    }
295}
296
297enum Handler<'a> {
298    Source(&'a mut dyn Source),
299    Format(&'a mut dyn Format),
300    Destination(&'a mut dyn Destination),
301}
302
303impl Handler<'_> {
304    fn kind(&self) -> Kind {
305        match self {
306            Handler::Source(_) => Kind::Source,
307            Handler::Format(_) => Kind::Format,
308            Handler::Destination(_) => Kind::Destination,
309        }
310    }
311
312    fn option_fields(&self) -> Vec<OptionField> {
313        match self {
314            Handler::Source(_) => Vec::new(),
315            Handler::Format(f) => f.options(),
316            Handler::Destination(d) => d.options(),
317        }
318    }
319
320    /// Every problem with a config block: the declared checks, then the plugin's own once the
321    /// declared ones pass.
322    fn check(&self, name: &str, o: &Map<String, Value>) -> Vec<String> {
323        let errs = options::check(self.kind(), name, &self.option_fields(), o);
324        if !errs.is_empty() {
325            return errs;
326        }
327        match self {
328            Handler::Source(_) => Vec::new(),
329            Handler::Format(f) => f.validate(o),
330            Handler::Destination(d) => d.validate(o),
331        }
332    }
333
334    /// The config block's problems as one error, for `write` and `deliver`.
335    fn checked(&self, name: &str, o: &Map<String, Value>) -> Result<()> {
336        let errs = self.check(name, o);
337        if errs.is_empty() {
338            Ok(())
339        } else {
340            Err(errs.join("; ").into())
341        }
342    }
343}
344
345pub fn serve_source(about: About, s: impl Source + 'static) -> ! {
346    serve_package(vec![Plugin::Source(about, Box::new(s))])
347}
348
349pub fn serve_format(about: About, f: impl Format + 'static) -> ! {
350    serve_package(vec![Plugin::Format(about, Box::new(f))])
351}
352
353pub fn serve_destination(about: About, d: impl Destination + 'static) -> ! {
354    serve_package(vec![Plugin::Destination(about, Box::new(d))])
355}
356
357/// One of the plugins a package's executable serves (see [`serve_package`]).
358pub enum Plugin {
359    Source(About, Box<dyn Source>),
360    Format(About, Box<dyn Format>),
361    Destination(About, Box<dyn Destination>),
362}
363
364impl Plugin {
365    fn id(&self) -> PluginId {
366        let (kind, about) = match self {
367            Plugin::Source(a, _) => (Kind::Source, a),
368            Plugin::Format(a, _) => (Kind::Format, a),
369            Plugin::Destination(a, _) => (Kind::Destination, a),
370        };
371        PluginId::new(kind, about.name)
372    }
373
374    fn handler(&mut self) -> (&About, Handler<'_>) {
375        match self {
376            Plugin::Source(a, s) => (a, Handler::Source(s.as_mut())),
377            Plugin::Format(a, f) => (a, Handler::Format(f.as_mut())),
378            Plugin::Destination(a, d) => (a, Handler::Destination(d.as_mut())),
379        }
380    }
381}
382
383/// Serve every plugin of a package from one executable. Core's `hello` names the plugin it
384/// wants; without one, the first is served.
385pub fn serve_package(mut plugins: Vec<Plugin>) -> ! {
386    assert!(!plugins.is_empty(), "a package serves at least one plugin");
387    let mut input = Input {
388        r: BufReader::new(Box::new(std::io::stdin())),
389    };
390    let mut out = Output {
391        w: BufWriter::new(std::io::stdout()),
392    };
393    let provides: Vec<PluginId> = plugins.iter().map(Plugin::id).collect();
394    let chosen = loop {
395        let req = match next_request(&mut input, &mut out) {
396            Some(r) => r,
397            None => continue,
398        };
399        let Request::Hello {
400            min_version,
401            max_version,
402            plugin,
403            ..
404        } = req
405        else {
406            out.send(&Response::Error {
407                message: "the first request must be `hello`".into(),
408            });
409            continue;
410        };
411        let Some(chosen) = negotiate((min_version, max_version), (MIN_VERSION, MAX_VERSION)) else {
412            out.send(&Response::VersionMismatch {
413                min_version: MIN_VERSION,
414                max_version: MAX_VERSION,
415            });
416            std::process::exit(1);
417        };
418        let i = match &plugin {
419            None => 0,
420            Some(want) => match provides.iter().position(|p| p == want) {
421                Some(i) => i,
422                None => {
423                    let list: Vec<String> = provides.iter().map(|p| p.to_string()).collect();
424                    out.send(&Response::Error {
425                        message: format!("this executable provides {}, not {want}", list.join(", ")),
426                    });
427                    std::process::exit(1);
428                }
429            },
430        };
431        let (about, h) = plugins[i].handler();
432        out.send(&Response::Hello {
433            protocol_version: chosen,
434            kind: h.kind(),
435            name: about.name.to_string(),
436            version: about.version.to_string(),
437            capabilities: about
438                .capabilities
439                .iter()
440                .chain(std::iter::once(&CAP_VALIDATE))
441                .map(|c| c.to_string())
442                .collect(),
443            provides: if provides.len() > 1 {
444                provides.clone()
445            } else {
446                Vec::new()
447            },
448        });
449        break i;
450    };
451    let (about, h) = plugins[chosen].handler();
452    serve(about.name, h, input, out)
453}
454
455/// The next request, or `None` after answering one that can't be read. Exits at end of input.
456fn next_request(input: &mut Input, out: &mut Output) -> Option<Request> {
457    let frame = match frame::read_frame(&mut input.r) {
458        Ok(f) => f,
459        Err(FrameError::Eof) => std::process::exit(0),
460        Err(e) => {
461            eprintln!("{e}");
462            out.send(&Response::Error {
463                message: e.to_string(),
464            });
465            std::process::exit(2);
466        }
467    };
468    match frame {
469        Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
470            Ok(r) => Some(r),
471            Err(_) => {
472                let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
473                out.send(&Response::Error {
474                    message: format!("unsupported request `{t}`"),
475                });
476                None
477            }
478        },
479        Frame::Arrow(_) => {
480            out.send(&Response::Error {
481                message: "unexpected Arrow frame".into(),
482            });
483            None
484        }
485    }
486}
487
488/// Serve requests after the handshake.
489fn serve(name: &str, mut h: Handler<'_>, mut input: Input, mut out: Output) -> ! {
490    loop {
491        let frame = match frame::read_frame(&mut input.r) {
492            Ok(f) => f,
493            Err(FrameError::Eof) => {
494                if let Handler::Source(src) = &mut h {
495                    src.close();
496                }
497                std::process::exit(0)
498            }
499            Err(e) => {
500                eprintln!("{e}");
501                out.send(&Response::Error {
502                    message: e.to_string(),
503                });
504                std::process::exit(2);
505            }
506        };
507        let req = match frame {
508            Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
509                Ok(r) => r,
510                Err(_) => {
511                    let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
512                    out.send(&Response::Error {
513                        message: format!("unsupported request `{t}`"),
514                    });
515                    continue;
516                }
517            },
518            Frame::Arrow(_) => {
519                out.send(&Response::Error {
520                    message: "unexpected Arrow frame".into(),
521                });
522                continue;
523            }
524        };
525        if let Request::Hello { .. } = req {
526            out.send(&Response::Error {
527                message: "`hello` was already answered".into(),
528            });
529            continue;
530        }
531        if let Request::Close {} = req {
532            if let Handler::Source(src) = &mut h {
533                src.close();
534            }
535            out.send(&Response::Ok {});
536            std::process::exit(0);
537        }
538        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
539            handle(&mut h, name, req, &mut input, &mut out)
540        }));
541        match result {
542            Ok(Ok(())) => {}
543            Ok(Err(e)) => out.send(&Response::Error {
544                message: e.to_string(),
545            }),
546            Err(p) => {
547                let msg = p
548                    .downcast_ref::<String>()
549                    .cloned()
550                    .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()))
551                    .unwrap_or_else(|| "unknown panic".into());
552                out.send(&Response::Error {
553                    message: format!("plugin panicked: {msg}"),
554                });
555                std::process::exit(101);
556            }
557        }
558    }
559}
560
561fn handle(h: &mut Handler<'_>, name: &str, req: Request, input: &mut Input, out: &mut Output) -> Result<()> {
562    if let Request::Describe {} = req {
563        let connection_fields = match h {
564            Handler::Source(s) => s.connection_fields(),
565            Handler::Destination(d) => d.connection_fields(),
566            Handler::Format(_) => Vec::new(),
567        };
568        let identifier_quote = match h {
569            Handler::Source(s) => s.identifier_quote().map(str::to_string),
570            _ => None,
571        };
572        let message_limit = match h {
573            Handler::Destination(d) => d.message_limit(),
574            _ => None,
575        };
576        out.send(&Response::Describe {
577            connection_fields,
578            option_fields: h.option_fields(),
579            identifier_quote,
580            message_limit,
581        });
582        return Ok(());
583    }
584    if let Request::Validate { options } = req {
585        out.send(&Response::Validated {
586            errors: h.check(name, &options),
587        });
588        return Ok(());
589    }
590    let checked = match &req {
591        Request::Write { options, .. } | Request::Deliver { options, .. } => h.checked(name, options),
592        _ => Ok(()),
593    };
594    match (h, req) {
595        (
596            Handler::Source(s),
597            Request::Open {
598                connection,
599                read_only,
600            },
601        ) => {
602            s.open(&connection, read_only)?;
603            out.send(&Response::Ok {});
604        }
605        (Handler::Source(s), Request::Load { name }) => {
606            let ipc = match input.read()? {
607                Frame::Arrow(ipc) => ipc,
608                Frame::Json(v) => return Err(format!("expected the rows to load, got {v}").into()),
609            };
610            let (schema, batches) = frame::decode_batches(&ipc)?;
611            let mut data = ResultSet {
612                meta: ResultSetMeta {
613                    name: name.clone(),
614                    query: name.clone(),
615                    result_index: 1,
616                    anchor: None,
617                    header: None,
618                    columns: Default::default(),
619                },
620                schema,
621                first: Some(batches),
622                input,
623                done: false,
624            };
625            let r = s.load(&name, &mut data);
626            data.drain()?;
627            let loaded = r?;
628            out.send(&Response::Loaded {
629                relation: loaded.relation,
630                rows: loaded.rows,
631                warning: loaded.warning,
632            });
633        }
634        (Handler::Source(s), Request::Check { sql }) => {
635            s.check(&sql)?;
636            out.send(&Response::Ok {});
637        }
638        (Handler::Source(s), Request::Execute { sql, row_limit }) => {
639            let mut sink = FrameSink {
640                out,
641                row_limit,
642                rows: 0,
643                state: SinkState::Idle,
644                schema: None,
645                sent_any: false,
646            };
647            let r = s.execute(&sql, row_limit, &mut sink);
648            match (r, sink.state) {
649                (Err(e), _) => return Err(e),
650                (Ok(()), SinkState::Idle) => {
651                    return Err("the source produced neither a result nor `no_result`".into());
652                }
653                (Ok(()), SinkState::NoResult) => {}
654                (Ok(()), SinkState::Result) => {
655                    if !sink.sent_any {
656                        let schema = sink.schema.clone().unwrap();
657                        sink.out.batch(&RecordBatch::new_empty(schema))?;
658                    }
659                    let rows = sink.rows;
660                    sink.out.send(&Response::ResultEnd { rows });
661                }
662            }
663        }
664        (
665            Handler::Format(f),
666            Request::Write {
667                path,
668                format,
669                options,
670                result_sets,
671                template,
672            },
673        ) => {
674            let req = WriteRequest {
675                path,
676                format,
677                options,
678                result_sets: result_sets.clone(),
679                template,
680            };
681            let mut sets = ResultSets {
682                metas: result_sets.into_iter(),
683                input,
684                warnings: Vec::new(),
685            };
686            let written = checked.and_then(|()| f.write(&req, &mut sets));
687            // A failed write is reported at once, so core can stop streaming; this is the
688            // request's one reply.
689            if let Err(e) = &written {
690                out.send(&Response::Error {
691                    message: e.to_string(),
692                });
693            }
694            // Consume whatever the format didn't read (after an error, possibly the rest of a
695            // result set), so the stream stays in sync.
696            loop {
697                match sets.input.read()? {
698                    Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("finish") => break,
699                    Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("result_set_end") => {}
700                    Frame::Arrow(_) => {}
701                    other => return Err(format!("expected `finish`, got {other:?}").into()),
702                }
703            }
704            if let Ok(files) = written {
705                out.send(&Response::Written {
706                    files,
707                    warnings: sets.warnings,
708                });
709            }
710        }
711        (
712            Handler::Destination(d),
713            Request::Deliver {
714                local_path,
715                remote_path,
716                files,
717                connection,
718                options,
719                message,
720            },
721        ) => {
722            checked?;
723            let files = match (local_path, files.is_empty(), &message) {
724                (Some(local), true, None) => vec![DeliveryFile {
725                    local: PathBuf::from(local),
726                    remote: remote_path,
727                }],
728                (None, false, _) | (None, true, Some(_)) => files
729                    .into_iter()
730                    .map(|f| DeliveryFile {
731                        local: PathBuf::from(f.local_path),
732                        remote: f.remote_path,
733                    })
734                    .collect(),
735                _ => return Err("`deliver` needs exactly one of `local_path` or `files`".into()),
736            };
737            let delivery = Delivery {
738                files,
739                connection,
740                options,
741                message,
742            };
743            let location = match &delivery.message {
744                Some(m) => d.deliver_message(&delivery, m)?,
745                None => d.deliver_files(&delivery)?,
746            };
747            out.send(&Response::Delivered { location });
748        }
749        (h, req) => {
750            let t = serde_json::to_value(&req)
751                .ok()
752                .and_then(|v| v.get("type").cloned())
753                .unwrap_or_default();
754            return Err(format!("a {} plugin doesn't handle {t} requests", h.kind()).into());
755        }
756    }
757    Ok(())
758}
759
760#[derive(PartialEq)]
761enum SinkState {
762    Idle,
763    NoResult,
764    Result,
765}
766
767struct FrameSink<'a> {
768    out: &'a mut Output,
769    row_limit: Option<u64>,
770    rows: u64,
771    state: SinkState,
772    schema: Option<SchemaRef>,
773    sent_any: bool,
774}
775
776impl ResultSink for FrameSink<'_> {
777    fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()> {
778        if self.state != SinkState::Idle {
779            return Err("`no_result` after the result started".into());
780        }
781        self.state = SinkState::NoResult;
782        self.out.send(&Response::NoResult { rows_affected });
783        Ok(())
784    }
785
786    fn begin(&mut self, schema: SchemaRef) -> Result<()> {
787        if self.state != SinkState::Idle {
788            return Err("a statement can only produce one result".into());
789        }
790        self.state = SinkState::Result;
791        self.out.send(&Response::Result {
792            columns: schema.fields().iter().map(|f| f.name().clone()).collect(),
793        });
794        self.schema = Some(schema);
795        Ok(())
796    }
797
798    fn batch(&mut self, b: RecordBatch) -> Result<bool> {
799        if self.state != SinkState::Result {
800            return Err("`batch` before `begin`".into());
801        }
802        if self.row_limit.is_some_and(|l| self.rows >= l) {
803            return Ok(false);
804        }
805        let b = match self.row_limit {
806            Some(l) if self.rows + b.num_rows() as u64 > l => b.slice(0, (l - self.rows) as usize),
807            _ => b,
808        };
809        if b.num_rows() > 0 || !self.sent_any {
810            self.rows += b.num_rows() as u64;
811            self.out.batch(&b)?;
812            self.sent_any = true;
813        }
814        Ok(self.row_limit.is_none_or(|l| self.rows < l))
815    }
816}
817
818/// The highest version in both ranges.
819fn negotiate(core: (u32, u32), ours: (u32, u32)) -> Option<u32> {
820    let (lo, hi) = (core.0.max(ours.0), core.1.min(ours.1));
821    (lo <= hi).then_some(hi)
822}
823
824/// Read a string field from a connection map.
825pub fn conn_str<'a>(c: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
826    c.get(key).and_then(Value::as_str).filter(|s| !s.is_empty())
827}
828
829/// Read a required string field from a connection map.
830pub fn conn_required<'a>(c: &'a Map<String, Value>, key: &str) -> Result<&'a str> {
831    conn_str(c, key).ok_or_else(|| format!("the profile output needs a `{key}` field").into())
832}
833
834/// Read a flag that may be written as a boolean or a string.
835pub fn conn_bool(c: &Map<String, Value>, key: &str) -> Option<bool> {
836    match c.get(key)? {
837        Value::Bool(b) => Some(*b),
838        Value::String(s) => Some(matches!(s.as_str(), "true" | "yes" | "1")),
839        _ => None,
840    }
841}