New Expression type to handle VRL in one place

Author: ardatan

Cargo.lock

@@ -6,7 +6,9 @@ use std::{
 
 use bytes::{BufMut, Bytes, BytesMut};
 use dashmap::DashMap;
-use hive_router_config::{override_subgraph_urls::UrlOrExpression, HiveRouterConfig};
+use hive_router_config::{
+    override_subgraph_urls::UrlOrExpression, primitives::expression::Expression, HiveRouterConfig,
+};
 use http::Uri;
 use hyper_tls::HttpsConnector;
 use hyper_util::{
@@ -15,16 +17,7 @@ use hyper_util::{
 };
 use tokio::sync::{OnceCell, Semaphore};
 use tracing::error;
-use vrl::{
-    compiler::compile as vrl_compile,
-    compiler::Program as VrlProgram,
-    compiler::TargetValue as VrlTargetValue,
-    core::Value as VrlValue,
-    prelude::Function as VrlFunction,
-    prelude::{state::RuntimeState as VrlState, Context as VrlContext, TimeZone as VrlTimeZone},
-    stdlib::all as vrl_build_functions,
-    value::Secrets as VrlSecrets,
-};
+use vrl::core::Value as VrlValue;
 
 use crate::{
     execution::client_request_details::ClientRequestDetails,
@@ -44,7 +37,7 @@ type SubgraphEndpoint = String;
 type ExecutorsBySubgraphMap =
     DashMap<SubgraphName, DashMap<SubgraphEndpoint, SubgraphExecutorBoxedArc>>;
 type EndpointsBySubgraphMap = DashMap<SubgraphName, SubgraphEndpoint>;
-type ExpressionsBySubgraphMap = HashMap<SubgraphName, VrlProgram>;
+type ExpressionsBySubgraphMap = HashMap<SubgraphName, Expression>;
 
 pub struct SubgraphExecutorMap {
     executors_by_subgraph: ExecutorsBySubgraphMap,
@@ -54,8 +47,6 @@ pub struct SubgraphExecutorMap {
     /// Mapping from subgraph name to VRL expression program
     expressions_by_subgraph: ExpressionsBySubgraphMap,
     config: Arc<HiveRouterConfig>,
-    /// Precompiled VRL functions to be used in endpoint expressions.
-    vrl_functions: Vec<Box<dyn VrlFunction>>,
     client: Arc<HttpClient>,
     semaphores_by_origin: DashMap<String, Arc<Semaphore>>,
     max_connections_per_host: usize,
@@ -80,7 +71,6 @@ impl SubgraphExecutorMap {
             static_endpoints_by_subgraph: Default::default(),
             expressions_by_subgraph: Default::default(),
             config,
-            vrl_functions: vrl_build_functions(),
             client: Arc::new(client),
             semaphores_by_origin: Default::default(),
             max_connections_per_host,
@@ -101,7 +91,7 @@ impl SubgraphExecutorMap {
 
             let endpoint_str = match endpoint_str {
                 Some(UrlOrExpression::Url(url)) => url,
-                Some(UrlOrExpression::Expression { expression }) => {
+                Some(UrlOrExpression::Expression(expression)) => {
                     subgraph_executor_map.register_expression(&subgraph_name, expression)?;
                     &original_endpoint_str
                 }
@@ -194,21 +184,13 @@ impl SubgraphExecutorMap {
                         SubgraphExecutorError::StaticEndpointNotFound(subgraph_name.to_string())
                     })?,
             ));
-            let mut target = VrlTargetValue {
-                value: VrlValue::Object(BTreeMap::from([
-                    ("request".into(), client_request.into()),
-                    ("original_url".into(), original_url_value),
-                ])),
-                metadata: VrlValue::Object(BTreeMap::new()),
-                secrets: VrlSecrets::default(),
-            };
-
-            let mut state = VrlState::default();
-            let timezone = VrlTimeZone::default();
-            let mut ctx = VrlContext::new(&mut target, &mut state, &timezone);
+            let value = VrlValue::Object(BTreeMap::from([
+                ("request".into(), client_request.into()),
+                ("original_url".into(), original_url_value),
+            ]));
 
             // Resolve the expression to get an endpoint URL.
-            let endpoint_result = expression.resolve(&mut ctx).map_err(|err| {
+            let endpoint_result = expression.execute(value).map_err(|err| {
                 SubgraphExecutorError::new_endpoint_expression_resolution_failure(
                     subgraph_name.to_string(),
                     err,
@@ -267,14 +249,10 @@ impl SubgraphExecutorMap {
     fn register_expression(
         &mut self,
         subgraph_name: &str,
-        expression: &str,
+        expression: &Expression,
     ) -> Result<(), SubgraphExecutorError> {
-        let compilation_result = vrl_compile(expression, &self.vrl_functions).map_err(|e| {
-            SubgraphExecutorError::new_endpoint_expression_build(subgraph_name.to_string(), e)
-        })?;
-
         self.expressions_by_subgraph
-            .insert(subgraph_name.to_string(), compilation_result.program);
+            .insert(subgraph_name.to_string(), expression.clone());
 
         Ok(())
     }

lib/executor/src/executors/map.rs

@@ -12,43 +12,13 @@ use crate::headers::{
 use hive_router_config::headers as config;
 use http::HeaderName;
 use regex_automata::{meta, util::syntax::Config as SyntaxConfig};
-use vrl::{
-    compiler::compile as vrl_compile, prelude::Function as VrlFunction,
-    stdlib::all as vrl_build_functions,
-};
-
-pub struct HeaderRuleCompilerContext {
-    vrl_functions: Vec<Box<dyn VrlFunction>>,
-}
-
-impl Default for HeaderRuleCompilerContext {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl HeaderRuleCompilerContext {
-    pub fn new() -> Self {
-        Self {
-            vrl_functions: vrl_build_functions(),
-        }
-    }
-}
 
 pub trait HeaderRuleCompiler<A> {
-    fn compile(
-        &self,
-        ctx: &HeaderRuleCompilerContext,
-        actions: &mut A,
-    ) -> Result<(), HeaderRuleCompileError>;
+    fn compile(&self, actions: &mut A) -> Result<(), HeaderRuleCompileError>;
 }
 
 impl HeaderRuleCompiler<Vec<RequestHeaderRule>> for config::RequestHeaderRule {
-    fn compile(
-        &self,
-        ctx: &HeaderRuleCompilerContext,
-        actions: &mut Vec<RequestHeaderRule>,
-    ) -> Result<(), HeaderRuleCompileError> {
+    fn compile(&self, actions: &mut Vec<RequestHeaderRule>) -> Result<(), HeaderRuleCompileError> {
         match self {
             config::RequestHeaderRule::Propagate(rule) => {
                 let spec = materialize_match_spec(
@@ -78,16 +48,11 @@ impl HeaderRuleCompiler<Vec<RequestHeaderRule>> for config::RequestHeaderRule {
                         value: build_header_value(&rule.name, value)?,
                     }));
                 }
-                config::InsertSource::Expression { expression } => {
-                    let compilation_result =
-                        vrl_compile(expression, &ctx.vrl_functions).map_err(|e| {
-                            HeaderRuleCompileError::new_expression_build(rule.name.clone(), e)
-                        })?;
-
+                config::InsertSource::Expression(expression) => {
                     actions.push(RequestHeaderRule::InsertExpression(
                         RequestInsertExpression {
                             name: build_header_name(&rule.name)?,
-                            expression: Box::new(compilation_result.program),
+                            expression: Box::new(expression.clone()),
                         },
                     ));
                 }
@@ -112,11 +77,7 @@ impl HeaderRuleCompiler<Vec<RequestHeaderRule>> for config::RequestHeaderRule {
 }
 
 impl HeaderRuleCompiler<Vec<ResponseHeaderRule>> for config::ResponseHeaderRule {
-    fn compile(
-        &self,
-        ctx: &HeaderRuleCompilerContext,
-        actions: &mut Vec<ResponseHeaderRule>,
-    ) -> Result<(), HeaderRuleCompileError> {
+    fn compile(&self, actions: &mut Vec<ResponseHeaderRule>) -> Result<(), HeaderRuleCompileError> {
         match self {
             config::ResponseHeaderRule::Propagate(rule) => {
                 let aggregation_strategy = rule.algorithm.into();
@@ -153,21 +114,16 @@ impl HeaderRuleCompiler<Vec<ResponseHeaderRule>> for config::ResponseHeaderRule
                             strategy: aggregation_strategy,
                         }));
                     }
-                    config::InsertSource::Expression { expression } => {
+                    config::InsertSource::Expression(expression) => {
                         // NOTE: In case we ever need to improve performance and not pass the whole context
                         // to VRL expressions, we can use:
                         // - compilation_result.program.info().target_assignments
                         // - compilation_result.program.info().target_queries
                         // to determine what parts of the context are actually needed by the expression
-                        let compilation_result = vrl_compile(expression, &ctx.vrl_functions)
-                            .map_err(|e| {
-                                HeaderRuleCompileError::new_expression_build(rule.name.clone(), e)
-                            })?;
-
                         actions.push(ResponseHeaderRule::InsertExpression(
                             ResponseInsertExpression {
                                 name: build_header_name(&rule.name)?,
-                                expression: Box::new(compilation_result.program),
+                                expression: Box::new(expression.clone()),
                                 strategy: aggregation_strategy,
                             },
                         ));
@@ -196,19 +152,18 @@ impl HeaderRuleCompiler<Vec<ResponseHeaderRule>> for config::ResponseHeaderRule
 pub fn compile_headers_plan(
     cfg: &config::HeadersConfig,
 ) -> Result<HeaderRulesPlan, HeaderRuleCompileError> {
-    let ctx = HeaderRuleCompilerContext::new();
     let mut request_plan = RequestHeaderRules::default();
     let mut response_plan = ResponseHeaderRules::default();
 
     if let Some(global_rules) = &cfg.all {
-        request_plan.global = compile_request_header_rules(&ctx, global_rules)?;
-        response_plan.global = compile_response_header_rules(&ctx, global_rules)?;
+        request_plan.global = compile_request_header_rules(global_rules)?;
+        response_plan.global = compile_response_header_rules(global_rules)?;
     }
 
     if let Some(subgraph_rules_map) = &cfg.subgraphs {
         for (subgraph_name, subgraph_rules) in subgraph_rules_map {
-            let request_actions = compile_request_header_rules(&ctx, subgraph_rules)?;
-            let response_actions = compile_response_header_rules(&ctx, subgraph_rules)?;
+            let request_actions = compile_request_header_rules(subgraph_rules)?;
+            let response_actions = compile_response_header_rules(subgraph_rules)?;
             request_plan
                 .by_subgraph
                 .insert(subgraph_name.clone(), request_actions);
@@ -225,26 +180,24 @@ pub fn compile_headers_plan(
 }
 
 fn compile_request_header_rules(
-    ctx: &HeaderRuleCompilerContext,
     header_rules: &config::HeaderRules,
 ) -> Result<Vec<RequestHeaderRule>, HeaderRuleCompileError> {
     let mut request_actions = Vec::new();
     if let Some(request_rule_entries) = &header_rules.request {
         for request_rule in request_rule_entries {
-            request_rule.compile(ctx, &mut request_actions)?;
+            request_rule.compile(&mut request_actions)?;
         }
     }
     Ok(request_actions)
 }
 
 fn compile_response_header_rules(
-    ctx: &HeaderRuleCompilerContext,
     header_rules: &config::HeaderRules,
 ) -> Result<Vec<ResponseHeaderRule>, HeaderRuleCompileError> {
     let mut response_actions = Vec::new();
     if let Some(response_rule_entries) = &header_rules.response {
         for response_rule in response_rule_entries {
-            response_rule.compile(ctx, &mut response_actions)?;
+            response_rule.compile(&mut response_actions)?;
         }
     }
     Ok(response_actions)
@@ -358,7 +311,7 @@ mod tests {
     use http::HeaderName;
 
     use crate::headers::{
-        compile::{build_header_value, HeaderRuleCompiler, HeaderRuleCompilerContext},
+        compile::{build_header_value, HeaderRuleCompiler},
         errors::HeaderRuleCompileError,
         plan::{HeaderAggregationStrategy, RequestHeaderRule, ResponseHeaderRule},
     };
@@ -378,9 +331,8 @@ mod tests {
             rename: None,
             default: None,
         });
-        let ctx = HeaderRuleCompilerContext::new();
         let mut actions = Vec::new();
-        rule.compile(&ctx, &mut actions).unwrap();
+        rule.compile(&mut actions).unwrap();
         assert_eq!(actions.len(), 1);
         match &actions[0] {
             RequestHeaderRule::PropagateNamed(data) => {
@@ -401,8 +353,7 @@ mod tests {
             },
         });
         let mut actions = Vec::new();
-        let ctx = HeaderRuleCompilerContext::new();
-        rule.compile(&ctx, &mut actions).unwrap();
+        rule.compile(&mut actions).unwrap();
         assert_eq!(actions.len(), 1);
         match &actions[0] {
             RequestHeaderRule::InsertStatic(data) => {
@@ -423,8 +374,7 @@ mod tests {
             },
         });
         let mut actions = Vec::new();
-        let ctx = HeaderRuleCompilerContext::new();
-        rule.compile(&ctx, &mut actions).unwrap();
+        rule.compile(&mut actions).unwrap();
         assert_eq!(actions.len(), 1);
         match &actions[0] {
             RequestHeaderRule::RemoveNamed(data) => {
@@ -449,8 +399,7 @@ mod tests {
             default: Some("def".to_string()),
         });
         let mut actions = Vec::new();
-        let ctx = HeaderRuleCompilerContext::new();
-        let err = rule.compile(&ctx, &mut actions).unwrap_err();
+        let err = rule.compile(&mut actions).unwrap_err();
         match err {
             HeaderRuleCompileError::InvalidDefault => {}
             _ => panic!("Expected InvalidDefault error"),
@@ -470,8 +419,7 @@ mod tests {
             algorithm: config::AggregationAlgo::First,
         });
         let mut actions = Vec::new();
-        let ctx = HeaderRuleCompilerContext::new();
-        rule.compile(&ctx, &mut actions).unwrap();
+        rule.compile(&mut actions).unwrap();
         assert_eq!(actions.len(), 1);
         match &actions[0] {
             ResponseHeaderRule::PropagateNamed(data) => {

lib/executor/src/headers/compile.rs

@@ -1,6 +1,6 @@
 use http::header::{InvalidHeaderName, InvalidHeaderValue};
 use regex_automata::meta::BuildError;
-use vrl::{diagnostic::DiagnosticList, prelude::ExpressionError};
+use vrl::prelude::ExpressionError;
 
 #[derive(thiserror::Error, Debug)]
 pub enum HeaderRuleCompileError {
@@ -27,16 +27,8 @@ pub enum HeaderRuleRuntimeError {
 }
 
 impl HeaderRuleCompileError {
-    pub fn new_expression_build(header_name: String, diagnostics: DiagnosticList) -> Self {
-        HeaderRuleCompileError::ExpressionBuild(
-            header_name,
-            diagnostics
-                .errors()
-                .into_iter()
-                .map(|d| d.code.to_string() + ": " + &d.message)
-                .collect::<Vec<_>>()
-                .join(", "),
-        )
+    pub fn new_expression_build(header_name: String, err: String) -> Self {
+        HeaderRuleCompileError::ExpressionBuild(header_name, err)
     }
 }
 

lib/executor/src/headers/errors.rs

@@ -1,7 +1,7 @@
 use ahash::HashMap;
+use hive_router_config::primitives::expression::Expression;
 use http::{HeaderName, HeaderValue};
 use regex_automata::meta::Regex;
-use vrl::compiler::Program as VrlProgram;
 
 #[derive(Clone)]
 pub struct HeaderRulesPlan {
@@ -62,13 +62,13 @@ pub struct ResponseInsertStatic {
 #[derive(Clone)]
 pub struct RequestInsertExpression {
     pub name: HeaderName,
-    pub expression: Box<VrlProgram>,
+    pub expression: Box<Expression>,
 }
 
 #[derive(Clone)]
 pub struct ResponseInsertExpression {
     pub name: HeaderName,
-    pub expression: Box<VrlProgram>,
+    pub expression: Box<Expression>,
     pub strategy: HeaderAggregationStrategy,
 }