Merge pull request #165 from influxdata/crepererum/env-vars

feat: environment variables for guests

Author: crepererum

guests/python/README.md

@@ -197,6 +197,9 @@ def compute(x: int) -> int:
 ## I/O
 All I/O operations go through the host, there is no direct interaction with the host operating system.
 
+### Environment Variables
+Hosts can pass environment variables to the guest if they want. By default, NO variables are available for the guest though (i.e. there is NO implicit pass-through). The standard Python library can read these environment variables, e.g. via [`os.environ`].
+
 ### Filesystem
 The [Python Standard Library] is mounted as a read-only filesystem. The host file system (incl. special paths like `/proc`) are NOT exposed to the guest.
 
@@ -234,6 +237,7 @@ There is NO other I/O available that escapes the sandbox.
 [`timedelta`]: https://docs.python.org/3/library/datetime.html#datetime.timedelta
 [`Duration`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html#variant.Duration
 [`Microsecond`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.TimeUnit.html#variant.Microsecond
+[`os.environ`]: https://docs.python.org/3/library/os.html#os.environ
 [Python 3.14.0]: https://www.python.org/downloads/release/python-3140
 [Python Standard Library]: https://docs.python.org/3/library/index.html
 [`requests`]: https://pypi.org/project/requests/

host/src/lib.rs

@@ -2,7 +2,7 @@
 //!
 //!
 //! [DataFusion]: https://datafusion.apache.org/
-use std::{any::Any, ops::DerefMut, sync::Arc};
+use std::{any::Any, collections::BTreeMap, ops::DerefMut, sync::Arc};
 
 use ::http::HeaderName;
 use arrow::datatypes::DataType;
@@ -247,6 +247,9 @@ pub struct WasmPermissions {
 
     /// Virtual file system limits.
     vfs: VfsLimits,
+
+    /// Environment variables.
+    envs: BTreeMap<String, String>,
 }
 
 impl WasmPermissions {
@@ -261,6 +264,7 @@ impl Default for WasmPermissions {
         Self {
             http: Arc::new(RejectAllHttpRequests),
             vfs: VfsLimits::default(),
+            envs: BTreeMap::default(),
         }
     }
 }
@@ -284,6 +288,12 @@ impl WasmPermissions {
             ..self
         }
     }
+
+    /// Add environment variable.
+    pub fn with_env(mut self, key: String, value: String) -> Self {
+        self.envs.insert(key, value);
+        self
+    }
 }
 
 /// A [`ScalarUDFImpl`] that wraps a WebAssembly payload.
@@ -329,12 +339,18 @@ impl WasmScalarUdf {
         // Create in-memory VFS
         let vfs_state = VfsState::new(&permissions.vfs);
 
+        // set up WASI p2 context
         let stderr = MemoryOutputPipe::new(1024);
-        let wasi_ctx = WasiCtx::builder().stderr(stderr.clone()).build();
+        let mut wasi_ctx_builder = WasiCtx::builder();
+        wasi_ctx_builder.stderr(stderr.clone());
+        permissions.envs.iter().for_each(|(k, v)| {
+            wasi_ctx_builder.env(k, v);
+        });
+
         let state = WasmStateImpl {
             vfs_state,
             stderr,
-            wasi_ctx,
+            wasi_ctx: wasi_ctx_builder.build(),
             wasi_http_ctx: WasiHttpCtx::new(),
             resource_table: ResourceTable::new(),
             http_validator: Arc::clone(&permissions.http),

host/tests/integration_tests/python/runtime/env.rs

@@ -0,0 +1,71 @@
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field};
+use datafusion_common::{cast::as_string_array, config::ConfigOptions};
+use datafusion_expr::{ScalarFunctionArgs, async_udf::AsyncScalarUDFImpl};
+use datafusion_udf_wasm_host::{WasmPermissions, WasmScalarUdf};
+use tokio::runtime::Handle;
+
+use crate::integration_tests::python::test_utils::python_component;
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_env() {
+    assert_env_roundrip(&[]).await;
+    assert_env_roundrip(&[("FOO", "BAR")]).await;
+    assert_env_roundrip(&[("foo", "bar")]).await;
+    assert_env_roundrip(&[("", "")]).await;
+    assert_env_roundrip(&[("FOO", "BAR"), ("X", "Y")]).await;
+}
+
+pub(crate) async fn assert_env_roundrip(env: &[(&'static str, &'static str)]) {
+    const CODE: &str = r#"
+import os
+
+def env() -> str:
+    return ",".join((
+        f"{k}:{v}"
+        for k, v in os.environ.items()
+    ))
+"#;
+
+    let component = python_component().await;
+
+    let mut permissions = WasmPermissions::default();
+    for (k, v) in env {
+        permissions = permissions.with_env((*k).to_owned(), (*v).to_owned());
+    }
+
+    let udfs = WasmScalarUdf::new(component, &permissions, Handle::current(), CODE.to_owned())
+        .await
+        .unwrap();
+    assert_eq!(udfs.len(), 1);
+    let udf = udfs.into_iter().next().unwrap();
+
+    let array = udf
+        .invoke_async_with_args(
+            ScalarFunctionArgs {
+                args: vec![],
+                arg_fields: vec![],
+                number_rows: 1,
+                return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
+            },
+            &ConfigOptions::default(),
+        )
+        .await
+        .unwrap();
+
+    let actual = as_string_array(&array)
+        .unwrap()
+        .into_iter()
+        .next()
+        .unwrap()
+        .unwrap();
+
+    let expected = env
+        .iter()
+        .map(|(k, v)| format!("{k}:{v}"))
+        .collect::<Vec<_>>();
+    let expected = expected.join(",");
+
+    assert_eq!(actual, expected);
+}

host/tests/integration_tests/python/runtime/mod.rs

@@ -1,4 +1,5 @@
 mod dependencies;
+mod env;
 mod errors;
 mod fs;
 mod http;