Skip to content
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.

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

feat(exec source): add support for customizing command environment #18223

Merged
merged 5 commits into from
Aug 14, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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: 102 additions & 0 deletions src/sources/exec/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use std::{
collections::HashMap,
io::{Error, ErrorKind},
path::PathBuf,
process::ExitStatus,
Expand Down Expand Up @@ -61,6 +62,17 @@ pub struct ExecConfig {
#[configurable(metadata(docs::examples = "echo", docs::examples = "Hello World!"))]
pub command: Vec<String>,

/// Custom environment variables to set or update when running the command.
hhromic marked this conversation as resolved.
Show resolved Hide resolved
/// If a variable name already exists in the environment, its value is replaced.
#[serde(default)]
#[configurable(metadata(docs::additional_props_description = "An environment variable."))]
#[configurable(metadata(docs::examples = "environment_examples()"))]
pub environment: Option<HashMap<String, String>>,

/// Whether or not to clear the environment before setting custom environment variables.
#[serde(default = "default_clear_environment")]
pub clear_environment: bool,

/// The directory in which to run the command.
pub working_directory: Option<PathBuf>,

Expand Down Expand Up @@ -141,6 +153,8 @@ impl Default for ExecConfig {
}),
streaming: None,
command: vec!["echo".to_owned(), "Hello World!".to_owned()],
environment: None,
clear_environment: default_clear_environment(),
working_directory: None,
include_stderr: default_include_stderr(),
maximum_buffer_size_bytes: default_maximum_buffer_size(),
Expand Down Expand Up @@ -168,10 +182,25 @@ const fn default_respawn_on_exit() -> bool {
true
}

const fn default_clear_environment() -> bool {
false
}

const fn default_include_stderr() -> bool {
true
}

fn environment_examples() -> HashMap<String, String> {
HashMap::<_, _>::from_iter(
[
("LANG".to_owned(), "es_ES.UTF-8".to_owned()),
("TZ".to_owned(), "Etc/UTC".to_owned()),
("PATH".to_owned(), "/bin:/usr/bin:/usr/local/bin".to_owned()),
]
.into_iter(),
)
}

fn get_hostname() -> Option<String> {
crate::get_hostname().ok()
}
Expand Down Expand Up @@ -610,6 +639,16 @@ fn build_command(config: &ExecConfig) -> Command {

command.kill_on_drop(true);

// Clear environment variables if needed
if config.clear_environment {
command.env_clear();
}

// Configure environment variables if needed
if let Some(envs) = &config.environment {
command.envs(envs);
}

// Explicitly set the current dir if needed
if let Some(current_dir) = &config.working_directory {
command.current_dir(current_dir);
Expand Down Expand Up @@ -726,6 +765,7 @@ mod tests {
use super::*;
use crate::{event::LogEvent, test_util::trace_init};
use bytes::Bytes;
use std::ffi::OsStr;
use std::io::Cursor;
use vector_core::event::EventMetadata;
use vrl::value;
Expand Down Expand Up @@ -900,6 +940,8 @@ mod tests {
respawn_interval_secs: default_respawn_interval_secs(),
}),
command: vec!["./runner".to_owned(), "arg1".to_owned(), "arg2".to_owned()],
environment: None,
clear_environment: default_clear_environment(),
working_directory: Some(PathBuf::from("/tmp")),
include_stderr: default_include_stderr(),
maximum_buffer_size_bytes: default_maximum_buffer_size(),
Expand All @@ -922,6 +964,64 @@ mod tests {
assert_eq!(expected_command_string, command_string);
}

#[test]
fn test_build_command_custom_environment() {
let config = ExecConfig {
mode: Mode::Streaming,
scheduled: None,
streaming: Some(StreamingConfig {
respawn_on_exit: default_respawn_on_exit(),
respawn_interval_secs: default_respawn_interval_secs(),
}),
command: vec!["./runner".to_owned(), "arg1".to_owned(), "arg2".to_owned()],
environment: Some(HashMap::from([("FOO".to_owned(), "foo".to_owned())])),
clear_environment: default_clear_environment(),
working_directory: Some(PathBuf::from("/tmp")),
include_stderr: default_include_stderr(),
maximum_buffer_size_bytes: default_maximum_buffer_size(),
framing: None,
decoding: default_decoding(),
log_namespace: None,
};

let command = build_command(&config);
let cmd = command.as_std();

let idx = cmd
.get_envs()
.position(|v| v == (OsStr::new("FOO"), Some(OsStr::new("foo"))));

assert_ne!(idx, None);
}

#[test]
fn test_build_command_clear_environment() {
let config = ExecConfig {
mode: Mode::Streaming,
scheduled: None,
streaming: Some(StreamingConfig {
respawn_on_exit: default_respawn_on_exit(),
respawn_interval_secs: default_respawn_interval_secs(),
}),
command: vec!["./runner".to_owned(), "arg1".to_owned(), "arg2".to_owned()],
environment: Some(HashMap::from([("FOO".to_owned(), "foo".to_owned())])),
clear_environment: true,
working_directory: Some(PathBuf::from("/tmp")),
include_stderr: default_include_stderr(),
maximum_buffer_size_bytes: default_maximum_buffer_size(),
framing: None,
decoding: default_decoding(),
log_namespace: None,
};

let command = build_command(&config);
let cmd = command.as_std();

let envs: Vec<_> = cmd.get_envs().collect();

assert_eq!(envs.len(), 1);
}

#[tokio::test]
async fn test_spawn_reader_thread() {
trace_init();
Expand Down Expand Up @@ -1112,6 +1212,8 @@ mod tests {
respawn_interval_secs: default_respawn_interval_secs(),
}),
command: vec!["yes".to_owned()],
environment: None,
clear_environment: default_clear_environment(),
working_directory: None,
include_stderr: default_include_stderr(),
maximum_buffer_size_bytes: default_maximum_buffer_size(),
Expand Down
24 changes: 24 additions & 0 deletions website/cue/reference/components/sources/base/exec.cue
Original file line number Diff line number Diff line change
@@ -1,6 +1,11 @@
package metadata

base: components: sources: exec: configuration: {
clear_environment: {
description: "Whether or not to clear the environment before setting custom environment variables."
required: false
type: bool: default: false
}
command: {
description: "The command to run, plus any arguments required."
required: true
Expand Down Expand Up @@ -143,6 +148,25 @@ base: components: sources: exec: configuration: {
}
}
}
environment: {
description: """
Custom environment variables to set or update when running the command.
If a variable name already exists in the environment, its value is replaced.
"""
required: false
type: object: {
examples: [{
LANG: "es_ES.UTF-8"
PATH: "/bin:/usr/bin:/usr/local/bin"
TZ: "Etc/UTC"
}]
options: "*": {
description: "An environment variable."
required: true
type: string: {}
}
}
}
framing: {
description: """
Framing configuration.
Expand Down