refactor(lsp_server): improve structure and robustness

Author: tembleking

src/app/lsp_server/lsp_server_inner.rs

@@ -4,8 +4,8 @@ use tower_lsp::lsp_types::{
     CodeActionOrCommand, CodeActionParams, CodeActionProviderCapability, CodeActionResponse,
     CodeLens, CodeLensOptions, CodeLensParams, Command, DidChangeConfigurationParams,
     DidChangeTextDocumentParams, DidOpenTextDocumentParams, ExecuteCommandOptions,
-    ExecuteCommandParams, InitializeParams, InitializeResult, InitializedParams, MessageType,
-    ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind,
+    ExecuteCommandParams, InitializeParams, InitializeResult, InitializedParams, Location,
+    MessageType, ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind,
 };
 use tracing::{debug, info};
 
@@ -23,7 +23,7 @@ use super::supported_commands::SupportedCommands;
 pub struct LSPServerInner<C> {
     interactor: LspInteractor<C>,
     query_executor: QueryExecutor,
-    component_factory: ComponentFactory,
+    component_factory: Option<ComponentFactory>,
 }
 
 impl<C> LSPServerInner<C> {
@@ -33,7 +33,7 @@ impl<C> LSPServerInner<C> {
         LSPServerInner {
             interactor: LspInteractor::new(client, document_database.clone()),
             query_executor: QueryExecutor::new(document_database.clone()),
-            component_factory: Default::default(), // to be initialized in the initialize method of the LSP
+            component_factory: None, // to be initialized in the initialize method of the LSP
         }
     }
 }
@@ -50,7 +50,9 @@ where
 
         debug!("updating with configuration: {config:?}");
 
-        self.component_factory.initialize_with(config);
+        let mut factory = ComponentFactory::default();
+        factory.initialize_with(config);
+        self.component_factory = Some(factory);
 
         debug!("updated configuration");
         Ok(())
@@ -61,6 +63,20 @@ impl<C> LSPServerInner<C>
 where
     C: LSPClient + Send + Sync + 'static,
 {
+    async fn get_commands_for_document(
+        &self,
+        uri: &tower_lsp::lsp_types::Url,
+    ) -> Result<Vec<command_generator::CommandInfo>> {
+        let Some(content) = self.query_executor.get_document_text(uri.as_str()).await else {
+            return Err(Error::internal_error().with_message(format!(
+                "unable to extract document content for document: {uri}"
+            )));
+        };
+
+        let commands = command_generator::generate_commands_for_uri(uri, &content);
+        Ok(commands)
+    }
+
     pub async fn initialize(
         &mut self,
         initialize_params: InitializeParams,
@@ -128,19 +144,9 @@ where
         &self,
         params: CodeActionParams,
     ) -> Result<Option<CodeActionResponse>> {
-        let Some(content) = self
-            .query_executor
-            .get_document_text(params.text_document.uri.as_str())
-            .await
-        else {
-            return Err(Error::internal_error().with_message(format!(
-                "unable to extract document content for document: {}",
-                &params.text_document.uri
-            )));
-        };
-
-        let commands =
-            command_generator::generate_commands_for_uri(&params.text_document.uri, &content);
+        let commands = self
+            .get_commands_for_document(&params.text_document.uri)
+            .await?;
         let code_actions: Vec<CodeActionOrCommand> = commands
             .into_iter()
             .filter(|cmd| cmd.range.start.line == params.range.start.line)
@@ -157,19 +163,9 @@ where
     }
 
     pub async fn code_lens(&self, params: CodeLensParams) -> Result<Option<Vec<CodeLens>>> {
-        let Some(content) = self
-            .query_executor
-            .get_document_text(params.text_document.uri.as_str())
-            .await
-        else {
-            return Err(Error::internal_error().with_message(format!(
-                "unable to extract document content for document: {}",
-                &params.text_document.uri
-            )));
-        };
-
-        let commands =
-            command_generator::generate_commands_for_uri(&params.text_document.uri, &content);
+        let commands = self
+            .get_commands_for_document(&params.text_document.uri)
+            .await?;
         let code_lenses = commands
             .into_iter()
             .map(|cmd| CodeLens {
@@ -186,49 +182,63 @@ where
         Ok(Some(code_lenses))
     }
 
+    fn component_factory_mut(&mut self) -> Result<&mut ComponentFactory> {
+        self.component_factory
+            .as_mut()
+            .ok_or_else(|| Error::internal_error().with_message("LSP not initialized"))
+    }
+
+    async fn execute_base_image_scan(
+        &mut self,
+        location: Location,
+        image: String,
+    ) -> Result<Option<Value>> {
+        let image_scanner = self.component_factory_mut()?.image_scanner().map_err(|e| {
+            Error::internal_error().with_message(format!("unable to create image scanner: {e}"))
+        })?;
+        let mut command =
+            ScanBaseImageCommand::new(&image_scanner, &self.interactor, location, image);
+        command.execute().await.map(|_| None)
+    }
+
+    async fn execute_build_and_scan(&mut self, location: Location) -> Result<Option<Value>> {
+        let factory = self.component_factory_mut()?;
+        let image_scanner = factory.image_scanner().map_err(|e| {
+            Error::internal_error().with_message(format!("unable to create image scanner: {e}"))
+        })?;
+        let image_builder = factory.image_builder().map_err(|e| {
+            Error::internal_error().with_message(format!("unable to create image builder: {e}"))
+        })?;
+        let mut command =
+            BuildAndScanCommand::new(&image_builder, &image_scanner, &self.interactor, location);
+        command.execute().await.map(|_| None)
+    }
+
     pub async fn execute_command(&mut self, params: ExecuteCommandParams) -> Result<Option<Value>> {
         let command: SupportedCommands = params.try_into()?;
 
         let result = match command.clone() {
             SupportedCommands::ExecuteBaseImageScan { location, image } => {
-                let image_scanner = self.component_factory.image_scanner().map_err(|e| {
-                    Error::internal_error()
-                        .with_message(format!("unable to create image scanner: {e}"))
-                })?;
-                let mut command =
-                    ScanBaseImageCommand::new(&image_scanner, &self.interactor, location, image);
-                command.execute().await.map(|_| None)
+                self.execute_base_image_scan(location, image).await
             }
 
             SupportedCommands::ExecuteBuildAndScan { location } => {
-                let image_scanner = self.component_factory.image_scanner().map_err(|e| {
-                    Error::internal_error()
-                        .with_message(format!("unable to create image scanner: {e}"))
-                })?;
-                let image_builder = self.component_factory.image_builder().map_err(|e| {
-                    Error::internal_error()
-                        .with_message(format!("unable to create image builder: {e}"))
-                })?;
-                let mut command = BuildAndScanCommand::new(
-                    &image_builder,
-                    &image_scanner,
-                    &self.interactor,
-                    location,
-                );
-                command.execute().await.map(|_| None)
+                self.execute_build_and_scan(location).await
             }
         };
 
-        match result {
-            Ok(_) => result,
-            Err(mut e) => {
-                self.interactor
-                    .show_message(MessageType::ERROR, e.to_string().as_str())
-                    .await;
-                e.message = format!("error calling command: '{command}': {e}").into();
-                Err(e)
-            }
+        if let Err(e) = &result {
+            self.interactor
+                .show_message(MessageType::ERROR, e.to_string().as_str())
+                .await;
+            return Err(Error {
+                code: e.code,
+                message: format!("error calling command: '{command}': {}", e.message).into(),
+                data: e.data.clone(),
+            });
         }
+
+        result
     }
 
     pub async fn shutdown(&self) -> Result<()> {