UTF-8 decoding and handling

Cargo.toml

@@ -14,6 +14,7 @@ anyhow = "1"
 async-trait = "0"
 aws-config = { version = "0" }
 aws-sdk-kinesis = { version = "0" }
+base64 = "0"
 chrono = { version = "0", features = ["clock", "std"] }
 clap = { version = "4", features = ["derive"] }
 colored = "2"

README.md

@@ -51,6 +51,7 @@ The [release page](https://github.com/grumlimited/kinesis-tailr/releases) provid
         -o, --output-file <OUTPUT_FILE>      Output file to write to
         -c, --concurrent <CONCURRENT>        Concurrent number of shards to tail
         -v, --verbose                        Display additional information
+        -n, --no-base64                      Do not base64 encode the payload upon invalid UTF-8 payloads. Print it raw instead
         -h, --help                           Print help
         -V, --version                        Print version
 
@@ -65,6 +66,14 @@ kinesis-tailr \
     --max-messages 2
 ```
 
+### UTF-8
+
+`kinesis-tailr` expects payloads to be UTF-8 encoded. If a payload is not UTF-8 encoded, it will be base64 encoded and printed as such.
+
+It might be useful to print the raw payload instead though. This can be achieved with the `--no-base64` flag.
+
+Properly UTF-8 encoded payloads will be printed as such and never base64 encoded.
+
 ### Logging
 
 General logging level for debugging can be turned on with:

src/cli_helpers.rs

@@ -87,6 +87,10 @@ pub struct Opt {
     /// Display additional information
     #[structopt(short, long)]
     pub verbose: bool,
+
+    /// Do not base64 encode the payload upon invalid UTF-8 payloads. Print it raw instead.
+    #[structopt(short, long)]
+    pub no_base64: bool,
 }
 
 pub(crate) fn selected_shards(

src/main.rs

@@ -67,6 +67,7 @@ async fn main() -> Result<()> {
                         opt.print_shard_id,
                         opt.print_timestamp,
                         opt.print_delimiter,
+                        opt.no_base64,
                         shard_count,
                         file,
                     )
@@ -82,6 +83,7 @@ async fn main() -> Result<()> {
                         opt.print_shard_id,
                         opt.print_timestamp,
                         opt.print_delimiter,
+                        opt.no_base64,
                         shard_count,
                     )
                     .run(tx_records, rx_records)

src/sink.rs

@@ -1,10 +1,11 @@
+use std::io;
+use std::io::{BufWriter, Write};
+
 use anyhow::Error;
 use anyhow::Result;
 use async_trait::async_trait;
 use chrono::TimeZone;
 use log::{debug, error, warn};
-use std::io;
-use std::io::{BufWriter, Write};
 use tokio::sync::mpsc::{Receiver, Sender};
 
 use crate::kinesis::models::{ProcessError, RecordResult, ShardProcessorADT};
@@ -22,6 +23,7 @@ pub struct SinkConfig {
     print_timestamp: bool,
     print_delimiter: bool,
     exit_after_termination: bool,
+    no_base64: bool,
 }
 
 pub trait Configurable {
@@ -225,9 +227,16 @@ where
     }
 
     fn format_record(&self, record_result: &RecordResult) -> String {
-        let data = std::str::from_utf8(record_result.data.as_slice())
-            .unwrap()
-            .to_string();
+        let data = match std::str::from_utf8(record_result.data.as_slice()) {
+            Ok(payload) => payload.to_string(),
+            Err(_) if self.get_config().no_base64 => {
+                String::from_utf8_lossy(record_result.data.as_slice()).to_string()
+            }
+            Err(_) => {
+                use base64::{engine::general_purpose, Engine as _};
+                general_purpose::STANDARD.encode(record_result.data.as_slice())
+            }
+        };
 
         let data = if self.get_config().print_key {
             let key = record_result.partition_key.to_string();

src/sink/console.rs

@@ -20,6 +20,7 @@ impl ConsoleSink {
         print_shard_id: bool,
         print_timestamp: bool,
         print_delimiter: bool,
+        no_base64: bool,
         shard_count: usize,
     ) -> Self {
         ConsoleSink {
@@ -31,6 +32,7 @@ impl ConsoleSink {
                 print_shard_id,
                 print_timestamp,
                 print_delimiter,
+                no_base64,
                 exit_after_termination: true,
             },
             shard_count,

src/sink/console_tests.rs

@@ -7,16 +7,7 @@ use tokio::sync::mpsc;
 #[test]
 fn format_nb_messages_ok() {
     let console = ConsoleSink {
-        config: SinkConfig {
-            max_messages: None,
-            no_color: false,
-            print_key: false,
-            print_sequence_number: false,
-            print_shard_id: false,
-            print_timestamp: false,
-            print_delimiter: false,
-            exit_after_termination: false,
-        },
+        config: SinkConfig::default(),
         shard_count: 1,
     };
 
@@ -54,6 +45,52 @@ fn format_outputs() {
     assert_eq!(bw_console.write_delimiter("data"), "data");
 }
 
+#[test]
+fn format_outputs_base64() {
+    let console = ConsoleSink {
+        config: SinkConfig {
+            no_color: true,
+            ..Default::default()
+        },
+        shard_count: 1,
+    };
+
+    let input = b"Hello \xF0\x90\x80World";
+
+    let record = RecordResult {
+        shard_id: "shard_id".to_string(),
+        sequence_id: "sequence_id".to_string(),
+        partition_key: "partition_key".to_string(),
+        datetime: DateTime::from_secs(1_000_000_i64),
+        data: input.to_vec(),
+    };
+
+    assert_eq!(console.format_record(&record), "SGVsbG8g8JCAV29ybGQ=");
+}
+
+#[test]
+fn format_outputs_no_base64() {
+    let console = ConsoleSink {
+        config: SinkConfig {
+            no_base64: true,
+            ..Default::default()
+        },
+        shard_count: 1,
+    };
+
+    let input = b"Hello \xF0\x90\x80World";
+
+    let record = RecordResult {
+        shard_id: "shard_id".to_string(),
+        sequence_id: "sequence_id".to_string(),
+        partition_key: "partition_key".to_string(),
+        datetime: DateTime::from_secs(1_000_000_i64),
+        data: input.to_vec(),
+    };
+
+    assert_eq!(console.format_record(&record), "Hello �World");
+}
+
 #[tokio::test]
 async fn expect_zero_messages_processed() {
     let (tx_records, rx_records) = mpsc::channel::<Result<ShardProcessorADT, ProcessError>>(1);

src/sink/file.rs

@@ -23,6 +23,7 @@ impl FileSink {
         print_shard_id: bool,
         print_timestamp: bool,
         print_delimiter: bool,
+        no_base64: bool,
         shard_count: usize,
         file: P,
     ) -> Self {
@@ -35,6 +36,7 @@ impl FileSink {
                 print_shard_id,
                 print_timestamp,
                 print_delimiter,
+                no_base64,
                 exit_after_termination: true,
             },
             file: file.into(),