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 write(&mut self, req: &WriteRequest, sets: &mut ResultSets<'_>) -> Result<Vec<String>>;
511}
512
513pub struct Delivery {
515 pub files: Vec<DeliveryFile>,
518 pub connection: Map<String, Value>,
519 pub options: Map<String, Value>,
521 pub message: Option<crate::msg::Message>,
523}
524
525pub struct DeliveryFile {
527 pub local: PathBuf,
528 pub remote: Option<String>,
529}
530
531pub trait Destination {
532 fn connection_fields(&self) -> Vec<ConnectionField> {
533 Vec::new()
534 }
535 fn options(&self) -> Vec<OptionField> {
539 Vec::new()
540 }
541 fn validate(&self, _options: &Map<String, Value>) -> Vec<String> {
544 Vec::new()
545 }
546 fn validate_connection(&self, _connection: &Map<String, Value>) -> Vec<String> {
550 Vec::new()
551 }
552 fn deliver(
555 &mut self,
556 _local: &Path,
557 _remote: Option<&str>,
558 _connection: &Map<String, Value>,
559 ) -> Result<String> {
560 Err("this destination doesn't implement `deliver`".into())
561 }
562 fn message_limit(&self) -> Option<u64> {
565 None
566 }
567 fn deliver_message(&mut self, _d: &Delivery, _m: &crate::msg::Message) -> Result<String> {
570 Err("this destination doesn't take messages".into())
571 }
572 fn deliver_files(&mut self, d: &Delivery) -> Result<String> {
575 match d.files.as_slice() {
576 [f] => self.deliver(&f.local, f.remote.as_deref(), &d.connection),
577 _ => Err("this destination takes one file per delivery".into()),
578 }
579 }
580}
581
582pub struct Input {
584 rx: Receiver<std::result::Result<Frame, FrameError>>,
585}
586
587impl Input {
588 fn next(&mut self) -> std::result::Result<Frame, FrameError> {
589 self.rx.recv().unwrap_or(Err(FrameError::Eof))
590 }
591
592 fn read(&mut self) -> Result<Frame> {
593 self.next().map_err(|e| -> Error {
594 match e {
595 FrameError::Eof => "core closed the connection".into(),
596 e => e.to_string().into(),
597 }
598 })
599 }
600}
601
602fn read_stdin(tx: Sender<std::result::Result<Frame, FrameError>>) {
605 let mut r = BufReader::new(std::io::stdin());
606 loop {
607 let f = frame::read_frame(&mut r);
608 match &f {
609 Ok(Frame::Json(v)) if v.get("type").and_then(Value::as_str) == Some("cancel") => {
610 cancel(v.get("id").and_then(Value::as_u64).unwrap_or(0));
611 continue;
612 }
613 Err(_) => cancel(shared().current.load(Ordering::SeqCst)),
614 Ok(_) => {}
615 }
616 let stop = f.is_err();
617 if tx.send(f).is_err() || stop {
618 break;
619 }
620 }
621}
622
623struct Output;
625
626impl Output {
627 fn send(&mut self, r: &Response) {
628 send(r)
629 }
630
631 fn batch(&mut self, b: &RecordBatch) -> Result<()> {
632 let ipc = frame::encode_batch(b)?;
633 let mut out = shared().out.lock().unwrap();
634 frame::write_arrow(&mut *out, &ipc).map_err(|_| -> Error { "core closed the connection".into() })
635 }
636}
637
638enum Handler<'a> {
639 Source(&'a mut dyn Source),
640 Format(&'a mut dyn Format),
641 Destination(&'a mut dyn Destination),
642}
643
644impl Handler<'_> {
645 fn kind(&self) -> Kind {
646 match self {
647 Handler::Source(_) => Kind::Source,
648 Handler::Format(_) => Kind::Format,
649 Handler::Destination(_) => Kind::Destination,
650 }
651 }
652
653 fn option_fields(&self) -> Vec<OptionField> {
654 match self {
655 Handler::Source(_) => Vec::new(),
656 Handler::Format(f) => f.options(),
657 Handler::Destination(d) => d.options(),
658 }
659 }
660
661 fn check(&self, name: &str, o: &Map<String, Value>) -> Vec<String> {
664 let errs = options::check(self.kind(), name, &self.option_fields(), o);
665 if !errs.is_empty() {
666 return errs;
667 }
668 match self {
669 Handler::Source(_) => Vec::new(),
670 Handler::Format(f) => f.validate(o),
671 Handler::Destination(d) => d.validate(o),
672 }
673 }
674
675 fn checked(&self, name: &str, o: &Map<String, Value>) -> Result<()> {
677 let errs = self.check(name, o);
678 if errs.is_empty() {
679 Ok(())
680 } else {
681 Err(PluginError::new(ErrorKind::Config, errs.join("; ")).into())
683 }
684 }
685}
686
687pub fn serve_source(about: About, s: impl Source + 'static) -> ! {
688 serve_package(vec![Plugin::Source(about, Box::new(s))])
689}
690
691pub fn serve_format(about: About, f: impl Format + 'static) -> ! {
692 serve_package(vec![Plugin::Format(about, Box::new(f))])
693}
694
695pub fn serve_destination(about: About, d: impl Destination + 'static) -> ! {
696 serve_package(vec![Plugin::Destination(about, Box::new(d))])
697}
698
699pub enum Plugin {
701 Source(About, Box<dyn Source>),
702 Format(About, Box<dyn Format>),
703 Destination(About, Box<dyn Destination>),
704}
705
706impl Plugin {
707 fn id(&self) -> PluginId {
708 let (kind, about) = match self {
709 Plugin::Source(a, _) => (Kind::Source, a),
710 Plugin::Format(a, _) => (Kind::Format, a),
711 Plugin::Destination(a, _) => (Kind::Destination, a),
712 };
713 PluginId::new(kind, about.name)
714 }
715
716 fn handler(&mut self) -> (&About, Handler<'_>) {
717 match self {
718 Plugin::Source(a, s) => (a, Handler::Source(s.as_mut())),
719 Plugin::Format(a, f) => (a, Handler::Format(f.as_mut())),
720 Plugin::Destination(a, d) => (a, Handler::Destination(d.as_mut())),
721 }
722 }
723}
724
725pub fn serve_package(mut plugins: Vec<Plugin>) -> ! {
728 assert!(!plugins.is_empty(), "a package serves at least one plugin");
729 shared();
731 let (tx, rx) = std::sync::mpsc::channel();
732 std::thread::spawn(move || read_stdin(tx));
733 let mut input = Input { rx };
734 let mut out = Output;
735 if ::log::set_logger(&FrameLogger).is_ok() {
736 ::log::set_max_level(::log::LevelFilter::Debug);
737 }
738 let provides: Vec<PluginId> = plugins.iter().map(Plugin::id).collect();
739 let chosen = loop {
740 let req = match next_request(&mut input, &mut out) {
741 Some(r) => r,
742 None => continue,
743 };
744 let Request::Hello {
745 min_version,
746 max_version,
747 plugin,
748 ..
749 } = req
750 else {
751 out.send(&Response::error("the first request must be `hello`"));
752 continue;
753 };
754 let Some(chosen) = negotiate((min_version, max_version), (MIN_VERSION, MAX_VERSION)) else {
755 out.send(&Response::VersionMismatch {
756 min_version: MIN_VERSION,
757 max_version: MAX_VERSION,
758 });
759 std::process::exit(1);
760 };
761 let i = match &plugin {
762 None => 0,
763 Some(want) => match provides.iter().position(|p| p == want) {
764 Some(i) => i,
765 None => {
766 let list: Vec<String> = provides.iter().map(|p| p.to_string()).collect();
767 out.send(&Response::error(format!(
768 "this executable provides {}, not {want}",
769 list.join(", ")
770 )));
771 std::process::exit(1);
772 }
773 },
774 };
775 let (about, h) = plugins[i].handler();
776 let _ = shared().name.set(about.name.to_string());
777 out.send(&Response::Hello {
778 protocol_version: chosen,
779 kind: h.kind(),
780 name: about.name.to_string(),
781 version: about.version.to_string(),
782 capabilities: about
783 .capabilities
784 .iter()
785 .chain(std::iter::once(&CAP_VALIDATE))
786 .map(|c| c.to_string())
787 .collect(),
788 provides: if provides.len() > 1 {
789 provides.clone()
790 } else {
791 Vec::new()
792 },
793 });
794 shared().version.store(chosen, Ordering::SeqCst);
795 break i;
796 };
797 let (about, h) = plugins[chosen].handler();
798 serve(about.name, h, input, out)
799}
800
801fn next_request(input: &mut Input, out: &mut Output) -> Option<Request> {
803 let frame = match input.next() {
804 Ok(f) => f,
805 Err(FrameError::Eof) => std::process::exit(0),
806 Err(e) => {
807 eprintln!("{e}");
808 out.send(&Response::error(e.to_string()));
809 std::process::exit(2);
810 }
811 };
812 match frame {
813 Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
814 Ok(r) => Some(r),
815 Err(_) => {
816 let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
817 out.send(&Response::error(format!("unsupported request `{t}`")));
818 None
819 }
820 },
821 Frame::Arrow(_) => {
822 out.send(&Response::error("unexpected Arrow frame"));
823 None
824 }
825 }
826}
827
828fn serve(name: &str, mut h: Handler<'_>, mut input: Input, mut out: Output) -> ! {
830 loop {
831 let frame = match input.next() {
832 Ok(f) => f,
833 Err(FrameError::Eof) => {
834 if let Handler::Source(src) = &mut h {
835 src.close();
836 }
837 std::process::exit(0)
838 }
839 Err(e) => {
840 eprintln!("{e}");
841 out.send(&Response::error(e.to_string()));
842 std::process::exit(2);
843 }
844 };
845 let s = shared();
846 let id = match &frame {
847 Frame::Json(v) => v.get("id").and_then(Value::as_u64).unwrap_or(0),
848 Frame::Arrow(_) => 0,
849 };
850 s.current.store(id, Ordering::SeqCst);
851 let req = match frame {
852 Frame::Json(v) => match serde_json::from_value::<Request>(v.clone()) {
853 Ok(r) => r,
854 Err(_) => {
855 let t = v.get("type").and_then(Value::as_str).unwrap_or("?").to_string();
856 out.send(&Response::error(format!("unsupported request `{t}`")));
857 s.current.store(0, Ordering::SeqCst);
858 continue;
859 }
860 },
861 Frame::Arrow(_) => {
862 out.send(&Response::error("unexpected Arrow frame"));
863 continue;
864 }
865 };
866 if let Request::Hello { .. } = req {
867 out.send(&Response::error("`hello` was already answered"));
868 s.current.store(0, Ordering::SeqCst);
869 continue;
870 }
871 if let Request::Close {} = req {
872 if let Handler::Source(src) = &mut h {
873 src.close();
874 }
875 out.send(&Response::Ok {});
876 std::process::exit(0);
877 }
878 let result = if cancelled() {
880 Ok(Err("cancelled before it started".into()))
881 } else {
882 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
883 handle(&mut h, name, req, &mut input, &mut out)
884 }))
885 };
886 match result {
887 Ok(Ok(())) => {}
888 Ok(Err(e)) => out.send(&error_reply(&e)),
889 Err(p) => {
890 let msg = p
891 .downcast_ref::<String>()
892 .cloned()
893 .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()))
894 .unwrap_or_else(|| "unknown panic".into());
895 out.send(&Response::Error {
896 message: format!("plugin panicked: {msg}"),
897 kind: Some("internal".into()),
898 code: None,
899 });
900 std::process::exit(101);
901 }
902 }
903 s.current.store(0, Ordering::SeqCst);
904 }
905}
906
907fn check_connection(
910 h: &Handler<'_>,
911 name: &str,
912 connection: &Map<String, Value>,
913 unresolved: &[String],
914) -> Result<crate::connection::Checked> {
915 let (kind, fields) = match h {
916 Handler::Source(s) => (Kind::Source, s.connection_fields()),
917 Handler::Destination(d) => (Kind::Destination, d.connection_fields()),
918 Handler::Format(_) => return Err("a format has no connection settings".into()),
919 };
920 let mut c = crate::connection::check(kind, name, &fields, connection, unresolved);
921 if c.errors.is_empty() {
922 c.errors = match h {
925 Handler::Source(s) => s.validate_connection(connection),
926 Handler::Destination(d) => d.validate_connection(connection),
927 Handler::Format(_) => Vec::new(),
928 };
929 }
930 c.errors = crate::connection::redact(c.errors, &fields, connection);
931 c.warnings = crate::connection::redact(c.warnings, &fields, connection);
932 Ok(c)
933}
934
935fn handle(h: &mut Handler<'_>, name: &str, req: Request, input: &mut Input, out: &mut Output) -> Result<()> {
936 if let Request::Describe {} = req {
937 let connection_fields = match h {
938 Handler::Source(s) => s.connection_fields(),
939 Handler::Destination(d) => d.connection_fields(),
940 Handler::Format(_) => Vec::new(),
941 };
942 let (identifier_quote, string_literal) = match h {
943 Handler::Source(s) => (
944 s.identifier_quote().map(str::to_string),
945 s.string_literal().map(str::to_string),
946 ),
947 _ => (None, None),
948 };
949 let message_limit = match h {
950 Handler::Destination(d) => d.message_limit(),
951 _ => None,
952 };
953 out.send(&Response::Describe {
954 connection_fields,
955 option_fields: h.option_fields(),
956 identifier_quote,
957 string_literal,
958 message_limit,
959 });
960 return Ok(());
961 }
962 if let Request::Validate { options, template } = req {
963 let mut errors = h.check(name, &options);
964 let mut warnings = Vec::new();
965 if let (Handler::Format(f), Some(t), true) = (&*h, &template, errors.is_empty()) {
966 let (e, w) = f.validate_template(&options, t);
967 errors.extend(e);
968 warnings.extend(w);
969 }
970 out.send(&Response::Validated { errors, warnings });
971 return Ok(());
972 }
973 if let Request::ValidateConnection {
974 connection,
975 unresolved,
976 } = req
977 {
978 let c = check_connection(h, name, &connection, &unresolved)?;
979 out.send(&Response::Validated {
980 errors: c.errors,
981 warnings: c.warnings,
982 });
983 return Ok(());
984 }
985 if let Request::Open { connection, .. } | Request::Deliver { connection, .. } = &req {
987 let c = check_connection(h, name, connection, &[])?;
988 for w in &c.warnings {
989 crate::log::warn!("{w}");
990 }
991 if !c.errors.is_empty() {
992 return Err(PluginError::new(ErrorKind::Config, c.errors.join("; ")).into());
993 }
994 }
995 let checked = match &req {
996 Request::Write { options, .. } | Request::Deliver { options, .. } => h.checked(name, options),
997 _ => Ok(()),
998 };
999 match (h, req) {
1000 (
1001 Handler::Source(s),
1002 Request::Open {
1003 connection,
1004 read_only,
1005 sandbox,
1006 },
1007 ) => {
1008 if sandbox {
1009 s.open_sandboxed(&connection)?;
1010 } else {
1011 s.open(&connection, read_only)?;
1012 }
1013 out.send(&Response::Ok {});
1014 }
1015 (Handler::Source(s), Request::Load { name }) => {
1016 let ipc = match input.read()? {
1017 Frame::Arrow(ipc) => ipc,
1018 Frame::Json(v) => return Err(format!("expected the rows to load, got {v}").into()),
1019 };
1020 let (schema, batches) = frame::decode_batches(&ipc)?;
1021 let mut data = ResultSet {
1022 meta: ResultSetMeta {
1023 name: name.clone(),
1024 query: name.clone(),
1025 result_index: 1,
1026 anchor: None,
1027 header: None,
1028 columns: Default::default(),
1029 autofit: None,
1030 style: None,
1031 },
1032 schema,
1033 first: Some(batches),
1034 input,
1035 done: false,
1036 };
1037 let r = s.load(&name, &mut data);
1038 data.drain()?;
1039 let loaded = r?;
1040 out.send(&Response::Loaded {
1041 relation: loaded.relation,
1042 rows: loaded.rows,
1043 warning: loaded.warning,
1044 });
1045 }
1046 (Handler::Source(s), Request::Check { sql }) => {
1047 s.check(&sql)?;
1048 out.send(&Response::Ok {});
1049 }
1050 (Handler::Source(s), Request::Execute { sql, row_limit }) => {
1051 let mut sink = FrameSink {
1052 out,
1053 row_limit,
1054 rows: 0,
1055 state: SinkState::Idle,
1056 schema: None,
1057 sent_any: false,
1058 };
1059 let r = s.execute(&sql, row_limit, &mut sink);
1060 match (r, sink.state) {
1061 (Err(e), _) => return Err(e),
1062 (Ok(()), SinkState::Idle) => {
1063 return Err("the source produced neither a result nor `no_result`".into());
1064 }
1065 (Ok(()), SinkState::NoResult) => {}
1066 (Ok(()), SinkState::Result) => {
1067 if !sink.sent_any {
1068 let schema = sink.schema.clone().unwrap();
1069 sink.out.batch(&RecordBatch::new_empty(schema))?;
1070 }
1071 let rows = sink.rows;
1072 sink.out.send(&Response::ResultEnd { rows });
1073 }
1074 }
1075 }
1076 (
1077 Handler::Format(f),
1078 Request::Write {
1079 path,
1080 format,
1081 options,
1082 result_sets,
1083 template,
1084 },
1085 ) => {
1086 let req = WriteRequest {
1087 path,
1088 format,
1089 options,
1090 result_sets: result_sets.clone(),
1091 template,
1092 };
1093 let mut sets = ResultSets {
1094 metas: result_sets.into_iter(),
1095 input,
1096 warnings: Vec::new(),
1097 };
1098 let written = checked.and_then(|()| f.write(&req, &mut sets));
1099 if let Err(e) = &written {
1102 out.send(&error_reply(e));
1103 }
1104 loop {
1107 match sets.input.read()? {
1108 Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("finish") => break,
1109 Frame::Json(v) if v.get("type").and_then(Value::as_str) == Some("result_set_end") => {}
1110 Frame::Arrow(_) => {}
1111 other => return Err(format!("expected `finish`, got {other:?}").into()),
1112 }
1113 }
1114 if let Ok(files) = written {
1115 out.send(&Response::Written {
1116 files,
1117 warnings: sets.warnings,
1118 });
1119 }
1120 }
1121 (
1122 Handler::Destination(d),
1123 Request::Deliver {
1124 local_path,
1125 remote_path,
1126 files,
1127 connection,
1128 options,
1129 message,
1130 },
1131 ) => {
1132 checked?;
1133 let files = match (local_path, files.is_empty(), &message) {
1134 (Some(local), true, None) => vec![DeliveryFile {
1135 local: PathBuf::from(local),
1136 remote: remote_path,
1137 }],
1138 (None, false, _) | (None, true, Some(_)) => files
1139 .into_iter()
1140 .map(|f| DeliveryFile {
1141 local: PathBuf::from(f.local_path),
1142 remote: f.remote_path,
1143 })
1144 .collect(),
1145 _ => return Err("`deliver` needs exactly one of `local_path` or `files`".into()),
1146 };
1147 let delivery = Delivery {
1148 files,
1149 connection,
1150 options,
1151 message,
1152 };
1153 crate::delivery::take_attempts();
1154 let location = match &delivery.message {
1155 Some(m) => d.deliver_message(&delivery, m)?,
1156 None => d.deliver_files(&delivery)?,
1157 };
1158 let attempts = Some(crate::delivery::take_attempts()).filter(|n| *n > 1);
1159 out.send(&Response::Delivered { location, attempts });
1160 }
1161 (h, req) => {
1162 let t = serde_json::to_value(&req)
1163 .ok()
1164 .and_then(|v| v.get("type").cloned())
1165 .unwrap_or_default();
1166 return Err(format!("a {} plugin doesn't handle {t} requests", h.kind()).into());
1167 }
1168 }
1169 Ok(())
1170}
1171
1172#[derive(PartialEq)]
1173enum SinkState {
1174 Idle,
1175 NoResult,
1176 Result,
1177}
1178
1179struct FrameSink<'a> {
1180 out: &'a mut Output,
1181 row_limit: Option<u64>,
1182 rows: u64,
1183 state: SinkState,
1184 schema: Option<SchemaRef>,
1185 sent_any: bool,
1186}
1187
1188impl ResultSink for FrameSink<'_> {
1189 fn no_result(&mut self, rows_affected: Option<u64>) -> Result<()> {
1190 if self.state != SinkState::Idle {
1191 return Err("`no_result` after the result started".into());
1192 }
1193 self.state = SinkState::NoResult;
1194 self.out.send(&Response::NoResult { rows_affected });
1195 Ok(())
1196 }
1197
1198 fn begin(&mut self, schema: SchemaRef) -> Result<()> {
1199 if self.state != SinkState::Idle {
1200 return Err("a statement can only produce one result".into());
1201 }
1202 self.state = SinkState::Result;
1203 self.out.send(&Response::Result {
1204 columns: schema.fields().iter().map(|f| f.name().clone()).collect(),
1205 });
1206 self.schema = Some(schema);
1207 Ok(())
1208 }
1209
1210 fn batch(&mut self, b: RecordBatch) -> Result<bool> {
1211 if self.state != SinkState::Result {
1212 return Err("`batch` before `begin`".into());
1213 }
1214 if self.row_limit.is_some_and(|l| self.rows >= l) {
1215 return Ok(false);
1216 }
1217 let b = match self.row_limit {
1218 Some(l) if self.rows + b.num_rows() as u64 > l => b.slice(0, (l - self.rows) as usize),
1219 _ => b,
1220 };
1221 if b.num_rows() > 0 || !self.sent_any {
1222 self.rows += b.num_rows() as u64;
1223 self.out.batch(&b)?;
1224 self.sent_any = true;
1225 }
1226 Ok(self.row_limit.is_none_or(|l| self.rows < l))
1227 }
1228}
1229
1230fn negotiate(core: (u32, u32), ours: (u32, u32)) -> Option<u32> {
1232 let (lo, hi) = (core.0.max(ours.0), core.1.min(ours.1));
1233 (lo <= hi).then_some(hi)
1234}
1235
1236pub fn conn_str<'a>(c: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
1238 c.get(key).and_then(Value::as_str).filter(|s| !s.is_empty())
1239}
1240
1241pub fn conn_required<'a>(c: &'a Map<String, Value>, key: &str) -> Result<&'a str> {
1243 conn_str(c, key).ok_or_else(|| format!("the profile output needs a `{key}` field").into())
1244}
1245
1246pub fn conn_bool(c: &Map<String, Value>, key: &str) -> Option<bool> {
1248 match c.get(key)? {
1249 Value::Bool(b) => Some(*b),
1250 Value::String(s) => Some(matches!(s.as_str(), "true" | "yes" | "1")),
1251 _ => None,
1252 }
1253}