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