1use std::io::{BufReader, BufWriter};
19use std::path::{Path, PathBuf};
20use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
21use std::sync::mpsc::{Receiver, Sender};
22use std::sync::{Arc, Mutex, OnceLock};
23
24use arrow::array::RecordBatch;
25use arrow::datatypes::SchemaRef;
26use serde_json::{Map, Value};
27
28use crate::frame::{self, Frame, FrameError};
29use crate::msg::{ConnectionField, Envelope, LogLevel, Request, Response, ResultSetMeta};
30use crate::options::{self, OptionField};
31use crate::{CAP_VALIDATE, Kind, MAX_VERSION, MIN_VERSION, PluginId};
32
33pub type Error = Box<dyn std::error::Error + Send + Sync>;
34pub type Result<T> = std::result::Result<T, Error>;
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum ErrorKind {
39 Config,
41 Plugin,
43 Refused,
45 Connection,
47 Auth,
49 Query,
51 Delivery,
53 Internal,
55 Cancelled,
56 TimedOut,
57}
58
59impl ErrorKind {
60 pub fn as_str(self) -> &'static str {
61 match self {
62 ErrorKind::Config => "config",
63 ErrorKind::Plugin => "plugin",
64 ErrorKind::Refused => "refused",
65 ErrorKind::Connection => "connection",
66 ErrorKind::Auth => "auth",
67 ErrorKind::Query => "query",
68 ErrorKind::Delivery => "delivery",
69 ErrorKind::Internal => "internal",
70 ErrorKind::Cancelled => "cancelled",
71 ErrorKind::TimedOut => "timed_out",
72 }
73 }
74}
75
76#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct PluginError {
80 pub kind: ErrorKind,
81 pub code: Option<String>,
82 pub message: String,
83}
84
85impl PluginError {
86 pub fn new(kind: ErrorKind, message: impl Into<String>) -> PluginError {
87 PluginError {
88 kind,
89 code: None,
90 message: message.into(),
91 }
92 }
93
94 #[must_use]
95 pub fn code(mut self, code: impl Into<String>) -> PluginError {
96 self.code = Some(code.into());
97 self
98 }
99}
100
101impl std::fmt::Display for PluginError {
102 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
103 f.write_str(&self.message)
104 }
105}
106
107impl std::error::Error for PluginError {}
108
109struct Shared {
112 out: Mutex<BufWriter<Box<dyn std::io::Write + Send>>>,
113 version: AtomicU32,
115 current: AtomicU64,
117 cancelled: AtomicU64,
119 hook: Mutex<Option<CancelHook>>,
120 name: OnceLock<String>,
122}
123
124type CancelHook = Arc<dyn Fn() + Send + Sync>;
125
126fn shared() -> &'static Shared {
127 static SHARED: OnceLock<Shared> = OnceLock::new();
128 SHARED.get_or_init(|| Shared {
129 out: Mutex::new(BufWriter::new(protocol_out())),
130 version: AtomicU32::new(0),
131 current: AtomicU64::new(0),
132 cancelled: AtomicU64::new(0),
133 hook: Mutex::new(None),
134 name: OnceLock::new(),
135 })
136}
137
138fn protocol_out() -> Box<dyn std::io::Write + Send> {
142 #[cfg(unix)]
143 {
144 use std::os::fd::FromRawFd;
145 unsafe {
147 let fd = libc::dup(1);
148 if fd >= 0 && libc::dup2(2, 1) >= 0 {
149 return Box::new(std::fs::File::from_raw_fd(fd));
150 }
151 }
152 }
153 #[cfg(windows)]
154 {
155 use std::os::windows::io::FromRawHandle;
156 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
157 use windows_sys::Win32::System::Console::{
158 GetStdHandle, STD_ERROR_HANDLE, STD_OUTPUT_HANDLE, SetStdHandle,
159 };
160 unsafe {
163 let out = GetStdHandle(STD_OUTPUT_HANDLE);
164 let err = GetStdHandle(STD_ERROR_HANDLE);
165 if !out.is_null() && out != INVALID_HANDLE_VALUE && SetStdHandle(STD_OUTPUT_HANDLE, err) != 0 {
166 return Box::new(std::fs::File::from_raw_handle(out as _));
167 }
168 }
169 }
170 Box::new(std::io::stdout())
171}
172
173pub fn cancelled() -> bool {
175 let s = shared();
176 let current = s.current.load(Ordering::SeqCst);
177 current != 0 && s.cancelled.load(Ordering::SeqCst) == current
178}
179
180pub fn on_cancel(hook: impl Fn() + Send + Sync + 'static) {
188 *shared().hook.lock().unwrap() = Some(Arc::new(hook));
189}
190
191fn cancel(id: u64) {
194 let s = shared();
195 s.cancelled.store(id, Ordering::SeqCst);
196 if id != 0 && s.current.load(Ordering::SeqCst) == id {
197 let hook = s.hook.lock().unwrap().clone();
198 if let Some(hook) = hook {
199 hook();
200 }
201 }
202}
203
204pub fn log_fields(level: LogLevel, message: &str, fields: Map<String, Value>) {
206 let s = shared();
207 if s.version.load(Ordering::SeqCst) == 0 {
208 let prefix = match level {
210 LogLevel::Error | LogLevel::Warn => "warning: ",
211 LogLevel::Info => "info: ",
212 LogLevel::Debug | LogLevel::Trace => "",
213 };
214 eprintln!("{prefix}{message}");
215 return;
216 }
217 send(&Response::Log {
218 level,
219 message: message.to_string(),
220 fields,
221 });
222}
223
224pub fn log(level: LogLevel, message: &str) {
226 log_fields(level, message, Map::new())
227}
228
229pub fn progress(message: Option<&str>, done: Option<u64>, total: Option<u64>) {
231 if shared().version.load(Ordering::SeqCst) == 0 {
232 if let Some(m) = message {
233 log(LogLevel::Info, m);
234 }
235 return;
236 }
237 send(&Response::Progress {
238 message: message.map(str::to_string),
239 done,
240 total,
241 });
242}
243
244struct FrameLogger;
246
247impl ::log::Log for FrameLogger {
248 fn enabled(&self, m: &::log::Metadata) -> bool {
251 m.level() <= ::log::Level::Info || (m.level() <= ::log::Level::Debug && m.target().starts_with("dre"))
252 }
253
254 fn log(&self, r: &::log::Record) {
255 if !self.enabled(r.metadata()) {
256 return;
257 }
258 let level = match r.level() {
259 ::log::Level::Error => LogLevel::Error,
260 ::log::Level::Warn => LogLevel::Warn,
261 ::log::Level::Info => LogLevel::Info,
262 ::log::Level::Debug => LogLevel::Debug,
263 ::log::Level::Trace => LogLevel::Trace,
264 };
265 let level = if r.target().starts_with("dre") {
267 level
268 } else {
269 level.max(LogLevel::Debug)
270 };
271 log(level, &r.args().to_string());
272 }
273
274 fn flush(&self) {}
275}
276
277fn send(r: &Response) {
279 let s = shared();
280 let id = match s.version.load(Ordering::SeqCst) {
281 0 => None,
282 _ => Some(s.current.load(Ordering::SeqCst)).filter(|id| *id != 0),
283 };
284 let mut out = s.out.lock().unwrap();
285 if frame::write_json(&mut *out, &Envelope::new(id, r)).is_err() {
286 std::process::exit(1);
288 }
289}
290
291fn error_reply(e: &Error) -> Response {
294 if cancelled() {
295 return Response::Error {
296 message: e.to_string(),
297 kind: Some("cancelled".into()),
298 code: None,
299 };
300 }
301 match e.downcast_ref::<PluginError>() {
302 Some(pe) => Response::Error {
303 message: pe.message.clone(),
304 kind: Some(pe.kind.as_str().into()),
305 code: pe.code.as_ref().map(|c| match shared().name.get() {
306 Some(name) if !c.contains('/') => format!("{name}/{c}"),
307 _ => c.clone(),
308 }),
309 },
310 None => Response::error(e.to_string()),
311 }
312}
313
314#[derive(Debug, Clone)]
316pub struct About {
317 pub name: &'static str,
318 pub version: &'static str,
319 pub capabilities: &'static [&'static str],
320}
321
322impl About {
323 pub fn new(name: &'static str, version: &'static str) -> About {
325 let version = if version == "0.0.0" { "unreleased" } else { version };
326 About {
327 name,
328 version,
329 capabilities: &[],
330 }
331 }
332
333 #[must_use]
334 pub fn capabilities(mut self, caps: &'static [&'static str]) -> About {
335 self.capabilities = caps;
336 self
337 }
338}
339
340pub trait ResultSink {
343 fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()>;
344 fn begin(&mut self, schema: SchemaRef) -> Result<()>;
345 fn batch(&mut self, batch: RecordBatch) -> Result<bool>;
347}
348
349pub trait Source {
350 fn connection_fields(&self) -> Vec<ConnectionField> {
351 Vec::new()
352 }
353 fn identifier_quote(&self) -> Option<&'static str> {
356 None
357 }
358 fn string_literal(&self) -> Option<&'static str> {
362 None
363 }
364 fn validate_connection(&self, _connection: &Map<String, Value>) -> Vec<String> {
369 Vec::new()
370 }
371 fn open(&mut self, connection: &Map<String, Value>, read_only: bool) -> Result<()>;
372 fn open_sandboxed(&mut self, connection: &Map<String, Value>) -> Result<()> {
375 self.open(connection, true)
376 }
377 fn execute(&mut self, sql: &str, row_limit: Option<u64>, out: &mut dyn ResultSink) -> Result<()>;
379 fn check(&mut self, _sql: &str) -> Result<()> {
381 Err("this source can't check statements".into())
382 }
383 fn load(&mut self, _name: &str, _data: &mut ResultSet<'_>) -> Result<Loaded> {
386 Err("this source can't load rows".into())
387 }
388 fn close(&mut self) {}
390}
391
392pub struct Loaded {
394 pub relation: String,
396 pub rows: u64,
397 pub warning: Option<String>,
399}
400
401pub struct ResultSet<'a> {
403 pub meta: ResultSetMeta,
404 pub schema: SchemaRef,
405 first: Option<Vec<RecordBatch>>,
406 input: &'a mut Input,
407 done: bool,
408}
409
410impl ResultSet<'_> {
411 pub fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
413 loop {
414 if let Some(first) = self.first.as_mut() {
415 if !first.is_empty() {
416 return Ok(Some(first.remove(0)));
417 }
418 self.first = None;
419 }
420 if self.done {
421 return Ok(None);
422 }
423 match self.input.read()? {
424 Frame::Arrow(ipc) => {
425 let (_, batches) = frame::decode_batches(&ipc)?;
426 self.first = Some(batches);
427 }
428 Frame::Json(v) => match serde_json::from_value::<Request>(v)? {
429 Request::ResultSetEnd {} => {
430 self.done = true;
431 return Ok(None);
432 }
433 other => return Err(format!("expected result data, got {other:?}").into()),
434 },
435 }
436 }
437 }
438
439 fn drain(&mut self) -> Result<()> {
440 while self.next_batch()?.is_some() {}
441 Ok(())
442 }
443}
444
445pub struct WriteRequest {
447 pub path: String,
448 pub format: String,
449 pub options: Map<String, Value>,
450 pub result_sets: Vec<ResultSetMeta>,
451 pub template: Option<Value>,
452}
453
454pub struct ResultSets<'a> {
456 metas: std::vec::IntoIter<ResultSetMeta>,
457 input: &'a mut Input,
458 warnings: Vec<String>,
459}
460
461impl ResultSets<'_> {
462 pub fn warn(&mut self, message: impl Into<String>) {
464 self.warnings.push(message.into());
465 }
466
467 pub fn next_set(&mut self) -> Result<Option<ResultSet<'_>>> {
470 let Some(meta) = self.metas.next() else {
471 return Ok(None);
472 };
473 let ipc = match self.input.read()? {
474 Frame::Arrow(ipc) => ipc,
475 Frame::Json(v) => {
476 return Err(format!(
477 "expected an Arrow frame starting result set `{}`, got {v}",
478 meta.name
479 )
480 .into());
481 }
482 };
483 let (schema, batches) = frame::decode_batches(&ipc)?;
484 Ok(Some(ResultSet {
485 meta,
486 schema,
487 first: Some(batches),
488 input: &mut *self.input,
489 done: false,
490 }))
491 }
492}
493
494pub trait Format {
495 fn options(&self) -> Vec<OptionField> {
498 Vec::new()
499 }
500 fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
503 Vec::new()
504 }
505 fn validate_template(
508 &self,
509 _options: &Map<String, Value>,
510 _template: &Value,
511 ) -> (Vec<String>, Vec<String>) {
512 (Vec::new(), Vec::new())
513 }
514 fn unknown_option_hint(&self, _key: &str) -> Option<String> {
517 None
518 }
519 fn write(&mut self, req: &WriteRequest, sets: &mut ResultSets<'_>) -> Result<Vec<String>>;
522}
523
524pub struct Delivery {
526 pub files: Vec<DeliveryFile>,
529 pub connection: Map<String, Value>,
530 pub options: Map<String, Value>,
532 pub message: Option<crate::msg::Message>,
534}
535
536pub struct DeliveryFile {
538 pub local: PathBuf,
539 pub remote: Option<String>,
540}
541
542pub trait Destination {
543 fn connection_fields(&self) -> Vec<ConnectionField> {
544 Vec::new()
545 }
546 fn options(&self) -> Vec<OptionField> {
550 Vec::new()
551 }
552 fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
555 Vec::new()
556 }
557 fn validate_connection(&self, _connection: &Map<String, Value>) -> Vec<String> {
561 Vec::new()
562 }
563 fn deliver(
566 &mut self,
567 _local: &Path,
568 _remote: Option<&str>,
569 _connection: &Map<String, Value>,
570 ) -> Result<String> {
571 Err("this destination doesn't implement `deliver`".into())
572 }
573 fn message_limit(&self) -> Option<u64> {
576 None
577 }
578 fn deliver_message(&mut self, _d: &Delivery, _m: &crate::msg::Message) -> Result<String> {
581 Err("this destination doesn't take messages".into())
582 }
583 fn deliver_files(&mut self, d: &Delivery) -> Result<String> {
586 match d.files.as_slice() {
587 [f] => self.deliver(&f.local, f.remote.as_deref(), &d.connection),
588 _ => Err("this destination takes one file per delivery".into()),
589 }
590 }
591}
592
593pub struct Input {
595 rx: Receiver<std::result::Result<Frame, FrameError>>,
596}
597
598impl Input {
599 fn next(&mut self) -> std::result::Result<Frame, FrameError> {
600 self.rx.recv().unwrap_or(Err(FrameError::Eof))
601 }
602
603 fn read(&mut self) -> Result<Frame> {
604 self.next().map_err(|e| -> Error {
605 match e {
606 FrameError::Eof => "core closed the connection".into(),
607 e => e.to_string().into(),
608 }
609 })
610 }
611}
612
613fn read_stdin(tx: Sender<std::result::Result<Frame, FrameError>>) {
616 let mut r = BufReader::new(std::io::stdin());
617 loop {
618 let f = frame::read_frame(&mut r);
619 match &f {
620 Ok(Frame::Json(v)) if v.get("type").and_then(Value::as_str) == Some("cancel") => {
621 cancel(v.get("id").and_then(Value::as_u64).unwrap_or(0));
622 continue;
623 }
624 Err(_) => cancel(shared().current.load(Ordering::SeqCst)),
625 Ok(_) => {}
626 }
627 let stop = f.is_err();
628 if tx.send(f).is_err() || stop {
629 break;
630 }
631 }
632}
633
634struct Output;
636
637impl Output {
638 fn send(&mut self, r: &Response) {
639 send(r)
640 }
641
642 fn batch(&mut self, b: &RecordBatch) -> Result<()> {
643 let ipc = frame::encode_batch(b)?;
644 let mut out = shared().out.lock().unwrap();
645 frame::write_arrow(&mut *out, &ipc).map_err(|_| -> Error { "core closed the connection".into() })
646 }
647}
648
649enum Handler<'a> {
650 Source(&'a mut dyn Source),
651 Format(&'a mut dyn Format),
652 Destination(&'a mut dyn Destination),
653}
654
655impl Handler<'_> {
656 fn kind(&self) -> Kind {
657 match self {
658 Handler::Source(_) => Kind::Source,
659 Handler::Format(_) => Kind::Format,
660 Handler::Destination(_) => Kind::Destination,
661 }
662 }
663
664 fn option_fields(&self) -> Vec<OptionField> {
665 match self {
666 Handler::Source(_) => Vec::new(),
667 Handler::Format(f) => f.options(),
668 Handler::Destination(d) => d.options(),
669 }
670 }
671
672 fn check(&self, name: &str, o: &Map<String, Value>) -> Vec<String> {
675 let mut errs = options::check(self.kind(), name, &self.option_fields(), o);
676 if let Handler::Format(f) = self {
677 for e in &mut errs {
678 let hint = o
679 .keys()
680 .find(|k| e.starts_with(&format!("unknown option `{k}`")))
681 .and_then(|k| f.unknown_option_hint(k));
682 if let Some(h) = hint {
683 e.push_str("; ");
684 e.push_str(&h);
685 }
686 }
687 }
688 if !errs.is_empty() {
689 return errs;
690 }
691 match self {
692 Handler::Source(_) => Vec::new(),
693 Handler::Format(f) => f.validate(o),
694 Handler::Destination(d) => d.validate(o),
695 }
696 }
697
698 fn checked(&self, name: &str, o: &Map<String, Value>) -> Result<()> {
700 let errs = self.check(name, o);
701 if errs.is_empty() {
702 Ok(())
703 } else {
704 Err(PluginError::new(ErrorKind::Config, errs.join("; ")).into())
706 }
707 }
708}
709
710pub fn serve_source(about: About, s: impl Source + 'static) -> ! {
711 serve_package(vec![Plugin::Source(about, Box::new(s))])
712}
713
714pub fn serve_format(about: About, f: impl Format + 'static) -> ! {
715 serve_package(vec![Plugin::Format(about, Box::new(f))])
716}
717
718pub fn serve_destination(about: About, d: impl Destination + 'static) -> ! {
719 serve_package(vec![Plugin::Destination(about, Box::new(d))])
720}
721
722pub enum Plugin {
724 Source(About, Box<dyn Source>),
725 Format(About, Box<dyn Format>),
726 Destination(About, Box<dyn Destination>),
727}
728
729impl Plugin {
730 fn id(&self) -> PluginId {
731 let (kind, about) = match self {
732 Plugin::Source(a, _) => (Kind::Source, a),
733 Plugin::Format(a, _) => (Kind::Format, a),
734 Plugin::Destination(a, _) => (Kind::Destination, a),
735 };
736 PluginId::new(kind, about.name)
737 }
738
739 fn handler(&mut self) -> (&About, Handler<'_>) {
740 match self {
741 Plugin::Source(a, s) => (a, Handler::Source(s.as_mut())),
742 Plugin::Format(a, f) => (a, Handler::Format(f.as_mut())),
743 Plugin::Destination(a, d) => (a, Handler::Destination(d.as_mut())),
744 }
745 }
746}
747
748pub fn serve_package(mut plugins: Vec<Plugin>) -> ! {
755 assert!(!plugins.is_empty(), "a package serves at least one plugin");
756 shared();
758 let (tx, rx) = std::sync::mpsc::channel();
759 std::thread::spawn(move || read_stdin(tx));
760 let mut input = Input { rx };
761 let mut out = Output;
762 if ::log::set_logger(&FrameLogger).is_ok() {
763 ::log::set_max_level(::log::LevelFilter::Debug);
764 }
765 let provides: Vec<PluginId> = plugins.iter().map(Plugin::id).collect();
766 let chosen = loop {
767 let Some(req) = next_request(&mut input, &mut out) else {
768 continue;
769 };
770 let Request::Hello {
771 min_version,
772 max_version,
773 plugin,
774 ..
775 } = req
776 else {
777 out.send(&Response::error("the first request must be `hello`"));
778 continue;
779 };
780 let Some(chosen) = negotiate((min_version, max_version), (MIN_VERSION, MAX_VERSION)) else {
781 out.send(&Response::VersionMismatch {
782 min_version: MIN_VERSION,
783 max_version: MAX_VERSION,
784 });
785 std::process::exit(1);
786 };
787 let i = match &plugin {
788 None => 0,
789 Some(want) => match provides.iter().position(|p| p == want) {
790 Some(i) => i,
791 None => {
792 let list: Vec<String> = provides.iter().map(|p| p.to_string()).collect();
793 out.send(&Response::error(format!(
794 "this executable provides {}, not {want}",
795 list.join(", ")
796 )));
797 std::process::exit(1);
798 }
799 },
800 };
801 let (about, h) = plugins[i].handler();
803 let _ = shared().name.set(about.name.to_string());
804 out.send(&Response::Hello {
805 protocol_version: chosen,
806 kind: h.kind(),
807 name: about.name.to_string(),
808 version: about.version.to_string(),
809 capabilities: about
810 .capabilities
811 .iter()
812 .chain(std::iter::once(&CAP_VALIDATE))
813 .map(|c| c.to_string())
814 .collect(),
815 provides: if provides.len() > 1 { provides } else { Vec::new() },
816 });
817 shared().version.store(chosen, Ordering::SeqCst);
818 break i;
819 };
820 let (about, h) = plugins[chosen].handler();
821 serve(about.name, h, input, out)
822}
823
824fn next_request(input: &mut Input, out: &mut Output) -> Option<Request> {
826 let frame = match input.next() {
827 Ok(f) => f,
828 Err(FrameError::Eof) => std::process::exit(0),
829 Err(e) => {
830 eprintln!("{e}");
831 out.send(&Response::error(e.to_string()));
832 std::process::exit(2);
833 }
834 };
835 match frame {
836 Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
837 Ok(r) => Some(r),
838 Err(_) => {
839 let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
840 out.send(&Response::error(format!("unsupported request `{t}`")));
841 None
842 }
843 },
844 Frame::Arrow(_) => {
845 out.send(&Response::error("unexpected Arrow frame"));
846 None
847 }
848 }
849}
850
851fn serve(name: &str, mut h: Handler<'_>, mut input: Input, mut out: Output) -> ! {
853 loop {
854 let frame = match input.next() {
855 Ok(f) => f,
856 Err(FrameError::Eof) => {
857 if let Handler::Source(src) = &mut h {
858 src.close();
859 }
860 std::process::exit(0)
861 }
862 Err(e) => {
863 eprintln!("{e}");
864 out.send(&Response::error(e.to_string()));
865 std::process::exit(2);
866 }
867 };
868 let s = shared();
869 let id = match &frame {
870 Frame::Json(v) => v.get("id").and_then(Value::as_u64).unwrap_or(0),
871 Frame::Arrow(_) => 0,
872 };
873 s.current.store(id, Ordering::SeqCst);
874 let req = match frame {
875 Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
876 Ok(r) => r,
877 Err(_) => {
878 let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
879 out.send(&Response::error(format!("unsupported request `{t}`")));
880 s.current.store(0, Ordering::SeqCst);
881 continue;
882 }
883 },
884 Frame::Arrow(_) => {
885 out.send(&Response::error("unexpected Arrow frame"));
886 continue;
887 }
888 };
889 if let Request::Hello { .. } = req {
890 out.send(&Response::error("`hello` was already answered"));
891 s.current.store(0, Ordering::SeqCst);
892 continue;
893 }
894 if let Request::Close {} = req {
895 if let Handler::Source(src) = &mut h {
896 src.close();
897 }
898 out.send(&Response::Ok {});
899 std::process::exit(0);
900 }
901 let result = if cancelled() {
903 Ok(Err("cancelled before it started".into()))
904 } else {
905 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
908 handle(&mut h, name, req, &mut input, &mut out)
909 }))
910 };
911 match result {
912 Ok(Ok(())) => {}
913 Ok(Err(e)) => out.send(&error_reply(&e)),
914 Err(p) => {
915 let msg = p
916 .downcast_ref::<String>()
917 .cloned()
918 .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()))
919 .unwrap_or_else(|| "unknown panic".into());
920 out.send(&Response::Error {
921 message: format!("plugin panicked: {msg}"),
922 kind: Some("internal".into()),
923 code: None,
924 });
925 std::process::exit(101);
926 }
927 }
928 s.current.store(0, Ordering::SeqCst);
929 }
930}
931
932fn check_connection(
935 h: &Handler<'_>,
936 name: &str,
937 connection: &Map<String, Value>,
938 unresolved: &[String],
939) -> Result<crate::connection::Checked> {
940 let (kind, fields) = match h {
941 Handler::Source(s) => (Kind::Source, s.connection_fields()),
942 Handler::Destination(d) => (Kind::Destination, d.connection_fields()),
943 Handler::Format(_) => return Err("a format has no connection settings".into()),
944 };
945 let mut c = crate::connection::check(kind, name, &fields, connection, unresolved);
946 if c.errors.is_empty() {
947 c.errors = match h {
950 Handler::Source(s) => s.validate_connection(connection),
951 Handler::Destination(d) => d.validate_connection(connection),
952 Handler::Format(_) => Vec::new(),
953 };
954 }
955 c.errors = crate::connection::redact(c.errors, &fields, connection);
956 c.warnings = crate::connection::redact(c.warnings, &fields, connection);
957 Ok(c)
958}
959
960fn handle(h: &mut Handler<'_>, name: &str, req: Request, input: &mut Input, out: &mut Output) -> Result<()> {
961 if let Request::Describe {} = req {
962 let connection_fields = match h {
963 Handler::Source(s) => s.connection_fields(),
964 Handler::Destination(d) => d.connection_fields(),
965 Handler::Format(_) => Vec::new(),
966 };
967 let (identifier_quote, string_literal) = match h {
968 Handler::Source(s) => (
969 s.identifier_quote().map(str::to_string),
970 s.string_literal().map(str::to_string),
971 ),
972 _ => (None, None),
973 };
974 let message_limit = match h {
975 Handler::Destination(d) => d.message_limit(),
976 _ => None,
977 };
978 out.send(&Response::Describe {
979 connection_fields,
980 option_fields: h.option_fields(),
981 identifier_quote,
982 string_literal,
983 message_limit,
984 });
985 return Ok(());
986 }
987 if let Request::Validate { options, template } = req {
988 let mut errors = h.check(name, &options);
989 let mut warnings = Vec::new();
990 if let (Handler::Format(f), Some(t), true) = (&*h, &template, errors.is_empty()) {
991 let (e, w) = f.validate_template(&options, t);
992 errors.extend(e);
993 warnings.extend(w);
994 }
995 out.send(&Response::Validated { errors, warnings });
996 return Ok(());
997 }
998 if let Request::ValidateConnection {
999 connection,
1000 unresolved,
1001 } = req
1002 {
1003 let c = check_connection(h, name, &connection, &unresolved)?;
1004 out.send(&Response::Validated {
1005 errors: c.errors,
1006 warnings: c.warnings,
1007 });
1008 return Ok(());
1009 }
1010 if let Request::Open { connection, .. } | Request::Deliver { connection, .. } = &req {
1012 let c = check_connection(h, name, connection, &[])?;
1013 for w in &c.warnings {
1014 crate::log::warn!("{w}");
1015 }
1016 if !c.errors.is_empty() {
1017 return Err(PluginError::new(ErrorKind::Config, c.errors.join("; ")).into());
1018 }
1019 }
1020 let checked = match &req {
1021 Request::Write { options, .. } | Request::Deliver { options, .. } => h.checked(name, options),
1022 _ => Ok(()),
1023 };
1024 match (h, req) {
1025 (
1026 Handler::Source(s),
1027 Request::Open {
1028 connection,
1029 read_only,
1030 sandbox,
1031 },
1032 ) => {
1033 if sandbox {
1034 s.open_sandboxed(&connection)?;
1035 } else {
1036 s.open(&connection, read_only)?;
1037 }
1038 out.send(&Response::Ok {});
1039 }
1040 (Handler::Source(s), Request::Load { name }) => {
1041 let ipc = match input.read()? {
1042 Frame::Arrow(ipc) => ipc,
1043 Frame::Json(v) => return Err(format!("expected the rows to load, got {v}").into()),
1044 };
1045 let (schema, batches) = frame::decode_batches(&ipc)?;
1046 let mut data = ResultSet {
1047 meta: ResultSetMeta {
1048 name: name.clone(),
1049 query: name.clone(),
1050 result_index: 1,
1051 anchor: None,
1052 header: None,
1053 columns: Default::default(),
1054 autofit: None,
1055 style: None,
1056 },
1057 schema,
1058 first: Some(batches),
1059 input,
1060 done: false,
1061 };
1062 let r = s.load(&name, &mut data);
1063 data.drain()?;
1064 let loaded = r?;
1065 out.send(&Response::Loaded {
1066 relation: loaded.relation,
1067 rows: loaded.rows,
1068 warning: loaded.warning,
1069 });
1070 }
1071 (Handler::Source(s), Request::Check { sql }) => {
1072 s.check(&sql)?;
1073 out.send(&Response::Ok {});
1074 }
1075 (Handler::Source(s), Request::Execute { sql, row_limit }) => {
1076 let mut sink = FrameSink {
1077 out,
1078 row_limit,
1079 rows: 0,
1080 state: SinkState::Idle,
1081 schema: None,
1082 sent_any: false,
1083 };
1084 let r = s.execute(&sql, row_limit, &mut sink);
1085 match (r, sink.state) {
1086 (Err(e), _) => return Err(e),
1087 (Ok(()), SinkState::Idle) => {
1088 return Err("the source produced neither a result nor `no_result`".into());
1089 }
1090 (Ok(()), SinkState::NoResult) => {}
1091 (Ok(()), SinkState::Result) => {
1092 if !sink.sent_any {
1093 let schema = sink.schema.clone().expect("a result sets the schema");
1094 sink.out.batch(&RecordBatch::new_empty(schema))?;
1095 }
1096 let rows = sink.rows;
1097 sink.out.send(&Response::ResultEnd { rows });
1098 }
1099 }
1100 }
1101 (
1102 Handler::Format(f),
1103 Request::Write {
1104 path,
1105 format,
1106 options,
1107 result_sets,
1108 template,
1109 },
1110 ) => {
1111 let req = WriteRequest {
1112 path,
1113 format,
1114 options,
1115 result_sets: result_sets.clone(),
1116 template,
1117 };
1118 let mut sets = ResultSets {
1119 metas: result_sets.into_iter(),
1120 input,
1121 warnings: Vec::new(),
1122 };
1123 let written = checked.and_then(|()| f.write(&req, &mut sets));
1124 if let Err(e) = &written {
1127 out.send(&error_reply(e));
1128 }
1129 loop {
1132 match sets.input.read()? {
1133 Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("finish") => break,
1134 Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("result_set_end") => {}
1135 Frame::Arrow(_) => {}
1136 other => return Err(format!("expected `finish`, got {other:?}").into()),
1137 }
1138 }
1139 if let Ok(files) = written {
1140 out.send(&Response::Written {
1141 files,
1142 warnings: sets.warnings,
1143 });
1144 }
1145 }
1146 (
1147 Handler::Destination(d),
1148 Request::Deliver {
1149 local_path,
1150 remote_path,
1151 files,
1152 connection,
1153 options,
1154 message,
1155 },
1156 ) => {
1157 checked?;
1158 let files = match (local_path, files.is_empty(), &message) {
1159 (Some(local), true, None) => vec![DeliveryFile {
1160 local: PathBuf::from(local),
1161 remote: remote_path,
1162 }],
1163 (None, false, _) | (None, true, Some(_)) => files
1164 .into_iter()
1165 .map(|f| DeliveryFile {
1166 local: PathBuf::from(f.local_path),
1167 remote: f.remote_path,
1168 })
1169 .collect(),
1170 _ => return Err("`deliver` needs exactly one of `local_path` or `files`".into()),
1171 };
1172 let delivery = Delivery {
1173 files,
1174 connection,
1175 options,
1176 message,
1177 };
1178 crate::delivery::take_attempts();
1179 let location = match &delivery.message {
1180 Some(m) => d.deliver_message(&delivery, m)?,
1181 None => d.deliver_files(&delivery)?,
1182 };
1183 let attempts = Some(crate::delivery::take_attempts()).filter(|n| *n > 1);
1184 out.send(&Response::Delivered { location, attempts });
1185 }
1186 (h, req) => {
1187 let t = serde_json::to_value(&req)
1188 .ok()
1189 .and_then(|v| v.get("type").cloned())
1190 .unwrap_or_default();
1191 return Err(format!("a {} plugin doesn't handle {t} requests", h.kind()).into());
1192 }
1193 }
1194 Ok(())
1195}
1196
1197#[derive(PartialEq)]
1198enum SinkState {
1199 Idle,
1200 NoResult,
1201 Result,
1202}
1203
1204struct FrameSink<'a> {
1205 out: &'a mut Output,
1206 row_limit: Option<u64>,
1207 rows: u64,
1208 state: SinkState,
1209 schema: Option<SchemaRef>,
1210 sent_any: bool,
1211}
1212
1213impl ResultSink for FrameSink<'_> {
1214 fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()> {
1215 if self.state != SinkState::Idle {
1216 return Err("`no_result` after the result started".into());
1217 }
1218 self.state = SinkState::NoResult;
1219 self.out.send(&Response::NoResult { rows_affected });
1220 Ok(())
1221 }
1222
1223 fn begin(&mut self, schema: SchemaRef) -> Result<()> {
1224 if self.state != SinkState::Idle {
1225 return Err("a statement can only produce one result".into());
1226 }
1227 self.state = SinkState::Result;
1228 self.out.send(&Response::Result {
1229 columns: schema.fields().iter().map(|f| f.name().clone()).collect(),
1230 });
1231 self.schema = Some(schema);
1232 Ok(())
1233 }
1234
1235 fn batch(&mut self, b: RecordBatch) -> Result<bool> {
1236 if self.state != SinkState::Result {
1237 return Err("`batch` before `begin`".into());
1238 }
1239 if self.row_limit.is_some_and(|l| self.rows >= l) {
1240 return Ok(false);
1241 }
1242 let b = match self.row_limit {
1243 Some(l) if self.rows + b.num_rows() as u64 > l => b.slice(0, (l - self.rows) as usize),
1244 _ => b,
1245 };
1246 if b.num_rows() > 0 || !self.sent_any {
1247 self.rows += b.num_rows() as u64;
1248 self.out.batch(&b)?;
1249 self.sent_any = true;
1250 }
1251 Ok(self.row_limit.is_none_or(|l| self.rows < l))
1252 }
1253}
1254
1255fn negotiate(core: (u32, u32), ours: (u32, u32)) -> Option<u32> {
1257 let (lo, hi) = (core.0.max(ours.0), core.1.min(ours.1));
1258 (lo <= hi).then_some(hi)
1259}
1260
1261pub fn conn_str<'a>(c: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
1263 c.get(key).and_then(Value::as_str).filter(|s| !s.is_empty())
1264}
1265
1266pub fn conn_required<'a>(c: &'a Map<String, Value>, key: &str) -> Result<&'a str> {
1268 conn_str(c, key).ok_or_else(|| format!("the profile output needs a `{key}` field").into())
1269}
1270
1271pub fn conn_bool(c: &Map<String, Value>, key: &str) -> Option<bool> {
1273 match c.get(key)? {
1274 Value::Bool(b) => Some(*b),
1275 Value::String(s) => Some(matches!(s.as_str(), "true" | "yes" | "1")),
1276 _ => None,
1277 }
1278}