Mirror/brk
1
0
Fork 0
mirror of https://github.com/bitcoinresearchkit/brk.git synced 2026-05-05 19:59:09 -07:00
Code Issues Packages Projects Releases Wiki Activity

cli + query: improvements

Browse Source
This commit is contained in:
nym21
2025-03-02 20:15:01 +01:00
parent ceefc8ffc6
commit be2012f28d
15 changed files with 308 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/brk_cli/Cargo.toml
View File

@@ -14,9 +14,11 @@ brk_logger = { workspace = true }
brk_parser = { workspace = true }
brk_query = { workspace = true }
brk_server = { workspace = true }
clap = { workspace = true , features = ["string"] }
clap = { workspace = true, features = ["string"] }
color-eyre = { workspace = true }
log = { workspace = true }
tabled = { workspace = true }
terminal_size = "0.4.1"
[[bin]]
name = "brk"

85
crates/brk_cli/src/main.rs
View File

@@ -1,13 +1,14 @@
use std::{path::Path, thread::sleep, time::Duration};
use std::path::Path;
use brk_computer::Computer;
use brk_exit::Exit;
use brk_indexer::Indexer;
use brk_parser::rpc::{self, RpcApi};
use brk_query::{Index, Params as QueryParams, Query};
use brk_server::tokio;
use brk_query::Params as QueryArgs;
use clap::{Parser, Subcommand};
use log::info;
use query::query;
use run::{RunArgs, run};
mod query;
mod run;
#[derive(Parser)]
#[command(version, about)]
@@ -19,13 +20,10 @@ struct Cli {
#[derive(Subcommand, Debug)]
enum Commands {
/// Run the indexer, computer and server
Run(RunArgs),
Query(QueryParams),
}
#[derive(Parser, Debug)]
struct RunArgs {
name: Option<String>,
/// Query generated datasets via the `run` command in a similar fashion as the server's API
Query(QueryArgs),
}
fn main() -> color_eyre::Result<()> {
@@ -37,67 +35,12 @@ fn main() -> color_eyre::Result<()> {
let outputs_dir = Path::new("../../_outputs");
let mut indexer = Indexer::import(&outputs_dir.join("indexed"))?;
let indexer = Indexer::import(&outputs_dir.join("indexed"))?;
let mut computer = Computer::import(&outputs_dir.join("computed"))?;
let computer = Computer::import(&outputs_dir.join("computed"))?;
match &cli.command {
Commands::Run(_) => {
let data_dir = Path::new("../../../bitcoin");
let rpc = Box::leak(Box::new(rpc::Client::new(
"http://localhost:8332",
rpc::Auth::CookieFile(Path::new(data_dir).join(".cookie")),
)?));
let exit = Exit::new();
let parser = brk_parser::Parser::new(data_dir, rpc);
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?
.block_on(async {
let served_indexer = indexer.clone();
let served_computer = computer.clone();
tokio::spawn(async move {
brk_server::main(served_indexer, served_computer).await.unwrap();
});
loop {
let block_count = rpc.get_block_count()?;
info!("{block_count} blocks found.");
let starting_indexes = indexer.index(&parser, rpc, &exit)?;
computer.compute(&mut indexer, starting_indexes, &exit)?;
info!("Waiting for new blocks...");
while block_count == rpc.get_block_count()? {
sleep(Duration::from_secs(1))
}
}
#[allow(unreachable_code)]
Ok(())
})
}
Commands::Query(args) => {
let query = Query::build(&indexer, &computer);
println!(
"{}",
query.search(
Index::try_from(args.index.as_str())?,
&args.values.iter().flat_map(|v| v.split(",")).collect::<Vec<_>>(),
args.from,
args.to,
args.format
)?
);
Ok(())
}
Commands::Run(_) => run(indexer, computer),
Commands::Query(args) => query(indexer, computer, args),
}
}

40
crates/brk_cli/src/query.rs Normal file
View File

@@ -0,0 +1,40 @@
use brk_computer::Computer;
use brk_indexer::Indexer;
use brk_query::{Index, Output, Params as QueryParams, Query, Tabled, Value};
use tabled::settings::Style;
pub fn query(indexer: Indexer, computer: Computer, params: &QueryParams) -> color_eyre::Result<()> {
let query = Query::build(&indexer, &computer);
let ids = params.values.iter().flat_map(|v| v.split(",")).collect::<Vec<_>>();
let index = Index::try_from(params.index.as_str())?;
let res = query.search(index, &ids, params.from, params.to, params.format)?;
if params.format.is_some() {
println!("{}", res);
} else {
println!(
"{}",
match res {
Output::Json(v) => match v {
Value::Single(v) => v.to_string(),
v => {
let v = match v {
Value::Single(_) => unreachable!("Already processed"),
Value::List(v) => vec![v],
Value::Matrix(v) => v,
};
let mut table = v.to_table(ids.iter().map(|id| id.to_string()).collect::<Vec<_>>());
table.with(Style::psql());
table.to_string()
}
},
_ => unreachable!(),
}
);
}
Ok(())
}

58
crates/brk_cli/src/run.rs Normal file
View File

@@ -0,0 +1,58 @@
use std::{path::Path, thread::sleep, time::Duration};
use brk_computer::Computer;
use brk_exit::Exit;
use brk_indexer::Indexer;
use brk_parser::rpc::{self, RpcApi};
use brk_server::tokio;
use clap::Parser;
use log::info;
#[derive(Parser, Debug)]
pub struct RunArgs {
name: Option<String>,
}
pub fn run(mut indexer: Indexer, mut computer: Computer) -> color_eyre::Result<()> {
let data_dir = Path::new("../../../bitcoin");
let rpc = Box::leak(Box::new(rpc::Client::new(
"http://localhost:8332",
rpc::Auth::CookieFile(Path::new(data_dir).join(".cookie")),
)?));
let exit = Exit::new();
let parser = brk_parser::Parser::new(data_dir, rpc);
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?
.block_on(async {
let served_indexer = indexer.clone();
let served_computer = computer.clone();
tokio::spawn(async move {
brk_server::main(served_indexer, served_computer).await.unwrap();
});
loop {
let block_count = rpc.get_block_count()?;
info!("{block_count} blocks found.");
let starting_indexes = indexer.index(&parser, rpc, &exit)?;
computer.compute(&mut indexer, starting_indexes, &exit)?;
info!("Waiting for new blocks...");
while block_count == rpc.get_block_count()? {
sleep(Duration::from_secs(1))
}
}
#[allow(unreachable_code)]
Ok(())
})
}

1
crates/brk_query/Cargo.toml
View File

@@ -15,3 +15,4 @@ color-eyre = { workspace = true }
derive_deref = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tabled = { workspace = true }

5
crates/brk_query/src/format.rs
View File

@@ -8,6 +8,7 @@ pub enum Format {
JSON,
CSV,
TSV,
MD,
}
impl TryFrom<Option<String>> for Format {
@@ -16,7 +17,9 @@ impl TryFrom<Option<String>> for Format {
if let Some(value) = value {
let value = value.to_lowercase();
let value = value.as_str();
if value == "csv" {
if value == "md" || value == "markdown" {
Ok(Self::MD)
} else if value == "csv" {
Ok(Self::CSV)
} else if value == "tsv" {
Ok(Self::TSV)

2
crates/brk_query/src/index.rs
View File

@@ -56,7 +56,7 @@ impl Index {
}
}
pub fn ids() -> Vec<String> {
pub fn possible_values() -> Vec<String> {
Self::all()
.iter()
.flat_map(|i| i.self_to_ids().iter().map(|s| s.to_string()))

89
crates/brk_query/src/lib.rs
View File

@@ -3,19 +3,22 @@
#![doc = include_str!("main.rs")]
#![doc = "```"]
mod format;
mod index;
mod params;
mod tree;
use std::fmt;
use brk_computer::Computer;
use brk_indexer::Indexer;
use tabled::settings::Style;
mod format;
mod index;
mod output;
mod params;
mod table;
mod tree;
pub use format::Format;
pub use index::Index;
pub use output::{Output, Value};
pub use params::Params;
use serde::Serialize;
pub use table::Tabled;
use tree::VecIdToIndexToVec;
pub struct Query<'a> {
@@ -41,12 +44,12 @@ impl<'a> Query<'a> {
pub fn search(
&self,
index: Index,
values: &[&str],
ids: &[&str],
from: Option<i64>,
to: Option<i64>,
format: Option<Format>,
) -> color_eyre::Result<QueryResponse> {
let ids = values
) -> color_eyre::Result<Output> {
let tuples = ids
.iter()
.map(|s| {
(
@@ -58,27 +61,27 @@ impl<'a> Query<'a> {
.map(|(id, vec)| (id, vec.unwrap()))
.collect::<Vec<_>>();
if ids.is_empty() {
return Ok(QueryResponse::default(format));
if tuples.is_empty() {
return Ok(Output::default(format));
}
let mut values = ids
let mut values = tuples
.iter()
.flat_map(|(_, i_to_v)| i_to_v.get(&index))
.map(|vec| -> brk_vec::Result<Vec<serde_json::Value>> { vec.collect_range_values(from, to) })
.collect::<brk_vec::Result<Vec<_>>>()?;
if values.is_empty() {
return Ok(QueryResponse::default(format));
return Ok(Output::default(format));
}
let ids_last_i = ids.len() - 1;
let ids_last_i = tuples.len() - 1;
Ok(match format {
Some(Format::CSV) | Some(Format::TSV) => {
let delimiter = if format == Some(Format::CSV) { ',' } else { '\t' };
let mut text = ids
let mut text = tuples
.into_iter()
.map(|(id, _)| id)
.collect::<Vec<_>>()
@@ -102,59 +105,31 @@ impl<'a> Query<'a> {
});
if format == Some(Format::CSV) {
QueryResponse::CSV(text)
Output::CSV(text)
} else {
QueryResponse::TSV(text)
Output::TSV(text)
}
}
Some(Format::MD) => {
let mut table = values.to_table(ids.iter().map(|s| s.to_string()).collect::<Vec<_>>());
table.with(Style::markdown());
Output::MD(table.to_string())
}
Some(Format::JSON) | None => {
if values.len() == 1 {
let mut values = values.pop().unwrap();
if values.len() == 1 {
let value = values.pop().unwrap();
QueryResponse::Json(Value::Single(value))
Output::Json(Value::Single(value))
} else {
QueryResponse::Json(Value::List(values))
Output::Json(Value::List(values))
}
} else {
QueryResponse::Json(Value::Matrix(values))
Output::Json(Value::Matrix(values))
}
}
})
}
}
#[derive(Debug)]
pub enum QueryResponse {
Json(Value),
CSV(String),
TSV(String),
}
#[derive(Debug, Serialize)]
#[serde(untagged)]
pub enum Value {
Matrix(Vec<Vec<serde_json::Value>>),
List(Vec<serde_json::Value>),
Single(serde_json::Value),
}
impl QueryResponse {
fn default(format: Option<Format>) -> Self {
match format {
Some(Format::CSV) => QueryResponse::CSV("".to_string()),
Some(Format::TSV) => QueryResponse::TSV("".to_string()),
_ => QueryResponse::Json(Value::Single(serde_json::Value::Null)),
}
}
}
impl fmt::Display for QueryResponse {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Json(value) => write!(f, "{}", serde_json::to_string_pretty(value).unwrap()),
Self::CSV(string) => write!(f, "{}", string),
Self::TSV(string) => write!(f, "{}", string),
}
}
}

43
crates/brk_query/src/output.rs Normal file
View File

@@ -0,0 +1,43 @@
use std::fmt;
use serde::Serialize;
use tabled::Tabled as TabledTabled;
use crate::Format;
#[derive(Debug)]
pub enum Output {
Json(Value),
CSV(String),
TSV(String),
MD(String),
}
#[derive(Debug, Serialize, TabledTabled)]
#[serde(untagged)]
pub enum Value {
Matrix(Vec<Vec<serde_json::Value>>),
List(Vec<serde_json::Value>),
Single(serde_json::Value),
}
impl Output {
pub fn default(format: Option<Format>) -> Self {
match format {
Some(Format::CSV) => Output::CSV("".to_string()),
Some(Format::TSV) => Output::TSV("".to_string()),
_ => Output::Json(Value::Single(serde_json::Value::Null)),
}
}
}
impl fmt::Display for Output {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Json(value) => write!(f, "{}", serde_json::to_string_pretty(value).unwrap()),
Self::CSV(string) => write!(f, "{}", string),
Self::TSV(string) => write!(f, "{}", string),
Self::MD(string) => write!(f, "{}", string),
}
}
}

7
crates/brk_query/src/params.rs
View File

@@ -5,14 +5,19 @@ use crate::{Format, Index};
#[derive(Debug, Deserialize, Parser)]
pub struct Params {
#[clap(short, long, value_parser = PossibleValuesParser::new(Index::ids()))]
#[clap(short, long, value_parser = PossibleValuesParser::new(Index::possible_values()))]
/// Index of the values requested
pub index: String,
#[clap(short, long, value_delimiter = ' ', num_args = 1..)]
/// Names of the values requested
pub values: Vec<String>,
#[clap(short, long, allow_hyphen_values = true)]
/// Inclusive starting index, if negative will be from the end
pub from: Option<i64>,
#[clap(short, long, allow_hyphen_values = true)]
/// Inclusive ending index, if negative will be from the end
pub to: Option<i64>,
#[clap(long)]
/// Format of the output
pub format: Option<Format>,
}

23
crates/brk_query/src/table.rs Normal file
View File

@@ -0,0 +1,23 @@
use tabled::{Table, builder::Builder};
pub trait Tabled {
fn to_table(&self, ids: Vec<String>) -> Table;
}
impl Tabled for Vec<Vec<serde_json::Value>> {
fn to_table(&self, ids: Vec<String>) -> Table {
let mut builder = Builder::default();
builder.push_record(ids);
if let Some(first) = self.first() {
let len = first.len();
(0..len).for_each(|index| {
builder.push_record(self.iter().map(|vec| vec.get(index).unwrap().to_string()));
});
}
builder.build()
}
}

2
crates/brk_server/src/api/vecs/mod.rs
View File

@@ -122,6 +122,7 @@ fn req_to_response_res(
csv.into_response()
}
Some(Format::MD) => "".into_response(),
Some(Format::JSON) | None => {
if values.len() == 1 {
let values = values.first().unwrap();
@@ -147,6 +148,7 @@ fn req_to_response_res(
headers.insert_content_disposition_attachment();
match format {
Format::CSV => headers.insert_content_type_text_csv(),
Format::MD => headers.insert_content_type_text_plain(),
Format::TSV => headers.insert_content_type_text_tsv(),
Format::JSON => headers.insert_content_type_application_json(),
}

10
crates/brk_vec/src/lib.rs
View File

@@ -325,15 +325,15 @@ where
}
});
let to = to.map_or(len, |to| {
let to = to.map_or(len - 1, |to| {
if to >= 0 {
to as usize
} else {
(len as i64 + to) as usize
((len - 1) as i64 + to) as usize
}
});
if from >= to {
if from > to {
return Err(Error::RangeFromAfterTo);
}
@@ -341,8 +341,8 @@ where
let mut buf = Self::create_buffer();
Ok((from..to)
.map(|_| Self::read_exact(&mut file, &mut buf).map(|v| v.to_owned()).unwrap())
Ok((from..=to)
.flat_map(|_| Self::read_exact(&mut file, &mut buf).map(|v| v.to_owned()))
.collect::<Vec<_>>())
}