omdb: Add one-shot query exec to oxql subcommand (#10154)

In some cases, we may execute an `oxql` query as part of an ad hoc shell
script, or have capture the queries executed in the shell history. The
mandatory `reedline` shell used by the `oxql` subcommand prevents this.

Add a `--execute` flag to provide a one-shot execution option to the
`oxql` subcommand.

Author: wfchandler

dev-tools/omdb/src/bin/omdb/oxql.rs

@@ -38,6 +38,10 @@ pub struct OxqlArgs {
     /// Print the total elapsed query duration.
     #[clap(long = "elapsed")]
     print_elapsed: bool,
+
+    /// Execute the query and exit.
+    #[clap(long, short, value_name = "QUERY")]
+    execute: Option<String>,
 }
 
 impl OxqlArgs {
@@ -54,13 +58,24 @@ impl OxqlArgs {
             print_elapsed: self.print_elapsed,
         };
 
-        oxql::shell(
-            native_addr.ip(),
-            native_addr.port(),
-            log.new(slog::o!("component" => "clickhouse-client")),
-            opts,
-        )
-        .await
+        if let Some(query) = &self.execute {
+            oxql::exec_query(
+                native_addr.ip(),
+                native_addr.port(),
+                log.new(slog::o!("component" => "clickhouse-client")),
+                opts,
+                query.to_owned(),
+            )
+            .await
+        } else {
+            oxql::shell(
+                native_addr.ip(),
+                native_addr.port(),
+                log.new(slog::o!("component" => "clickhouse-client")),
+                opts,
+            )
+            .await
+        }
     }
 
     /// Resolve the ClickHouse native TCP socket address.

dev-tools/omdb/tests/usage_errors.out

@@ -1205,6 +1205,7 @@ Options:
       --summaries              Print summaries of each SQL query run against the database
       --elapsed                Print the total elapsed query duration
       --color <COLOR>          Color output [default: auto] [possible values: auto, always, never]
+  -e, --execute <QUERY>        Execute the query and exit
   -h, --help                   Print help
 
 Connection Options:

oximeter/db/src/shells/oxql.rs

@@ -37,12 +37,7 @@ pub async fn shell(
     opts: ShellOptions,
 ) -> anyhow::Result<()> {
     // Create the client.
-    let client = make_client(address, port, &log).await?;
-
-    // A workaround to ensure the client has all available timeseries when the
-    // shell starts.
-    let dummy = "foo:bar".parse().unwrap();
-    let _ = client.schema_for_timeseries(&dummy).await;
+    let client = make_oxql_client(address, port, &log).await?;
 
     // Create the line-editor.
     let mut ed = Reedline::create();
@@ -157,6 +152,46 @@ pub async fn shell(
     }
 }
 
+/// Execute the provided OxQL query.
+pub async fn exec_query(
+    address: IpAddr,
+    port: u16,
+    log: Logger,
+    opts: ShellOptions,
+    statement: String,
+) -> anyhow::Result<()> {
+    // Create the client.
+    let client = make_oxql_client(address, port, &log).await?;
+
+    let result = client
+        .oxql_query(
+            statement.trim().trim_end_matches(';'),
+            QueryAuthzScope::Fleet,
+        )
+        .await?;
+
+    print_tables(&result.tables);
+    println!();
+    print_query_summary(&result, opts.print_elapsed, opts.print_summaries);
+
+    Ok(())
+}
+
+/// Create an OxQL client and prime its schema cache.
+async fn make_oxql_client(
+    address: IpAddr,
+    port: u16,
+    log: &Logger,
+) -> anyhow::Result<Client> {
+    let client = make_client(address, port, log).await?;
+
+    // Workaround to ensure the client has all available timeseries.
+    let dummy = "foo:bar".parse().unwrap();
+    let _ = client.schema_for_timeseries(&dummy).await;
+
+    Ok(client)
+}
+
 /// Describe a single timeseries.
 async fn describe_timeseries(
     client: &Client,