Skip to content

cargo fmt tower-cmd/ #337

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
socksy wants to merge 1 commit into
base: main
Choose a base branch
from
Open

cargo fmt tower-cmd/ #337

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 59 additions & 43 deletions crates/tower-cmd/src/catalogs.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -591,18 +591,19 @@ async fn execute_catalog_query(
// legitimate analytical scan over a large catalog can take minutes and a
// person is driving; the agent path uses `Limits::agent()`, which adds one.
let harden = !write;
let result = tokio::task::spawn_blocking(move || -> Result<QueryResult, tower_duckdb::Error> {
let session = Session::open()?;
session.run_setup(&setup)?;
if harden {
// `Hardening::agent()` adds an engine memory ceiling on top of the
// lockdown. The result `limits` bound what comes back; only the engine
// can bound what a query spends producing it.
session.harden(&Hardening::agent())?;
}
session.query(&sql, [], &limits)
})
.await;
let result =
tokio::task::spawn_blocking(move || -> Result<QueryResult, tower_duckdb::Error> {
let session = Session::open()?;
session.run_setup(&setup)?;
if harden {
// `Hardening::agent()` adds an engine memory ceiling on top of the
// lockdown. The result `limits` bound what comes back; only the engine
// can bound what a query spends producing it.
session.harden(&Hardening::agent())?;
}
session.query(&sql, [], &limits)
})
.await;

match result {
Ok(Ok(query_result)) => {
Expand Down Expand Up @@ -687,13 +688,14 @@ pub(crate) async fn query_catalog_for_agent(
&response.credentials,
vend_catalog_credentials_body::Mode::Read,
);
let result = tokio::task::spawn_blocking(move || -> Result<QueryResult, tower_duckdb::Error> {
let session = Session::open()?;
session.run_setup(&setup)?;
session.harden(&Hardening::agent())?;
session.query(&sql, [], &Limits::agent())
})
.await;
let result =
tokio::task::spawn_blocking(move || -> Result<QueryResult, tower_duckdb::Error> {
let session = Session::open()?;
session.run_setup(&setup)?;
session.harden(&Hardening::agent())?;
session.query(&sql, [], &Limits::agent())
})
.await;

match result {
Ok(Ok(query_result)) => Ok(query_result),
Expand Down Expand Up @@ -844,7 +846,10 @@ async fn fetch_catalog_columns_via_rest(
rows.sort_by(|a, b| {
let key = |row: &[serde_json::Value]| {
(
row.first().and_then(|v| v.as_str()).unwrap_or("").to_owned(),
row.first()
.and_then(|v| v.as_str())
.unwrap_or("")
.to_owned(),
row.get(1).and_then(|v| v.as_str()).unwrap_or("").to_owned(),
)
};
Expand Down Expand Up @@ -1088,7 +1093,10 @@ fn iceberg_type_to_display(ty: &serde_json::Value) -> String {
format!("LIST({})", element)
}
Some("map") => {
let key = obj.get("key").map(iceberg_type_to_display).unwrap_or_default();
let key = obj
.get("key")
.map(iceberg_type_to_display)
.unwrap_or_default();
let value = obj
.get("value")
.map(iceberg_type_to_display)
Expand Down Expand Up @@ -1592,9 +1600,7 @@ pub async fn do_knowledge_show(out: &output::Out, config: Config, args: &ArgMatc
let catalog = args
.get_one::<String>("catalog_name")
.expect("catalog_name is required");
let name = args
.get_one::<String>("name")
.expect("name is required");
let name = args.get_one::<String>("name").expect("name is required");
let env = cmd::get_string_flag(args, "environment");

let response = out
Expand All @@ -1614,9 +1620,7 @@ pub async fn do_knowledge_set(out: &output::Out, config: Config, args: &ArgMatch
let catalog = args
.get_one::<String>("catalog_name")
.expect("catalog_name is required");
let name = args
.get_one::<String>("name")
.expect("name is required");
let name = args.get_one::<String>("name").expect("name is required");
let env = cmd::get_string_flag(args, "environment");
let statement = cmd::get_string_flag(args, "statement");
let scope = cmd::get_string_flag(args, "scope");
Expand Down Expand Up @@ -1662,9 +1666,7 @@ pub async fn do_knowledge_delete(out: &output::Out, config: Config, args: &ArgMa
let catalog = args
.get_one::<String>("catalog_name")
.expect("catalog_name is required");
let name = args
.get_one::<String>("name")
.expect("name is required");
let name = args.get_one::<String>("name").expect("name is required");
let env = cmd::get_string_flag(args, "environment");

out.with_spinner(
Expand Down Expand Up @@ -2127,7 +2129,10 @@ mod tests {
])
.expect("query --max-rows should parse");
let (_, query_args) = matches.subcommand().expect("expected query subcommand");
assert_eq!(query_args.get_one::<usize>("max_rows").copied(), Some(50_000));
assert_eq!(
query_args.get_one::<usize>("max_rows").copied(),
Some(50_000)
);

assert!(
catalogs_cmd()
Expand Down Expand Up @@ -2296,15 +2301,14 @@ mod tests {
.expect("knowledge delete should parse");

let (_, knowledge_args) = matches.subcommand().expect("expected knowledge subcommand");
let (_, delete_args) = knowledge_args.subcommand().expect("expected delete subcommand");
let (_, delete_args) = knowledge_args
.subcommand()
.expect("expected delete subcommand");
assert_eq!(
delete_args.get_one::<String>("catalog_name").unwrap(),
"my-catalog"
);
assert_eq!(
delete_args.get_one::<String>("name").unwrap(),
"my-entry"
);
assert_eq!(delete_args.get_one::<String>("name").unwrap(), "my-entry");
}

#[test]
Expand Down Expand Up @@ -2363,7 +2367,9 @@ mod tests {
.expect("knowledge list should parse");

let (_, knowledge_args) = matches.subcommand().expect("expected knowledge subcommand");
let (_, list_args) = knowledge_args.subcommand().expect("expected list subcommand");
let (_, list_args) = knowledge_args
.subcommand()
.expect("expected list subcommand");

assert_eq!(
list_args.get_one::<String>("catalog_name").unwrap(),
Expand Down Expand Up @@ -2451,7 +2457,9 @@ mod tests {
.expect("knowledge set with all flags should parse");

let (_, knowledge_args) = matches.subcommand().expect("expected knowledge subcommand");
let (_, set_args) = knowledge_args.subcommand().expect("expected set subcommand");
let (_, set_args) = knowledge_args
.subcommand()
.expect("expected set subcommand");

assert_eq!(set_args.get_one::<String>("scope").unwrap(), "column");
assert_eq!(
Expand Down Expand Up @@ -2484,7 +2492,9 @@ mod tests {
.expect("knowledge set should parse");

let (_, knowledge_args) = matches.subcommand().expect("expected knowledge subcommand");
let (_, set_args) = knowledge_args.subcommand().expect("expected set subcommand");
let (_, set_args) = knowledge_args
.subcommand()
.expect("expected set subcommand");

assert_eq!(set_args.get_one::<String>("scope").unwrap(), "catalog");
assert_eq!(
Expand All @@ -2498,8 +2508,13 @@ mod tests {

#[test]
fn knowledge_set_requires_statement() {
let result =
catalogs_cmd().try_get_matches_from(["catalogs", "knowledge", "set", "my-catalog", "f"]);
let result = catalogs_cmd().try_get_matches_from([
"catalogs",
"knowledge",
"set",
"my-catalog",
"f",
]);
assert!(result.is_err());
}

Expand All @@ -2513,7 +2528,9 @@ mod tests {
.try_get_matches_from(["catalogs", "knowledge", "show", "my-catalog", "my-entry"])
.expect("knowledge show should parse");
let (_, knowledge_args) = matches.subcommand().expect("expected knowledge subcommand");
let (_, show_args) = knowledge_args.subcommand().expect("expected show subcommand");
let (_, show_args) = knowledge_args
.subcommand()
.expect("expected show subcommand");
assert_eq!(show_args.get_one::<String>("name").unwrap(), "my-entry");
assert_eq!(
show_args.get_one::<String>("environment").unwrap(),
Expand All @@ -2529,5 +2546,4 @@ mod tests {
assert_eq!(truncated.chars().count(), 80);
assert!(truncated.ends_with('…'));
}

}
36 changes: 19 additions & 17 deletions crates/tower-cmd/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,24 +154,26 @@ impl App {
Some(("query", args)) => {
catalogs::do_query(&out, sessionized_config, args).await
}
Some(("knowledge", knowledge_matches)) => match knowledge_matches.subcommand() {
Some(("list", args)) => {
catalogs::do_knowledge_list(&out, sessionized_config, args).await
Some(("knowledge", knowledge_matches)) => {
match knowledge_matches.subcommand() {
Some(("list", args)) => {
catalogs::do_knowledge_list(&out, sessionized_config, args).await
}
Some(("show", args)) => {
catalogs::do_knowledge_show(&out, sessionized_config, args).await
}
Some(("set", args)) => {
catalogs::do_knowledge_set(&out, sessionized_config, args).await
}
Some(("delete", args)) => {
catalogs::do_knowledge_delete(&out, sessionized_config, args).await
}
_ => {
catalogs::catalogs_cmd().print_help().unwrap();
std::process::exit(2);
}
}
Some(("show", args)) => {
catalogs::do_knowledge_show(&out, sessionized_config, args).await
}
Some(("set", args)) => {
catalogs::do_knowledge_set(&out, sessionized_config, args).await
}
Some(("delete", args)) => {
catalogs::do_knowledge_delete(&out, sessionized_config, args).await
}
_ => {
catalogs::catalogs_cmd().print_help().unwrap();
std::process::exit(2);
}
},
}
_ => {
catalogs::catalogs_cmd().print_help().unwrap();
std::process::exit(2);
Expand Down
Loading