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`.
204    pub files: Vec<DeliveryFile>,
205    pub connection: Map<String, Value>,
206    /// The destination entry's plugin options, rendered by core (empty when none).
207    pub options: Map<String, Value>,
208}
209
210/// One file to deliver: the local copy in `target/run/` and its rendered remote path, if any.
211pub struct DeliveryFile {
212    pub local: PathBuf,
213    pub remote: Option<String>,
214}
215
216pub trait Destination {
217    fn connection_fields(&self) -> Vec<ConnectionField> {
218        Vec::new()
219    }
220    /// The options a report's destination entry takes: every key but `profile` and `path`.
221    /// String values may still hold Jinja when checked (see [`options::is_template`]); core
222    /// renders it before `deliver`.
223    fn options(&self) -> Vec<OptionField> {
224        Vec::new()
225    }
226    /// Problems the declared [`Destination::options`] can't catch, each a sentence naming the key.
227    /// Only called on options whose declared types already passed.
228    fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
229        Vec::new()
230    }
231    /// Deliver `local` to `remote` (rendered by core); return where it landed. Enough for a
232    /// destination that takes no options; others implement [`Destination::deliver_files`].
233    fn deliver(
234        &mut self,
235        _local: &Path,
236        _remote: Option<&str>,
237        _connection: &Map<String, Value>,
238    ) -> Result<String> {
239        Err("this destination doesn't implement `deliver`".into())
240    }
241    /// The whole request, options included (already checked). The default hands a single file
242    /// to [`Destination::deliver`]; several files arrive only with `multi_file` advertised.
243    fn deliver_files(&mut self, d: &Delivery) -> Result<String> {
244        match d.files.as_slice() {
245            [f] => self.deliver(&f.local, f.remote.as_deref(), &d.connection),
246            _ => Err("this destination takes one file per delivery".into()),
247        }
248    }
249}
250
251pub struct Input {
252    r: BufReader<Box<dyn Read>>,
253}
254
255impl Input {
256    fn read(&mut self) -> Result<Frame> {
257        frame::read_frame(&mut self.r).map_err(|e| -> Error {
258            match e {
259                FrameError::Eof => "core closed the connection".into(),
260                e => e.to_string().into(),
261            }
262        })
263    }
264}
265
266struct Output {
267    w: BufWriter<Stdout>,
268}
269
270impl Output {
271    fn send(&mut self, r: &Response) {
272        if frame::write_json(&mut self.w, r).is_err() {
273            // Core has gone; nothing left to talk to.
274            std::process::exit(1);
275        }
276    }
277
278    fn batch(&mut self, b: &RecordBatch) -> Result<()> {
279        let ipc = frame::encode_batch(b)?;
280        frame::write_arrow(&mut self.w, &ipc).map_err(|_| -> Error { "core closed the connection".into() })
281    }
282}
283
284enum Handler<'a> {
285    Source(&'a mut dyn Source),
286    Format(&'a mut dyn Format),
287    Destination(&'a mut dyn Destination),
288}
289
290impl Handler<'_> {
291    fn kind(&self) -> Kind {
292        match self {
293            Handler::Source(_) => Kind::Source,
294            Handler::Format(_) => Kind::Format,
295            Handler::Destination(_) => Kind::Destination,
296        }
297    }
298
299    fn option_fields(&self) -> Vec<OptionField> {
300        match self {
301            Handler::Source(_) => Vec::new(),
302            Handler::Format(f) => f.options(),
303            Handler::Destination(d) => d.options(),
304        }
305    }
306
307    /// Every problem with a config block: the declared checks, then the plugin's own once the
308    /// declared ones pass.
309    fn check(&self, name: &str, o: &Map<String, Value>) -> Vec<String> {
310        let errs = options::check(self.kind(), name, &self.option_fields(), o);
311        if !errs.is_empty() {
312            return errs;
313        }
314        match self {
315            Handler::Source(_) => Vec::new(),
316            Handler::Format(f) => f.validate(o),
317            Handler::Destination(d) => d.validate(o),
318        }
319    }
320
321    /// The config block's problems as one error, for `write` and `deliver`.
322    fn checked(&self, name: &str, o: &Map<String, Value>) -> Result<()> {
323        let errs = self.check(name, o);
324        if errs.is_empty() {
325            Ok(())
326        } else {
327            Err(errs.join("; ").into())
328        }
329    }
330}
331
332pub fn serve_source(about: About, s: impl Source + 'static) -> ! {
333    serve_package(vec![Plugin::Source(about, Box::new(s))])
334}
335
336pub fn serve_format(about: About, f: impl Format + 'static) -> ! {
337    serve_package(vec![Plugin::Format(about, Box::new(f))])
338}
339
340pub fn serve_destination(about: About, d: impl Destination + 'static) -> ! {
341    serve_package(vec![Plugin::Destination(about, Box::new(d))])
342}
343
344/// One of the plugins a package's executable serves (see [`serve_package`]).
345pub enum Plugin {
346    Source(About, Box<dyn Source>),
347    Format(About, Box<dyn Format>),
348    Destination(About, Box<dyn Destination>),
349}
350
351impl Plugin {
352    fn id(&self) -> PluginId {
353        let (kind, about) = match self {
354            Plugin::Source(a, _) => (Kind::Source, a),
355            Plugin::Format(a, _) => (Kind::Format, a),
356            Plugin::Destination(a, _) => (Kind::Destination, a),
357        };
358        PluginId::new(kind, about.name)
359    }
360
361    fn handler(&mut self) -> (&About, Handler<'_>) {
362        match self {
363            Plugin::Source(a, s) => (a, Handler::Source(s.as_mut())),
364            Plugin::Format(a, f) => (a, Handler::Format(f.as_mut())),
365            Plugin::Destination(a, d) => (a, Handler::Destination(d.as_mut())),
366        }
367    }
368}
369
370/// Serve every plugin of a package from one executable. Core's `hello` names the plugin it
371/// wants; without one, the first is served.
372pub fn serve_package(mut plugins: Vec<Plugin>) -> ! {
373    assert!(!plugins.is_empty(), "a package serves at least one plugin");
374    let mut input = Input {
375        r: BufReader::new(Box::new(std::io::stdin())),
376    };
377    let mut out = Output {
378        w: BufWriter::new(std::io::stdout()),
379    };
380    let provides: Vec<PluginId> = plugins.iter().map(Plugin::id).collect();
381    let chosen = loop {
382        let req = match next_request(&mut input, &mut out) {
383            Some(r) => r,
384            None => continue,
385        };
386        let Request::Hello {
387            min_version,
388            max_version,
389            plugin,
390            ..
391        } = req
392        else {
393            out.send(&Response::Error {
394                message: "the first request must be `hello`".into(),
395            });
396            continue;
397        };
398        let Some(chosen) = negotiate((min_version, max_version), (MIN_VERSION, MAX_VERSION)) else {
399            out.send(&Response::VersionMismatch {
400                min_version: MIN_VERSION,
401                max_version: MAX_VERSION,
402            });
403            std::process::exit(1);
404        };
405        let i = match &plugin {
406            None => 0,
407            Some(want) => match provides.iter().position(|p| p == want) {
408                Some(i) => i,
409                None => {
410                    let list: Vec<String> = provides.iter().map(|p| p.to_string()).collect();
411                    out.send(&Response::Error {
412                        message: format!("this executable provides {}, not {want}", list.join(", ")),
413                    });
414                    std::process::exit(1);
415                }
416            },
417        };
418        let (about, h) = plugins[i].handler();
419        out.send(&Response::Hello {
420            protocol_version: chosen,
421            kind: h.kind(),
422            name: about.name.to_string(),
423            version: about.version.to_string(),
424            capabilities: about
425                .capabilities
426                .iter()
427                .chain(std::iter::once(&CAP_VALIDATE))
428                .map(|c| c.to_string())
429                .collect(),
430            provides: if provides.len() > 1 {
431                provides.clone()
432            } else {
433                Vec::new()
434            },
435        });
436        break i;
437    };
438    let (about, h) = plugins[chosen].handler();
439    serve(about.name, h, input, out)
440}
441
442/// The next request, or `None` after answering one that can't be read. Exits at end of input.
443fn next_request(input: &mut Input, out: &mut Output) -> Option<Request> {
444    let frame = match frame::read_frame(&mut input.r) {
445        Ok(f) => f,
446        Err(FrameError::Eof) => std::process::exit(0),
447        Err(e) => {
448            eprintln!("{e}");
449            out.send(&Response::Error {
450                message: e.to_string(),
451            });
452            std::process::exit(2);
453        }
454    };
455    match frame {
456        Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
457            Ok(r) => Some(r),
458            Err(_) => {
459                let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
460                out.send(&Response::Error {
461                    message: format!("unsupported request `{t}`"),
462                });
463                None
464            }
465        },
466        Frame::Arrow(_) => {
467            out.send(&Response::Error {
468                message: "unexpected Arrow frame".into(),
469            });
470            None
471        }
472    }
473}
474
475/// Serve requests after the handshake.
476fn serve(name: &str, mut h: Handler<'_>, mut input: Input, mut out: Output) -> ! {
477    loop {
478        let frame = match frame::read_frame(&mut input.r) {
479            Ok(f) => f,
480            Err(FrameError::Eof) => {
481                if let Handler::Source(src) = &mut h {
482                    src.close();
483                }
484                std::process::exit(0)
485            }
486            Err(e) => {
487                eprintln!("{e}");
488                out.send(&Response::Error {
489                    message: e.to_string(),
490                });
491                std::process::exit(2);
492            }
493        };
494        let req = match frame {
495            Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
496                Ok(r) => r,
497                Err(_) => {
498                    let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
499                    out.send(&Response::Error {
500                        message: format!("unsupported request `{t}`"),
501                    });
502                    continue;
503                }
504            },
505            Frame::Arrow(_) => {
506                out.send(&Response::Error {
507                    message: "unexpected Arrow frame".into(),
508                });
509                continue;
510            }
511        };
512        if let Request::Hello { .. } = req {
513            out.send(&Response::Error {
514                message: "`hello` was already answered".into(),
515            });
516            continue;
517        }
518        if let Request::Close {} = req {
519            if let Handler::Source(src) = &mut h {
520                src.close();
521            }
522            out.send(&Response::Ok {});
523            std::process::exit(0);
524        }
525        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
526            handle(&mut h, name, req, &mut input, &mut out)
527        }));
528        match result {
529            Ok(Ok(())) => {}
530            Ok(Err(e)) => out.send(&Response::Error {
531                message: e.to_string(),
532            }),
533            Err(p) => {
534                let msg = p
535                    .downcast_ref::<String>()
536                    .cloned()
537                    .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()))
538                    .unwrap_or_else(|| "unknown panic".into());
539                out.send(&Response::Error {
540                    message: format!("plugin panicked: {msg}"),
541                });
542                std::process::exit(101);
543            }
544        }
545    }
546}
547
548fn handle(h: &mut Handler<'_>, name: &str, req: Request, input: &mut Input, out: &mut Output) -> Result<()> {
549    if let Request::Describe {} = req {
550        let connection_fields = match h {
551            Handler::Source(s) => s.connection_fields(),
552            Handler::Destination(d) => d.connection_fields(),
553            Handler::Format(_) => Vec::new(),
554        };
555        let identifier_quote = match h {
556            Handler::Source(s) => s.identifier_quote().map(str::to_string),
557            _ => None,
558        };
559        out.send(&Response::Describe {
560            connection_fields,
561            option_fields: h.option_fields(),
562            identifier_quote,
563        });
564        return Ok(());
565    }
566    if let Request::Validate { options } = req {
567        out.send(&Response::Validated {
568            errors: h.check(name, &options),
569        });
570        return Ok(());
571    }
572    let checked = match &req {
573        Request::Write { options, .. } | Request::Deliver { options, .. } => h.checked(name, options),
574        _ => Ok(()),
575    };
576    match (h, req) {
577        (
578            Handler::Source(s),
579            Request::Open {
580                connection,
581                read_only,
582            },
583        ) => {
584            s.open(&connection, read_only)?;
585            out.send(&Response::Ok {});
586        }
587        (Handler::Source(s), Request::Load { name }) => {
588            let ipc = match input.read()? {
589                Frame::Arrow(ipc) => ipc,
590                Frame::Json(v) => return Err(format!("expected the rows to load, got {v}").into()),
591            };
592            let (schema, batches) = frame::decode_batches(&ipc)?;
593            let mut data = ResultSet {
594                meta: ResultSetMeta {
595                    name: name.clone(),
596                    query: name.clone(),
597                    result_index: 1,
598                    anchor: None,
599                    header: None,
600                    columns: Default::default(),
601                },
602                schema,
603                first: Some(batches),
604                input,
605                done: false,
606            };
607            let r = s.load(&name, &mut data);
608            data.drain()?;
609            let loaded = r?;
610            out.send(&Response::Loaded {
611                relation: loaded.relation,
612                rows: loaded.rows,
613                warning: loaded.warning,
614            });
615        }
616        (Handler::Source(s), Request::Check { sql }) => {
617            s.check(&sql)?;
618            out.send(&Response::Ok {});
619        }
620        (Handler::Source(s), Request::Execute { sql, row_limit }) => {
621            let mut sink = FrameSink {
622                out,
623                row_limit,
624                rows: 0,
625                state: SinkState::Idle,
626                schema: None,
627                sent_any: false,
628            };
629            let r = s.execute(&sql, row_limit, &mut sink);
630            match (r, sink.state) {
631                (Err(e), _) => return Err(e),
632                (Ok(()), SinkState::Idle) => {
633                    return Err("the source produced neither a result nor `no_result`".into());
634                }
635                (Ok(()), SinkState::NoResult) => {}
636                (Ok(()), SinkState::Result) => {
637                    if !sink.sent_any {
638                        let schema = sink.schema.clone().unwrap();
639                        sink.out.batch(&RecordBatch::new_empty(schema))?;
640                    }
641                    let rows = sink.rows;
642                    sink.out.send(&Response::ResultEnd { rows });
643                }
644            }
645        }
646        (
647            Handler::Format(f),
648            Request::Write {
649                path,
650                format,
651                options,
652                result_sets,
653                template,
654            },
655        ) => {
656            let req = WriteRequest {
657                path,
658                format,
659                options,
660                result_sets: result_sets.clone(),
661                template,
662            };
663            let mut sets = ResultSets {
664                metas: result_sets.into_iter(),
665                input,
666                warnings: Vec::new(),
667            };
668            let written = checked.and_then(|()| f.write(&req, &mut sets));
669            // A failed write is reported at once, so core can stop streaming; this is the
670            // request's one reply.
671            if let Err(e) = &written {
672                out.send(&Response::Error {
673                    message: e.to_string(),
674                });
675            }
676            // Consume whatever the format didn't read (after an error, possibly the rest of a
677            // result set), so the stream stays in sync.
678            loop {
679                match sets.input.read()? {
680                    Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("finish") => break,
681                    Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("result_set_end") => {}
682                    Frame::Arrow(_) => {}
683                    other => return Err(format!("expected `finish`, got {other:?}").into()),
684                }
685            }
686            if let Ok(files) = written {
687                out.send(&Response::Written {
688                    files,
689                    warnings: sets.warnings,
690                });
691            }
692        }
693        (
694            Handler::Destination(d),
695            Request::Deliver {
696                local_path,
697                remote_path,
698                files,
699                connection,
700                options,
701            },
702        ) => {
703            checked?;
704            let files = match (local_path, files.is_empty()) {
705                (Some(local), true) => vec![DeliveryFile {
706                    local: PathBuf::from(local),
707                    remote: remote_path,
708                }],
709                (None, false) => files
710                    .into_iter()
711                    .map(|f| DeliveryFile {
712                        local: PathBuf::from(f.local_path),
713                        remote: f.remote_path,
714                    })
715                    .collect(),
716                _ => return Err("`deliver` needs exactly one of `local_path` or `files`".into()),
717            };
718            let location = d.deliver_files(&Delivery {
719                files,
720                connection,
721                options,
722            })?;
723            out.send(&Response::Delivered { location });
724        }
725        (h, req) => {
726            let t = serde_json::to_value(&req)
727                .ok()
728                .and_then(|v| v.get("type").cloned())
729                .unwrap_or_default();
730            return Err(format!("a {} plugin doesn't handle {t} requests", h.kind()).into());
731        }
732    }
733    Ok(())
734}
735
736#[derive(PartialEq)]
737enum SinkState {
738    Idle,
739    NoResult,
740    Result,
741}
742
743struct FrameSink<'a> {
744    out: &'a mut Output,
745    row_limit: Option<u64>,
746    rows: u64,
747    state: SinkState,
748    schema: Option<SchemaRef>,
749    sent_any: bool,
750}
751
752impl ResultSink for FrameSink<'_> {
753    fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()> {
754        if self.state != SinkState::Idle {
755            return Err("`no_result` after the result started".into());
756        }
757        self.state = SinkState::NoResult;
758        self.out.send(&Response::NoResult { rows_affected });
759        Ok(())
760    }
761
762    fn begin(&mut self, schema: SchemaRef) -> Result<()> {
763        if self.state != SinkState::Idle {
764            return Err("a statement can only produce one result".into());
765        }
766        self.state = SinkState::Result;
767        self.out.send(&Response::Result {
768            columns: schema.fields().iter().map(|f| f.name().clone()).collect(),
769        });
770        self.schema = Some(schema);
771        Ok(())
772    }
773
774    fn batch(&mut self, b: RecordBatch) -> Result<bool> {
775        if self.state != SinkState::Result {
776            return Err("`batch` before `begin`".into());
777        }
778        if self.row_limit.is_some_and(|l| self.rows >= l) {
779            return Ok(false);
780        }
781        let b = match self.row_limit {
782            Some(l) if self.rows + b.num_rows() as u64 > l => b.slice(0, (l - self.rows) as usize),
783            _ => b,
784        };
785        if b.num_rows() > 0 || !self.sent_any {
786            self.rows += b.num_rows() as u64;
787            self.out.batch(&b)?;
788            self.sent_any = true;
789        }
790        Ok(self.row_limit.is_none_or(|l| self.rows < l))
791    }
792}
793
794/// The highest version in both ranges.
795fn negotiate(core: (u32, u32), ours: (u32, u32)) -> Option<u32> {
796    let (lo, hi) = (core.0.max(ours.0), core.1.min(ours.1));
797    (lo <= hi).then_some(hi)
798}
799
800/// Read a string field from a connection map.
801pub fn conn_str<'a>(c: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
802    c.get(key).and_then(Value::as_str).filter(|s| !s.is_empty())
803}
804
805/// Read a required string field from a connection map.
806pub fn conn_required<'a>(c: &'a Map<String, Value>, key: &str) -> Result<&'a str> {
807    conn_str(c, key).ok_or_else(|| format!("the profile output needs a `{key}` field").into())
808}
809
810/// Read a flag that may be written as a boolean or a string.
811pub fn conn_bool(c: &Map<String, Value>, key: &str) -> Option<bool> {
812    match c.get(key)? {
813        Value::Bool(b) => Some(*b),
814        Value::String(s) => Some(matches!(s.as_str(), "true" | "yes" | "1")),
815        _ => None,
816    }
817}