1use 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#[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 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
49pub trait ResultSink {
52 fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()>;
53 fn begin(&mut self, schema: SchemaRef) -> Result<()>;
54 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 fn identifier_quote(&self) -> Option<&'static str> {
65 None
66 }
67 fn open(&mut self, connection: &Map<String, Value>, read_only: bool) -> Result<()>;
68 fn execute(&mut self, sql: &str, row_limit: Option<u64>, out: &mut dyn ResultSink) -> Result<()>;
70 fn check(&mut self, _sql: &str) -> Result<()> {
72 Err("this source can't check statements".into())
73 }
74 fn load(&mut self, _name: &str, _data: &mut ResultSet<'_>) -> Result<Loaded> {
77 Err("this source can't load rows".into())
78 }
79 fn close(&mut self) {}
81}
82
83pub struct Loaded {
85 pub relation: String,
87 pub rows: u64,
88 pub warning: Option<String>,
90}
91
92pub 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 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
136pub 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
145pub struct ResultSets<'a> {
147 metas: std::vec::IntoIter<ResultSetMeta>,
148 input: &'a mut Input,
149 warnings: Vec<String>,
150}
151
152impl ResultSets<'_> {
153 pub fn warn(&mut self, message: impl Into<String>) {
155 self.warnings.push(message.into());
156 }
157
158 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 fn options(&self) -> Vec<OptionField> {
189 Vec::new()
190 }
191 fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
194 Vec::new()
195 }
196 fn write(&mut self, req: &WriteRequest, sets: &mut ResultSets<'_>) -> Result<Vec<String>>;
199}
200
201pub struct Delivery {
203 pub files: Vec<DeliveryFile>,
206 pub connection: Map<String, Value>,
207 pub options: Map<String, Value>,
209 pub message: Option<crate::msg::Message>,
211}
212
213pub 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 fn options(&self) -> Vec<OptionField> {
227 Vec::new()
228 }
229 fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
232 Vec::new()
233 }
234 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 fn message_limit(&self) -> Option<u64> {
247 None
248 }
249 fn deliver_message(&mut self, _d: &Delivery, _m: &crate::msg::Message) -> Result<String> {
252 Err("this destination doesn't take messages".into())
253 }
254 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 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 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 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
357pub 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
383pub 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
455fn 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
488fn 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 if let Err(e) = &written {
690 out.send(&Response::Error {
691 message: e.to_string(),
692 });
693 }
694 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
818fn 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
824pub 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
829pub 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
834pub 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}