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