diff --git a/src/meta_plugin/env.rs b/src/meta_plugin/env.rs index e69de29..084e41a 100644 --- a/src/meta_plugin/env.rs +++ b/src/meta_plugin/env.rs @@ -0,0 +1,139 @@ +use crate::meta_plugin::{MetaPlugin, MetaPluginType, process_metadata_outputs}; +use log::debug; +use std::collections::HashMap; + +#[derive(Debug, Clone)] +pub struct EnvMetaPlugin { + is_finalized: bool, + outputs: std::collections::HashMap, + env_vars: Vec<(String, String)>, +} + +impl EnvMetaPlugin { + pub fn new( + _options: Option>, + outputs: Option>, + ) -> Self { + // Collect environment variables starting with KEEP_META_ + let mut env_vars = Vec::new(); + let mut outputs_map = std::collections::HashMap::new(); + + for (key, value) in std::env::vars() { + if let Some(stripped_key) = key.strip_prefix("KEEP_META_") { + // Add to env_vars to process later + env_vars.push((stripped_key.to_string(), value)); + // Add to outputs with default mapping to the stripped name + outputs_map.insert( + stripped_key.to_string(), + serde_yaml::Value::String(stripped_key.to_string()) + ); + } + } + + // Override with provided outputs + if let Some(provided_outputs) = outputs { + for (key, value) in provided_outputs { + outputs_map.insert(key, value); + } + } + + EnvMetaPlugin { + is_finalized: false, + outputs: outputs_map, + env_vars, + } + } +} + +impl MetaPlugin for EnvMetaPlugin { + fn meta_type(&self) -> MetaPluginType { + MetaPluginType::Env + } + + fn is_finalized(&self) -> bool { + self.is_finalized + } + + fn set_finalized(&mut self, finalized: bool) { + self.is_finalized = finalized; + } + + fn initialize(&mut self) -> crate::meta_plugin::MetaPluginResponse { + // If already finalized, don't process again + if self.is_finalized { + return crate::meta_plugin::MetaPluginResponse { + metadata: Vec::new(), + is_finalized: true, + }; + } + + // Process all collected environment variables + let mut metadata = Vec::new(); + for (name, value) in &self.env_vars { + if let Some(meta_data) = process_metadata_outputs( + name, + serde_yaml::Value::String(value.clone()), + &self.outputs + ) { + metadata.push(meta_data); + } + } + + // Mark as finalized since this plugin only needs to run once + self.is_finalized = true; + + crate::meta_plugin::MetaPluginResponse { + metadata, + is_finalized: true, + } + } + + fn update(&mut self, _data: &[u8]) -> crate::meta_plugin::MetaPluginResponse { + // If already finalized, don't process more data + if self.is_finalized { + return crate::meta_plugin::MetaPluginResponse { + metadata: Vec::new(), + is_finalized: true, + }; + } + + crate::meta_plugin::MetaPluginResponse { + metadata: Vec::new(), + is_finalized: false, + } + } + + fn finalize(&mut self) -> crate::meta_plugin::MetaPluginResponse { + // If not already finalized, we can call initialize + if !self.is_finalized { + return self.initialize(); + } + + crate::meta_plugin::MetaPluginResponse { + metadata: Vec::new(), + is_finalized: true, + } + } + + fn outputs(&self) -> &std::collections::HashMap { + &self.outputs + } + + fn outputs_mut(&mut self) -> &mut std::collections::HashMap { + &mut self.outputs + } + + fn default_outputs(&self) -> Vec { + self.env_vars.iter() + .map(|(name, _)| name.clone()) + .collect() + } + + fn options(&self) -> &std::collections::HashMap { + &std::collections::HashMap::new() + } + + fn options_mut(&mut self) -> &mut std::collections::HashMap { + panic!("options_mut() not implemented for EnvMetaPlugin") + } +} diff --git a/src/meta_plugin/mod.rs b/src/meta_plugin/mod.rs index f04b1bc..50a7dcd 100644 --- a/src/meta_plugin/mod.rs +++ b/src/meta_plugin/mod.rs @@ -13,6 +13,7 @@ pub mod user; pub mod shell; pub mod shell_pid; pub mod keep_pid; +pub mod env; use crate::meta_plugin::exec::MetaPluginExec; use crate::meta_plugin::digest::DigestMetaPlugin; @@ -26,6 +27,7 @@ use crate::meta_plugin::user::UserMetaPlugin; use crate::meta_plugin::shell::ShellMetaPlugin; use crate::meta_plugin::shell_pid::ShellPidMetaPlugin; use crate::meta_plugin::keep_pid::KeepPidMetaPlugin; +use crate::meta_plugin::env::EnvMetaPlugin; /// Represents metadata to be stored #[derive(Debug, Clone, Serialize, Deserialize)] @@ -131,6 +133,7 @@ pub enum MetaPluginType { ReadRate, Hostname, Exec, + Env, } /// Central function to handle metadata output with name mapping @@ -330,5 +333,8 @@ pub fn get_meta_plugin( options, outputs)) } + MetaPluginType::Env => { + Box::new(EnvMetaPlugin::new(options, outputs)) + } } }