use std::{
fs::File,
io::{self, BufWriter},
num::NonZeroU64,
path::{Path, PathBuf},
};
use anyhow::{anyhow, Context};
use clap::{ArgAction, Parser, ValueEnum};
use dbn::{
encode::SplitDuration,
enums::{Compression, Encoding},
Schema, VersionUpgradePolicy,
};
pub mod encode;
pub mod filter;
#[derive(Clone, Copy, Debug, ValueEnum)]
pub enum OutputEncoding {
Infer,
Dbn,
Csv,
Tsv,
Json,
DbnFragment,
}
#[derive(Clone, Copy, Debug, ValueEnum)]
pub enum SplitBy {
Symbol,
Schema,
Day,
Week,
Month,
}
impl SplitBy {
pub fn duration(self) -> Option<SplitDuration> {
match self {
SplitBy::Day => Some(SplitDuration::Day),
SplitBy::Week => Some(SplitDuration::Week),
SplitBy::Month => Some(SplitDuration::Month),
SplitBy::Symbol | SplitBy::Schema => None,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct InferredEncoding {
pub encoding: Encoding,
pub compression: Compression,
pub delimiter: u8,
pub is_fragment: bool,
}
#[derive(Debug, Parser)]
#[clap(name = "dbn", version, about)]
#[cfg_attr(test, derive(Default))]
pub struct Args {
#[clap(
help = "One or more DBN or legacy DBZ files to decode. Passing multiple files will result in a merge. Pass '-' to read from standard input",
value_name = "FILE...",
value_delimiter = ' ',
num_args = 1..,
required = true,
)]
pub input: Vec<PathBuf>,
#[clap(
short,
long,
help = "Saves the result to FILE. If no path is specified, the output will be written to standard output",
value_name = "FILE"
)]
pub output: Option<PathBuf>,
#[clap(
short = 'O',
long,
help = "Saves the result of file splitting to paths according to PATTERN",
requires = "split_by",
conflicts_with = "output",
value_name = "PATTERN"
)]
pub output_pattern: Option<String>,
#[clap(
short = 'S',
long,
help = "How to optionally split the output across files",
requires = "output_pattern",
value_name = "SPLIT_BY"
)]
pub split_by: Option<SplitBy>,
#[clap(
short = 'J',
long,
action = ArgAction::SetTrue,
default_value = "false",
group = "output_encoding",
help = "Output the result as JSON lines"
)]
pub json: bool,
#[clap(
short = 'C',
long,
action = ArgAction::SetTrue,
default_value = "false",
group = "output_encoding",
help = "Output the result as CSV"
)]
pub csv: bool,
#[clap(
short = 'T',
long,
action = ArgAction::SetTrue,
default_value = "false",
group = "output_encoding",
help = "Output the result as tab-separated values (TSV)"
)]
pub tsv: bool,
#[clap(
short = 'D',
long,
action = ArgAction::SetTrue,
default_value = "false",
group = "output_encoding",
help = "Output the result as DBN"
)]
pub dbn: bool,
#[clap(
short = 'F',
long,
action = ArgAction::SetTrue,
default_value = "false",
group = "output_encoding",
help = "Output the result as a DBN fragment (no metadata)"
)]
pub fragment: bool,
#[clap(short, long, action = ArgAction::SetTrue, default_value = "false", help = "Zstd compress the output")]
pub zstd: bool,
#[clap(
short = 'u',
long = "upgrade",
default_value = "false",
action = ArgAction::SetTrue,
help = "Upgrade data when decoding previous DBN versions. By default data is decoded as-is."
)]
pub should_upgrade: bool,
#[clap(
short,
long,
action = ArgAction::SetTrue,
default_value = "false",
help = "Allow overwriting of existing files, such as the output file"
)]
pub force: bool,
#[clap(
short = 'm',
long = "metadata",
action = ArgAction::SetTrue,
default_value = "false",
conflicts_with_all = ["csv", "dbn", "fragment"],
help = "Output the metadata section instead of the body of the DBN file. Only valid for JSON output encoding"
)]
pub should_output_metadata: bool,
#[clap(
short = 'p',
long = "pretty",
action = ArgAction::SetTrue,
default_value = "false",
conflicts_with_all = ["dbn", "fragment"],
help ="Make the CSV or JSON output easier to read by converting timestamps to ISO 8601 and prices to decimals"
)]
pub should_pretty_print: bool,
#[clap(
short = 's',
long = "map-symbols",
action = ArgAction::SetTrue,
default_value = "false",
conflicts_with_all = ["input_fragment", "dbn", "fragment"],
help ="Use symbology mappings from the metadata to create a 'symbol' field mapping the instrument ID to its requested symbol."
)]
pub map_symbols: bool,
#[clap(
short = 'l',
long = "limit",
value_name = "NUM_RECORDS",
help = "Limit the number of records in the output to the specified number"
)]
pub limit: Option<NonZeroU64>,
#[clap(
long = "input-fragment",
action = ArgAction::SetTrue,
default_value = "false",
group = "input_fragment",
conflicts_with_all = ["is_input_zstd_fragment", "should_output_metadata", "dbn"],
help = "Interpret the input as an uncompressed DBN fragment, i.e. records without metadata. Only valid with text output encodings"
)]
pub is_input_fragment: bool,
#[clap(
long = "input-zstd-fragment",
action = ArgAction::SetTrue,
default_value = "false",
group = "input_fragment",
conflicts_with_all = ["should_output_metadata", "dbn"],
help = "Interpret the input as a Zstd-compressed DBN fragment, i.e. records without metadata. Only valid with text output encodings"
)]
pub is_input_zstd_fragment: bool,
#[clap(
long = "input-dbn-version",
help = "Specify the DBN version of the fragment. By default the fragment is assumed to be of the current version",
value_name = "DBN_VERSION",
value_parser = clap::value_parser!(u8).range(1..=3),
requires = "input_fragment"
)]
pub input_dbn_version_override: Option<u8>,
#[clap(
long = "schema",
help = "Only encode records of this schema. This is particularly useful for transcoding mixed-schema DBN to CSV, which doesn't support mixing schemas",
value_name = "SCHEMA"
)]
pub schema_filter: Option<Schema>,
#[clap(
long = "omit-header",
action = ArgAction::SetFalse,
default_value = "true",
conflicts_with_all = ["json", "dbn", "fragment"],
help = "Skip encoding the header. Only valid when encoding CSV or TSV."
)]
pub write_header: bool,
}
impl Args {
pub fn output_encoding(&self) -> OutputEncoding {
if self.json {
OutputEncoding::Json
} else if self.csv {
OutputEncoding::Csv
} else if self.tsv {
OutputEncoding::Tsv
} else if self.dbn {
OutputEncoding::Dbn
} else if self.fragment {
OutputEncoding::DbnFragment
} else {
OutputEncoding::Infer
}
}
pub fn upgrade_policy(&self) -> VersionUpgradePolicy {
if self.should_upgrade {
VersionUpgradePolicy::UpgradeToV3
} else {
VersionUpgradePolicy::AsIs
}
}
pub fn input_version(&self) -> u8 {
self.input_dbn_version_override.unwrap_or(dbn::DBN_VERSION)
}
}
pub fn infer_encoding(args: &Args) -> anyhow::Result<InferredEncoding> {
let compression = if args.zstd {
Compression::Zstd
} else {
Compression::None
};
match args.output_encoding() {
OutputEncoding::DbnFragment => Ok(InferredEncoding {
encoding: Encoding::Dbn,
compression,
delimiter: 0,
is_fragment: true,
}),
OutputEncoding::Dbn => Ok(InferredEncoding {
encoding: Encoding::Dbn,
compression,
delimiter: 0,
is_fragment: false,
}),
OutputEncoding::Csv => Ok(InferredEncoding {
encoding: Encoding::Csv,
compression,
delimiter: b',',
is_fragment: false,
}),
OutputEncoding::Tsv => Ok(InferredEncoding {
encoding: Encoding::Csv,
compression,
delimiter: b'\t',
is_fragment: false,
}),
OutputEncoding::Json => Ok(InferredEncoding {
encoding: Encoding::Json,
compression,
delimiter: 0,
is_fragment: false,
}),
OutputEncoding::Infer => {
let output = args
.output
.as_ref()
.map(|p| p.to_string_lossy().into_owned())
.or_else(|| args.output_pattern.clone());
if let Some(output) = output {
if output.ends_with(".dbn.frag.zst") {
Ok(InferredEncoding {
encoding: Encoding::Dbn,
compression: Compression::Zstd,
delimiter: 0,
is_fragment: true,
})
} else if output.ends_with(".dbn.frag") {
Ok(InferredEncoding {
encoding: Encoding::Dbn,
compression: Compression::None,
delimiter: 0,
is_fragment: true,
})
} else if output.ends_with(".dbn.zst") {
Ok(InferredEncoding {
encoding: Encoding::Dbn,
compression: Compression::Zstd,
delimiter: 0,
is_fragment: false,
})
} else if output.ends_with(".dbn") {
Ok(InferredEncoding {
encoding: Encoding::Dbn,
compression: Compression::None,
delimiter: 0,
is_fragment: false,
})
} else if output.ends_with(".csv.zst") {
Ok(InferredEncoding {
encoding: Encoding::Csv,
compression: Compression::Zstd,
delimiter: b',',
is_fragment: false,
})
} else if output.ends_with(".csv") {
Ok(InferredEncoding {
encoding: Encoding::Csv,
compression: Compression::None,
delimiter: b',',
is_fragment: false,
})
} else if output.ends_with(".tsv.zst") || output.ends_with(".xls.zst") {
Ok(InferredEncoding {
encoding: Encoding::Csv,
compression: Compression::Zstd,
delimiter: b'\t',
is_fragment: false,
})
} else if output.ends_with(".tsv") || output.ends_with(".xls") {
Ok(InferredEncoding {
encoding: Encoding::Csv,
compression: Compression::None,
delimiter: b'\t',
is_fragment: false,
})
} else if output.ends_with(".json.zst") {
Ok(InferredEncoding {
encoding: Encoding::Json,
compression: Compression::Zstd,
delimiter: 0,
is_fragment: false,
})
} else if output.ends_with(".json") {
Ok(InferredEncoding {
encoding: Encoding::Json,
compression: Compression::None,
delimiter: 0,
is_fragment: false,
})
} else {
Err(anyhow!(
"Unable to infer output encoding from output path '{output}'",
))
}
} else {
Err(anyhow!(
"Unable to infer output encoding when no output was specified"
))
}
}
}
}
pub fn output_from_args(args: &Args) -> anyhow::Result<Box<dyn io::Write>> {
output(args.output.as_deref(), args.force)
}
pub fn output(output: Option<&Path>, force: bool) -> anyhow::Result<Box<dyn io::Write>> {
if let Some(output) = output {
let output_file = open_output_file(output, force)?;
Ok(Box::new(BufWriter::new(output_file)))
} else {
Ok(Box::new(io::stdout().lock()))
}
}
fn open_output_file(path: &Path, force: bool) -> anyhow::Result<File> {
let mut options = File::options();
options.write(true).truncate(true);
if force {
options.create(true);
} else if path.exists() {
return Err(anyhow!(
"Output file exists. Pass --force flag to overwrite the existing file."
));
} else {
options.create_new(true);
}
options
.open(path)
.with_context(|| format!("Unable to open output file '{}'", path.display()))
}
#[cfg(test)]
mod tests {
#![allow(clippy::too_many_arguments)]
use rstest::*;
use super::*;
#[rstest]
#[case(true, false, false, false, false, Encoding::Json, Compression::None, 0)]
#[case(
false,
true,
false,
false,
false,
Encoding::Csv,
Compression::None,
b','
)]
#[case(
false,
false,
true,
false,
false,
Encoding::Csv,
Compression::None,
b'\t'
)]
#[case(false, false, false, true, false, Encoding::Dbn, Compression::None, 0)]
#[case(true, false, false, false, true, Encoding::Json, Compression::Zstd, 0)]
#[case(
false,
true,
false,
false,
true,
Encoding::Csv,
Compression::Zstd,
b','
)]
#[case(
false,
false,
true,
false,
true,
Encoding::Csv,
Compression::Zstd,
b'\t'
)]
#[case(false, false, false, true, true, Encoding::Dbn, Compression::Zstd, 0)]
fn test_infer_encoding_and_compression_explicit(
#[case] json: bool,
#[case] csv: bool,
#[case] tsv: bool,
#[case] dbn: bool,
#[case] zstd: bool,
#[case] exp_enc: Encoding,
#[case] exp_comp: Compression,
#[case] exp_sep: u8,
) {
let args = Args {
json,
csv,
tsv,
dbn,
zstd,
..Default::default()
};
assert_eq!(
infer_encoding(&args).unwrap(),
InferredEncoding {
encoding: exp_enc,
compression: exp_comp,
delimiter: exp_sep,
is_fragment: false,
}
);
}
#[rstest]
#[case("out.json", Encoding::Json, Compression::None, 0)]
#[case("out.csv", Encoding::Csv, Compression::None, b',')]
#[case("out.tsv", Encoding::Csv, Compression::None, b'\t')]
#[case("out.xls", Encoding::Csv, Compression::None, b'\t')]
#[case("out.dbn", Encoding::Dbn, Compression::None, 0)]
#[case("out.json.zst", Encoding::Json, Compression::Zstd, 0)]
#[case("out.csv.zst", Encoding::Csv, Compression::Zstd, b',')]
#[case("out.tsv.zst", Encoding::Csv, Compression::Zstd, b'\t')]
#[case("out.xls.zst", Encoding::Csv, Compression::Zstd, b'\t')]
#[case("out.dbn.zst", Encoding::Dbn, Compression::Zstd, 0)]
fn test_infer_encoding_and_compression_inference(
#[case] output: &str,
#[case] exp_enc: Encoding,
#[case] exp_comp: Compression,
#[case] exp_sep: u8,
) {
let args = Args {
output: Some(PathBuf::from(output)),
..Default::default()
};
assert_eq!(
infer_encoding(&args).unwrap(),
InferredEncoding {
encoding: exp_enc,
compression: exp_comp,
delimiter: exp_sep,
is_fragment: false,
}
);
}
#[test]
fn test_infer_encoding_and_compression_bad() {
let args = Args {
output: Some(PathBuf::from("out.pb")),
..Default::default()
};
assert!(
matches!(infer_encoding(&args), Err(e) if e.to_string().starts_with("Unable to infer"))
);
}
}