From a0f0372839d549bc6a4ea1b9c02aabbe9ae6651f Mon Sep 17 00:00:00 2001 From: Patrick Meredith Date: Sat, 24 Aug 2019 15:14:06 -0400 Subject: [PATCH] Add mostly working BSON support (missing some types) --- src/commands/from_bson.rs | 148 ++++++++++++++++++++++++++++++++++++++ src/commands/open.rs | 20 +++--- 2 files changed, 157 insertions(+), 11 deletions(-) create mode 100644 src/commands/from_bson.rs diff --git a/src/commands/from_bson.rs b/src/commands/from_bson.rs new file mode 100644 index 0000000000..d0f3e75485 --- /dev/null +++ b/src/commands/from_bson.rs @@ -0,0 +1,148 @@ +use crate::commands::WholeStreamCommand; +use crate::object::base::OF64; +use crate::object::{Primitive, TaggedDictBuilder, Value}; +use crate::prelude::*; +use bson::{decode_document, Bson}; + +pub struct FromBSON; + +impl WholeStreamCommand for FromBSON { + fn run( + &self, + args: CommandArgs, + registry: &CommandRegistry, + ) -> Result { + from_bson(args, registry) + } + + fn name(&self) -> &str { + "from-bson" + } + + fn signature(&self) -> Signature { + Signature::build("from-bson") + } +} + +fn convert_bson_value_to_nu_value(v: &Bson, tag: impl Into) -> Tagged { + let tag = tag.into(); + + match v { + Bson::FloatingPoint(n) => Value::Primitive(Primitive::Float(OF64::from(*n))).tagged(tag), + Bson::String(s) => Value::Primitive(Primitive::String(String::from(s))).tagged(tag), + Bson::Array(a) => Value::List( + a.iter() + .map(|x| convert_bson_value_to_nu_value(x, tag)) + .collect(), + ) + .tagged(tag), + Bson::Document(doc) => { + let mut collected = TaggedDictBuilder::new(tag); + for (k, v) in doc.iter() { + collected.insert_tagged(k.clone(), convert_bson_value_to_nu_value(v, tag)); + } + + collected.into_tagged_value() + } + Bson::Boolean(b) => Value::Primitive(Primitive::Boolean(*b)).tagged(tag), + Bson::Null => Value::Primitive(Primitive::String(String::from(""))).tagged(tag), + // Bson::RegExp(r, opts) => { + // let mut collected = TaggedDictBuilder::new(tag); + // collected.insert_tagged( + // "$regex".to_owned(), + // Value::Primitive(Primitive::String(String::from(r))), + // ); + // collected.insert_tagged( + // "$options".to_owned(), + // Value::Primitive(Primitive::String(String::from(opts))), + // ); + // collected.into_tagged_value() + // } + Bson::I32(n) => Value::Primitive(Primitive::Int(*n as i64)).tagged(tag), + Bson::I64(n) => Value::Primitive(Primitive::Int(*n as i64)).tagged(tag), + Bson::ObjectId(obj_id) => Value::Primitive(Primitive::String(obj_id.to_hex())).tagged(tag), + x => { + println!("{:?}", x); + panic!() + } + } +} + +#[derive(Debug)] +struct BytesReader { + pos: usize, + inner: Vec, +} + +impl BytesReader { + fn new(bytes: Vec) -> BytesReader { + BytesReader { + pos: 0, + inner: bytes, + } + } +} + +impl std::io::Read for BytesReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let src: &mut &[u8] = &mut self.inner[self.pos..].as_ref(); + let diff = src.read(buf)?; + self.pos += diff; + Ok(diff) + } +} + +pub fn from_bson_bytes_to_value( + bytes: Vec, + tag: impl Into + std::clone::Clone, +) -> bson::DecoderResult> { + let mut out = Vec::new(); + let mut b_reader = BytesReader::new(bytes); + while let Ok(v) = decode_document(&mut b_reader) { + out.push(convert_bson_value_to_nu_value( + &Bson::Document(v), + tag.clone(), + )); + } + Ok(Value::List(out).tagged(tag)) +} + +fn from_bson(args: CommandArgs, registry: &CommandRegistry) -> Result { + let args = args.evaluate_once(registry)?; + let span = args.name_span(); + let input = args.input; + + let stream = async_stream_block! { + let values: Vec> = input.values.collect().await; + + for value in values { + let value_tag = value.tag(); + latest_tag = Some(value_tag); + match value.item { + Value::Binary(vb) => + match from_bson_bytes_to_value(vb, span) { + Ok(x) => yield ReturnSuccess::value(x), + Err(_) => if let Some(last_tag) = latest_tag { + yield Err(ShellError::labeled_error_with_secondary( + "Could not parse as BSON", + "input cannot be parsed as BSON", + span, + "value originates from here", + last_tag.span, + )) + } + } + _ => yield Err(ShellError::labeled_error_with_secondary( + "Expected a string from pipeline", + "requires string input", + span, + "value originates from here", + value_tag.span, + )), + + } + } + }; + + Ok(stream.to_output_stream()) +} diff --git a/src/commands/open.rs b/src/commands/open.rs index 6fc3f01832..d499d6c73d 100644 --- a/src/commands/open.rs +++ b/src/commands/open.rs @@ -501,19 +501,17 @@ pub fn parse_binary_as_value( contents_tag: Tag, name_span: Span, ) -> Result, ShellError> { - println!("{:?}", extension); match extension { Some(x) if x == "bson" => { - Err(ShellError::labeled_error("Could not open as BSON", "Could not open as BSON", name_span)) - //crate::commands::from_json::from_bson_bytes_to_value(contents, contents_tag).map_err( - // move |_| { - // ShellError::labeled_error( - // "Could not open as BSON", - // "could not open as BSON", - // name_span, - // ) - // }, - // ) + crate::commands::from_bson::from_bson_bytes_to_value(contents, contents_tag).map_err( + move |_| { + ShellError::labeled_error( + "Could not open as BSON", + "could not open as BSON", + name_span, + ) + }, + ) } _ => Ok(Value::Binary(contents).tagged(contents_tag)), }