Simplify

Author: ardatan

lib/executor/src/executors/map.rs

@@ -108,21 +108,14 @@ impl SubgraphExecutorMap {
         client_request: &ClientRequestDetails<'a, 'req>,
     ) -> HttpExecutionResponse {
         match self.get_or_create_executor(subgraph_name, client_request) {
-            Ok(Some(executor)) => executor.execute(execution_request).await,
+            Ok(executor) => executor.execute(execution_request).await,
             Err(err) => {
                 error!(
                     "Subgraph executor error for subgraph '{}': {}",
                     subgraph_name, err,
                 );
                 self.internal_server_error_response(err.into(), subgraph_name)
             }
-            Ok(None) => {
-                error!(
-                    "Subgraph executor not found for subgraph '{}'",
-                    subgraph_name
-                );
-                self.internal_server_error_response("Internal server error".into(), subgraph_name)
-            }
         }
     }
 
@@ -151,15 +144,17 @@ impl SubgraphExecutorMap {
         &self,
         subgraph_name: &str,
         client_request: &ClientRequestDetails<'_, '_>,
-    ) -> Result<Option<SubgraphExecutorBoxedArc>, SubgraphExecutorError> {
-        let from_expression =
-            self.get_or_create_executor_from_expression(subgraph_name, client_request)?;
-
-        if from_expression.is_some() {
-            return Ok(from_expression);
-        }
-
-        Ok(self.get_executor_from_static_endpoint(subgraph_name))
+    ) -> Result<SubgraphExecutorBoxedArc, SubgraphExecutorError> {
+        self.expressions_by_subgraph
+            .get(subgraph_name)
+            .map(|expression| {
+                self.get_or_create_executor_from_expression(
+                    subgraph_name,
+                    expression,
+                    client_request,
+                )
+            })
+            .unwrap_or_else(|| self.get_executor_from_static_endpoint(subgraph_name))
     }
 
     /// Looks up a subgraph executor,
@@ -169,65 +164,50 @@ impl SubgraphExecutorMap {
     fn get_or_create_executor_from_expression(
         &self,
         subgraph_name: &str,
+        expression: &VrlProgram,
         client_request: &ClientRequestDetails<'_, '_>,
-    ) -> Result<Option<SubgraphExecutorBoxedArc>, SubgraphExecutorError> {
-        if let Some(expression) = self.expressions_by_subgraph.get(subgraph_name) {
-            let original_url_value = VrlValue::Bytes(Bytes::from(
-                self.static_endpoints_by_subgraph
-                    .get(subgraph_name)
-                    .map(|endpoint| endpoint.value().clone())
-                    .ok_or_else(|| {
-                        SubgraphExecutorError::StaticEndpointNotFound(subgraph_name.to_string())
-                    })?,
-            ));
-            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 =
-                execute_expression_with_value(expression, value).map_err(|err| {
-                    SubgraphExecutorError::new_endpoint_expression_resolution_failure(
-                        subgraph_name.to_string(),
-                        err,
-                    )
-                })?;
-            let endpoint_str = match endpoint_result.as_str() {
-                Some(s) => Ok(s),
-                None => Err(SubgraphExecutorError::EndpointExpressionWrongType(
-                    subgraph_name.to_string(),
-                )),
-            }?;
-
-            // Check if an executor for this endpoint already exists.
-            let existing_executor = self
-                .executors_by_subgraph
+    ) -> Result<SubgraphExecutorBoxedArc, SubgraphExecutorError> {
+        let original_url_value = VrlValue::Bytes(Bytes::from(
+            self.static_endpoints_by_subgraph
                 .get(subgraph_name)
-                .and_then(|endpoints| endpoints.get(endpoint_str.as_ref()).map(|e| e.clone()));
-
-            if let Some(executor) = existing_executor {
-                return Ok(Some(executor));
-            }
-
+                .map(|endpoint| endpoint.value().clone())
+                .ok_or_else(|| {
+                    SubgraphExecutorError::StaticEndpointNotFound(subgraph_name.to_string())
+                })?,
+        ));
+        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 = execute_expression_with_value(expression, value).map_err(|err| {
+            SubgraphExecutorError::new_endpoint_expression_resolution_failure(
+                subgraph_name.to_string(),
+                err,
+            )
+        })?;
+        let endpoint_str = match endpoint_result.as_str() {
+            Some(s) => Ok(s),
+            None => Err(SubgraphExecutorError::EndpointExpressionWrongType(
+                subgraph_name.to_string(),
+            )),
+        }?;
+
+        // Check if an executor for this endpoint already exists.
+        self.executors_by_subgraph
+            .get(subgraph_name)
+            .and_then(|endpoints| endpoints.get(endpoint_str.as_ref()).map(|e| e.clone()))
+            .map(Ok)
             // If not, create and register a new one.
-            self.register_executor(subgraph_name, &endpoint_str)?;
-
-            let endpoints = self
-                .executors_by_subgraph
-                .get(subgraph_name)
-                .expect("Executor was just registered, should be present");
-            return Ok(endpoints.get(endpoint_str.as_ref()).map(|e| e.clone()));
-        }
-
-        Ok(None)
+            .unwrap_or_else(|| self.register_executor(subgraph_name, endpoint_str.as_ref()))
     }
 
     /// Looks up a subgraph executor based on a static endpoint URL.
     fn get_executor_from_static_endpoint(
         &self,
         subgraph_name: &str,
-    ) -> Option<SubgraphExecutorBoxedArc> {
+    ) -> Result<SubgraphExecutorBoxedArc, SubgraphExecutorError> {
         self.static_endpoints_by_subgraph
             .get(subgraph_name)
             .and_then(|endpoint_ref| {
@@ -236,6 +216,7 @@ impl SubgraphExecutorMap {
                     .get(subgraph_name)
                     .and_then(|endpoints| endpoints.get(endpoint_str).map(|e| e.clone()))
             })
+            .ok_or_else(|| SubgraphExecutorError::StaticEndpointNotFound(subgraph_name.to_string()))
     }
 
     /// Registers a VRL expression for the given subgraph name.
@@ -269,7 +250,7 @@ impl SubgraphExecutorMap {
         &self,
         subgraph_name: &str,
         endpoint_str: &str,
-    ) -> Result<(), SubgraphExecutorError> {
+    ) -> Result<SubgraphExecutorBoxedArc, SubgraphExecutorError> {
         let endpoint_uri = endpoint_str.parse::<Uri>().map_err(|e| {
             SubgraphExecutorError::EndpointParseFailure(endpoint_str.to_string(), e.to_string())
         })?;
@@ -302,11 +283,13 @@ impl SubgraphExecutorMap {
             self.in_flight_requests.clone(),
         );
 
+        let executor_arc = executor.to_boxed_arc();
+
         self.executors_by_subgraph
             .entry(subgraph_name.to_string())
             .or_default()
-            .insert(endpoint_str.to_string(), executor.to_boxed_arc());
+            .insert(endpoint_str.to_string(), executor_arc.clone());
 
-        Ok(())
+        Ok(executor_arc)
     }
 }