feat: Add query API to Python plugins (#25766)

This ended up being a couple things rolled into one. In order to add a query API to the Python plugin, I had to pull the QueryExecutor trait out of server into a place so that the python crate could use it.

This implements the query API, but also fixes up the WAL plugin test CLI a bit. I've added a test in the CLI section so that it shows end-to-end operation of the WAL plugin test API and exercise of the entire Plugin API.

Closes #25757

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "influxdb3_clap_blocks",
     "influxdb3_client",
     "influxdb3_id",
+    "influxdb3_internal_api",
     "influxdb3_load_generator",
     "influxdb3_process",
     "influxdb3_py_api",

influxdb3/src/commands/test.rs

@@ -1,4 +1,5 @@
 use crate::commands::common::{InfluxDb3Config, SeparatedKeyValue, SeparatedList};
+use anyhow::Context;
 use influxdb3_client::plugin_development::WalPluginTestRequest;
 use influxdb3_client::Client;
 use secrecy::ExposeSecret;
@@ -53,34 +54,39 @@ pub struct WalPluginConfig {
     /// If given pass this map of string key/value pairs as input arguments
     #[clap(long = "input-arguments")]
     pub input_arguments: Option<SeparatedList<SeparatedKeyValue<String, String>>>,
-    /// The name of the plugin, which should match its file name on the server `<plugin-dir>/<name>.py`
+    /// The file name of the plugin, which should exist on the server in `<plugin-dir>/<filename>`.
+    /// The plugin-dir is provided on server startup.
     #[clap(required = true)]
-    pub name: String,
-}
-
-impl From<WalPluginConfig> for WalPluginTestRequest {
-    fn from(val: WalPluginConfig) -> Self {
-        let input_arguments = val.input_arguments.map(|a| {
-            a.into_iter()
-                .map(|SeparatedKeyValue((k, v))| (k, v))
-                .collect::<HashMap<String, String>>()
-        });
-
-        Self {
-            name: val.name,
-            input_lp: val.input_lp,
-            input_file: val.input_file,
-            input_arguments,
-        }
-    }
+    pub filename: String,
 }
 
 pub async fn command(config: Config) -> Result<(), Box<dyn Error>> {
     let client = config.get_client()?;
 
     match config.cmd {
         SubCommand::WalPlugin(plugin_config) => {
-            let wal_plugin_test_request: WalPluginTestRequest = plugin_config.into();
+            let input_arguments = plugin_config.input_arguments.map(|a| {
+                a.into_iter()
+                    .map(|SeparatedKeyValue((k, v))| (k, v))
+                    .collect::<HashMap<String, String>>()
+            });
+
+            let input_lp = match plugin_config.input_lp {
+                Some(lp) => lp,
+                None => {
+                    let file_path = plugin_config
+                        .input_file
+                        .context("either input_lp or input_file must be provided")?;
+                    std::fs::read_to_string(file_path).context("unable to read input file")?
+                }
+            };
+
+            let wal_plugin_test_request = WalPluginTestRequest {
+                filename: plugin_config.filename,
+                database: plugin_config.influxdb3_config.database_name,
+                input_lp,
+                input_arguments,
+            };
 
             let response = client.wal_plugin_test(wal_plugin_test_request).await?;
 

influxdb3/tests/server/cli.rs

@@ -849,3 +849,105 @@ async fn meta_cache_create_and_delete() {
 
     insta::assert_yaml_snapshot!(result);
 }
+
+#[cfg(feature = "system-py")]
+#[test_log::test(tokio::test)]
+async fn test_wal_plugin_test() {
+    use crate::ConfigProvider;
+    use influxdb3_client::Precision;
+
+    // Create plugin file
+    let plugin_file = create_plugin_file(
+        r#"
+def process_writes(influxdb3_local, table_batches, args=None):
+    influxdb3_local.info("arg1: " + args["arg1"])
+
+    query_params = {"host": args["host"]}
+    query_result = influxdb3_local.query_rows("SELECT * FROM cpu where host = $host", query_params)
+    influxdb3_local.info("query result: " + str(query_result))
+
+    for table_batch in table_batches:
+        influxdb3_local.info("table: " + table_batch["table_name"])
+
+        for row in table_batch["rows"]:
+            influxdb3_local.info("row: " + str(row))
+
+    line = LineBuilder("some_table")\
+        .tag("tag1", "tag1_value")\
+        .tag("tag2", "tag2_value")\
+        .int64_field("field1", 1)\
+        .float64_field("field2", 2.0)\
+        .string_field("field3", "number three")
+    influxdb3_local.write(line)
+
+    other_line = LineBuilder("other_table")
+    other_line.int64_field("other_field", 1)
+    other_line.float64_field("other_field2", 3.14)
+    other_line.time_ns(1302)
+
+    influxdb3_local.write_to_db("mytestdb", other_line)
+
+    influxdb3_local.info("done")"#,
+    );
+
+    let plugin_dir = plugin_file.path().parent().unwrap().to_str().unwrap();
+    let plugin_name = plugin_file.path().file_name().unwrap().to_str().unwrap();
+
+    let server = TestServer::configure()
+        .with_plugin_dir(plugin_dir)
+        .spawn()
+        .await;
+    let server_addr = server.client_addr();
+
+    server
+        .write_lp_to_db(
+            "foo",
+            "cpu,host=s1,region=us-east usage=0.9 1\n\
+            cpu,host=s2,region=us-east usage=0.89 2\n\
+            cpu,host=s1,region=us-east usage=0.85 3",
+            Precision::Nanosecond,
+        )
+        .await
+        .unwrap();
+
+    let db_name = "foo";
+
+    // Run the test
+    let result = run_with_confirmation(&[
+        "test",
+        "wal_plugin",
+        "--database",
+        db_name,
+        "--host",
+        &server_addr,
+        "--lp",
+        "test_input,tag1=tag1_value,tag2=tag2_value field1=1i 500",
+        "--input-arguments",
+        "arg1=arg1_value,host=s2",
+        plugin_name,
+    ]);
+    debug!(result = ?result, "test wal plugin");
+
+    let res = serde_json::from_str::<serde_json::Value>(&result).unwrap();
+
+    let expected_result = r#"{
+  "log_lines": [
+    "INFO: arg1: arg1_value",
+    "INFO: query result: [{'host': 's2', 'region': 'us-east', 'time': 2, 'usage': 0.89}]",
+    "INFO: table: test_input",
+    "INFO: row: {'tag1': 'tag1_value', 'tag2': 'tag2_value', 'field1': 1, 'time': 500}",
+    "INFO: done"
+  ],
+  "database_writes": {
+    "mytestdb": [
+      "other_table other_field=1i,other_field2=3.14 1302"
+    ],
+    "foo": [
+      "some_table,tag1=tag1_value,tag2=tag2_value field1=1i,field2=2.0,field3=\"number three\""
+    ]
+  },
+  "errors": []
+}"#;
+    let expected_result = serde_json::from_str::<serde_json::Value>(expected_result).unwrap();
+    assert_eq!(res, expected_result);
+}

influxdb3/tests/server/main.rs

@@ -48,6 +48,7 @@ trait ConfigProvider {
 pub struct TestConfig {
     auth_token: Option<(String, String)>,
     host_id: Option<String>,
+    plugin_dir: Option<String>,
 }
 
 impl TestConfig {
@@ -66,6 +67,12 @@ impl TestConfig {
         self.host_id = Some(host_id.into());
         self
     }
+
+    /// Set the plugin dir for this [`TestServer`]
+    pub fn with_plugin_dir<S: Into<String>>(mut self, plugin_dir: S) -> Self {
+        self.plugin_dir = Some(plugin_dir.into());
+        self
+    }
 }
 
 impl ConfigProvider for TestConfig {
@@ -74,6 +81,9 @@ impl ConfigProvider for TestConfig {
         if let Some((token, _)) = &self.auth_token {
             args.append(&mut vec!["--bearer-token".to_string(), token.to_owned()]);
         }
+        if let Some(plugin_dir) = &self.plugin_dir {
+            args.append(&mut vec!["--plugin-dir".to_string(), plugin_dir.to_owned()]);
+        }
         args.push("--host-id".to_string());
         if let Some(host) = &self.host_id {
             args.push(host.to_owned());

influxdb3_client/src/plugin_development.rs

@@ -6,9 +6,9 @@ use std::collections::HashMap;
 /// Request definition for `POST /api/v3/plugin_test/wal` API
 #[derive(Debug, Serialize, Deserialize)]
 pub struct WalPluginTestRequest {
-    pub name: String,
-    pub input_lp: Option<String>,
-    pub input_file: Option<String>,
+    pub filename: String,
+    pub database: String,
+    pub input_lp: String,
     pub input_arguments: Option<HashMap<String, String>>,
 }
 

influxdb3_internal_api/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "influxdb3_internal_api"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[dependencies]
+# Core Crates
+iox_query.workspace = true
+iox_query_params.workspace = true
+trace.workspace = true
+trace_http.workspace = true
+tracker.workspace = true
+
+# Local Crates
+
+# Crates.io dependencies
+async-trait.workspace = true
+datafusion.workspace = true
+thiserror.workspace = true
+
+[lints]
+workspace = true

influxdb3_internal_api/src/lib.rs

@@ -0,0 +1,4 @@
+//! This crate contains the internal API for use across the crates in this code base, mainly
+//! to get around circular dependency issues.
+
+pub mod query_executor;