@@ -4,26 +4,15 @@ use crate::{
 };
 use tonic::{Extensions, Request, transport::Channel};
 use tucana::{
-    aquila::{
-        DataTypeUpdateRequest, FlowTypeUpdateRequest, FunctionDefinitionUpdateRequest,
-        RuntimeFunctionDefinitionUpdateRequest, data_type_service_client::DataTypeServiceClient,
-        flow_type_service_client::FlowTypeServiceClient,
-        function_definition_service_client::FunctionDefinitionServiceClient,
-        runtime_function_definition_service_client::RuntimeFunctionDefinitionServiceClient,
-    },
-    shared::{
-        DefinitionDataType as DataType, FlowType, FunctionDefinition, RuntimeFunctionDefinition,
-    },
+    aquila::{ModuleUpdateRequest, module_service_client::ModuleServiceClient},
+    shared::Module,
 };
 
 pub mod auth;
 pub mod retry;
 
 pub struct FlowUpdateService {
-    data_types: Vec<DataType>,
-    runtime_functions: Vec<RuntimeFunctionDefinition>,
-    functions: Vec<FunctionDefinition>,
-    flow_types: Vec<FlowType>,
+    modules: Vec<Module>,
     channel: Channel,
     aquila_token: String,
     definition_source: Option<String>,
@@ -34,35 +23,19 @@ impl FlowUpdateService {
     ///
     /// This will read the definition files from the given path and initialize the service with the data types, runtime function definitions, function definitions, and flow types.
     pub async fn from_url(aquila_url: String, definition_path: &str, aquila_token: String) -> Self {
-        let mut data_types = Vec::new();
-        let mut runtime_functions = Vec::new();
-        let mut functions = Vec::new();
-        let mut flow_types = Vec::new();
-
         let reader = Reader::configure(definition_path.to_string(), true, vec![], None);
-
-        let features = match reader.read_features() {
-            Ok(features) => features,
+        let modules = match reader.read_modules() {
+            Ok(modules) => modules,
             Err(error) => {
                 log::error!("Error occurred while reading definitions: {:?}", error);
                 panic!("Error occurred while reading definitions")
             }
         };
 
-        for feature in features {
-            data_types.append(&mut feature.data_types.clone());
-            flow_types.append(&mut feature.flow_types.clone());
-            runtime_functions.append(&mut feature.runtime_functions.clone());
-            functions.append(&mut feature.functions.clone());
-        }
-
         let channel = create_channel_with_retry("Aquila", aquila_url).await;
 
         Self {
-            data_types,
-            runtime_functions,
-            functions,
-            flow_types,
+            modules,
             channel,
             aquila_token,
             definition_source: None,
@@ -74,194 +47,71 @@ impl FlowUpdateService {
         self
     }
 
-    pub fn with_flow_types(mut self, flow_types: Vec<FlowType>) -> Self {
-        self.flow_types = flow_types;
-        self
-    }
-
-    pub fn with_data_types(mut self, data_types: Vec<DataType>) -> Self {
-        self.data_types = data_types;
-        self
-    }
-
-    pub fn with_runtime_functions(
-        mut self,
-        runtime_functions: Vec<RuntimeFunctionDefinition>,
-    ) -> Self {
-        self.runtime_functions = runtime_functions;
-        self
-    }
-
-    pub fn with_functions(mut self, functions: Vec<FunctionDefinition>) -> Self {
-        self.functions = functions;
-        self
-    }
-
     pub async fn send(&mut self) {
         let _ = self.send_with_status().await;
     }
 
     pub async fn send_with_status(&mut self) -> bool {
-        let data_types_success = self.update_data_types().await;
-        let runtime_functions_success = self.update_runtime_functions().await;
-        let functions_success = self.update_functions().await;
-        let flow_types_success = self.update_flow_types().await;
-        data_types_success && runtime_functions_success && functions_success && flow_types_success
+        self.update().await
     }
 
-    async fn update_data_types(&mut self) -> bool {
-        if self.data_types.is_empty() {
-            log::info!("No DataTypes present.");
+    async fn update(&mut self) -> bool {
+        if self.modules.is_empty() {
+            log::info!("No Modules are present, aboarting update.");
             return true;
         }
 
+        let mut modules = self.modules.clone();
         if let Some(source) = &self.definition_source {
-            for data_type in self.data_types.iter_mut() {
-                data_type.definition_source = source.to_string();
-            }
+            modules = modules
+                .into_iter()
+                .map(|module| apply_definition_source_to_module(module, source.clone()))
+                .collect::<Vec<_>>();
         }
 
-        log::info!("Updating {} DataTypes.", self.data_types.len());
-        let mut client = DataTypeServiceClient::new(self.channel.clone());
+        log::info!("Updating {} Modules.", self.modules.len());
+        let mut client = ModuleServiceClient::new(self.channel.clone());
         let request = Request::from_parts(
             get_authorization_metadata(&self.aquila_token),
             Extensions::new(),
-            DataTypeUpdateRequest {
-                data_types: self.data_types.clone(),
-            },
+            ModuleUpdateRequest { modules },
         );
 
         match client.update(request).await {
             Ok(response) => {
                 let res = response.into_inner();
-                log::info!(
-                    "Was the update of the DataTypes accepted by Sagittarius? {}",
-                    res.success
-                );
-
-                res.success
-            }
-            Err(err) => {
-                log::error!("Failed to update data types: {:?}", err);
-                false
-            }
-        }
-    }
-
-    async fn update_functions(&mut self) -> bool {
-        if self.functions.is_empty() {
-            log::info!("No FunctionDefinitions present.");
-            return true;
-        }
 
-        if let Some(source) = &self.definition_source {
-            for function in self.functions.iter_mut() {
-                function.definition_source = source.to_string();
-            }
-        };
-
-        log::info!("Updating {} FunctionDefinitions.", self.functions.len());
-        let mut client = FunctionDefinitionServiceClient::new(self.channel.clone());
-        let request = Request::from_parts(
-            get_authorization_metadata(&self.aquila_token),
-            Extensions::new(),
-            FunctionDefinitionUpdateRequest {
-                functions: self.functions.clone(),
-            },
-        );
+                match res.success {
+                    true => log::info!("Module definition update has been successful"),
+                    false => log::warn!("Module definition update has been unsuccessful"),
+                };
 
-        match client.update(request).await {
-            Ok(response) => {
-                let res = response.into_inner();
-                log::info!(
-                    "Was the update of the FunctionDefinitions accepted by Sagittarius? {}",
-                    res.success
-                );
                 res.success
             }
             Err(err) => {
-                log::error!("Failed to update function definitions: {:?}", err);
+                log::error!("Module definition update failed. Reason: {:?}", err);
                 false
             }
         }
     }
+}
 
-    async fn update_runtime_functions(&mut self) -> bool {
-        if self.runtime_functions.is_empty() {
-            log::info!("No RuntimeFunctionDefinitions present.");
-            return true;
-        }
-
-        if let Some(source) = &self.definition_source {
-            for runtime_function in self.runtime_functions.iter_mut() {
-                runtime_function.definition_source = source.to_string();
-            }
-        }
-
-        log::info!(
-            "Updating {} RuntimeFunctionDefinitions.",
-            self.runtime_functions.len()
-        );
-        let mut client = RuntimeFunctionDefinitionServiceClient::new(self.channel.clone());
-        let request = Request::from_parts(
-            get_authorization_metadata(&self.aquila_token),
-            Extensions::new(),
-            RuntimeFunctionDefinitionUpdateRequest {
-                runtime_functions: self.runtime_functions.clone(),
-            },
-        );
-
-        match client.update(request).await {
-            Ok(response) => {
-                let res = response.into_inner();
-                log::info!(
-                    "Was the update of the RuntimeFunctionDefinitions accepted by Sagittarius? {}",
-                    res.success
-                );
-                res.success
-            }
-            Err(err) => {
-                log::error!("Failed to update runtime function definitions: {:?}", err);
-                false
-            }
-        }
+fn apply_definition_source_to_module(mut module: Module, source: String) -> Module {
+    for data_type in &mut module.definition_data_types {
+        data_type.definition_source = source.clone();
     }
-
-    async fn update_flow_types(&mut self) -> bool {
-        if self.flow_types.is_empty() {
-            log::info!("No FlowTypes present.");
-            return true;
-        }
-
-        if let Some(source) = &self.definition_source {
-            for flow_type in self.flow_types.iter_mut() {
-                flow_type.definition_source = Some(source.to_string());
-            }
-        }
-
-        log::info!("Updating {} FlowTypes.", self.flow_types.len());
-        let mut client = FlowTypeServiceClient::new(self.channel.clone());
-        let request = Request::from_parts(
-            get_authorization_metadata(&self.aquila_token),
-            Extensions::new(),
-            FlowTypeUpdateRequest {
-                flow_types: self.flow_types.clone(),
-            },
-        );
-
-        match client.update(request).await {
-            Ok(response) => {
-                let res = response.into_inner();
-                log::info!(
-                    "Was the update of the FlowTypes accepted by Sagittarius? {}",
-                    res.success
-                );
-                res.success
-            }
-            Err(err) => {
-                log::error!("Failed to update flow types: {:?}", err);
-                false
-            }
-        }
+    for flow_type in &mut module.flow_types {
+        flow_type.definition_source = Some(source.clone());
+    }
+    for runtime_flow_type in &mut module.runtime_flow_types {
+        runtime_flow_type.definition_source = Some(source.clone());
+    }
+    for function in &mut module.function_definitions {
+        function.definition_source = source.clone();
     }
+    for runtime_function in &mut module.runtime_function_definitions {
+        runtime_function.definition_source = source.clone();
+    }
+
+    module
 }