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>,
205 pub connection: Map<String, Value>,
206 pub options: Map<String, Value>,
208}
209
210pub 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 fn options(&self) -> Vec<OptionField> {
224 Vec::new()
225 }
226 fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
229 Vec::new()
230 }
231 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 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 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 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 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
344pub 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
370pub 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
442fn 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
475fn 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 if let Err(e) = &written {
672 out.send(&Response::Error {
673 message: e.to_string(),
674 });
675 }
676 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
794fn 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
800pub 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
805pub 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
810pub 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}