Skip to content

Commit

Permalink
feat: add binary data type support to bwrite command
Browse files Browse the repository at this point in the history
  • Loading branch information
tmyoda committed Oct 23, 2023
1 parent 9f58756 commit 09898f2
Show file tree
Hide file tree
Showing 5 changed files with 67 additions and 3 deletions.
7 changes: 4 additions & 3 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ bytes = "1.4.0"
itertools = "0.10.5"
console = "0.15.7"
brotli = "3.3.4"
base64 = "0.21.5"

[dev-dependencies]
assert_cmd = "2.0.2" # contains helpers make executing the main binary on integration tests easier.
Expand Down
28 changes: 28 additions & 0 deletions src/batch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
* limitations under the License.
*/

use base64::{engine::general_purpose, Engine as _};
use bytes::Bytes;
use log::{debug, error};
use rusoto_core::RusotoError;
use rusoto_dynamodb::{
Expand Down Expand Up @@ -431,6 +433,11 @@ fn ddbjson_val_to_attrval(ddb_jsonval: &JsonValue) -> Option<AttributeValue> {
n: Some(x.as_str().unwrap().to_string()),
..Default::default()
})
} else if let Some(x) = ddb_jsonval.get("B") {
Some(AttributeValue {
b: Some(json_binary_val_to_bytes(x)),
..Default::default()
})
} else if let Some(x) = ddb_jsonval.get("BOOL") {
Some(AttributeValue {
bool: Some(x.as_bool().unwrap()),
Expand All @@ -446,6 +453,18 @@ fn ddbjson_val_to_attrval(ddb_jsonval: &JsonValue) -> Option<AttributeValue> {
ns: Some(set_logic(x)),
..Default::default()
})
} else if let Some(x) = ddb_jsonval.get("BS") {
let binary_set = x
.as_array()
.expect("should be valid JSON array")
.iter()
.map(json_binary_val_to_bytes)
.collect::<Vec<Bytes>>();
debug!("Binary Set: {:?}", binary_set);
Some(AttributeValue {
bs: Some(binary_set),
..Default::default()
})
} else if let Some(x) = ddb_jsonval.get("L") {
let list_element = x
.as_array()
Expand Down Expand Up @@ -473,3 +492,12 @@ fn ddbjson_val_to_attrval(ddb_jsonval: &JsonValue) -> Option<AttributeValue> {
None
}
}

// Decodes a base64 encoded binary value to Bytes.
fn json_binary_val_to_bytes(v: &JsonValue) -> Bytes {
Bytes::from(
general_purpose::STANDARD
.decode(v.as_str().expect("binary inputs should be string value"))
.expect("binary inputs should be base64 with padding encoded"),
)
}
32 changes: 32 additions & 0 deletions tests/bwrite.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ pub mod util;

use assert_cmd::prelude::*; // Add methods on commands
use predicates::prelude::*; // Used for writing assertions
use serde_json::Value;
use std::collections::HashSet;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
Expand Down Expand Up @@ -108,5 +110,35 @@ async fn test_batch_write() -> Result<(), Box<dyn std::error::Error>> {
predicate::str::is_match("\"Dimensions\":")?.eval(String::from_utf8(output)?.as_str())
);

let mut c = tm.command()?;
let get_cmd = c.args(&[
"--region",
"local",
"--table",
&table_name,
"get",
"ichi",
"-o",
"raw",
]);

let output = get_cmd.output()?.stdout;
let data: Value = serde_json::from_str(&String::from_utf8(output)?)?;

let binary = data["Binary"]["B"].as_str().unwrap();
assert_eq!(binary, "dGhpcyB0ZXh0IGlzIGJhc2U2NC1lbmNvZGVk");

// The order of the values within a set is not preserved, so I will use HashSet to compare them.
let binary_set: HashSet<String> = data["BinarySet"]["BS"]
.as_array()
.unwrap()
.iter()
.filter_map(Value::as_str)
.map(|s| s.to_string())
.collect();
let binary_set_expected: HashSet<String> =
HashSet::from(["U3Vubnk=", "UmFpbnk=", "U25vd3k="].map(String::from));
assert_eq!(binary_set, binary_set_expected);

Ok(())
}
2 changes: 2 additions & 0 deletions tests/resources/test_batch_write.json
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
"Dimensions": { "SS": ["Giraffe", "Hippo" ,"Zebra"] },
"PageCount": { "NS": ["42.2", "-19", "7.5", "3.14"] },
"InPublication": { "BOOL": false },
"Binary": {"B": "dGhpcyB0ZXh0IGlzIGJhc2U2NC1lbmNvZGVk"},
"BinarySet": {"BS": ["U3Vubnk=", "UmFpbnk=", "U25vd3k="]},
"Nothing": { "NULL": true },
"Authors": {
"L": [
Expand Down

0 comments on commit 09898f2

Please sign in to comment.