feat: add code actions and code lens for docker-compose

Author: tembleking

src/app/lsp_server.rs

@@ -19,6 +19,7 @@ use super::commands::CommandExecutor;
 use super::component_factory::{ComponentFactory, Config};
 use super::queries::QueryExecutor;
 use super::{InMemoryDocumentDatabase, LSPClient};
+use crate::infra::parse_compose_file;
 
 pub struct LSPServer<C> {
     command_executor: CommandExecutor<C>,
@@ -83,6 +84,94 @@ impl TryFrom<&str> for SupportedCommands {
     }
 }
 
+struct CommandInfo {
+    title: String,
+    command: String,
+    arguments: Option<Vec<Value>>,
+    range: Range,
+}
+
+impl<C> LSPServer<C>
+where
+    C: LSPClient + Send + Sync + 'static,
+{
+    fn generate_commands_for_uri(
+        &self,
+        uri: &tower_lsp::lsp_types::Url,
+        content: &str,
+    ) -> Vec<CommandInfo> {
+        let file_uri = uri.as_str();
+
+        if file_uri.contains("docker-compose.yml")
+            || file_uri.contains("compose.yml")
+            || file_uri.contains("docker-compose.yaml")
+            || file_uri.contains("compose.yaml")
+        {
+            self.generate_compose_commands(uri, content)
+        } else {
+            self.generate_dockerfile_commands(uri, content)
+        }
+    }
+
+    fn generate_compose_commands(
+        &self,
+        uri: &tower_lsp::lsp_types::Url,
+        content: &str,
+    ) -> Vec<CommandInfo> {
+        let mut commands = vec![];
+        if let Ok(instructions) = parse_compose_file(content) {
+            for instruction in instructions {
+                commands.push(CommandInfo {
+                    title: "Scan base image".to_string(),
+                    command: SupportedCommands::ExecuteBaseImageScan.to_string(),
+                    arguments: Some(vec![json!(uri), json!(instruction.range.start.line)]),
+                    range: instruction.range,
+                });
+            }
+        }
+        commands
+    }
+
+    fn generate_dockerfile_commands(
+        &self,
+        uri: &tower_lsp::lsp_types::Url,
+        content: &str,
+    ) -> Vec<CommandInfo> {
+        let mut commands = vec![];
+        if let Some(last_line_starting_with_from_statement) = content
+            .lines()
+            .enumerate()
+            .filter(|(_, line)| line.trim_start().starts_with("FROM "))
+            .map(|(line_num, _)| line_num)
+            .last()
+        {
+            let range = Range::new(
+                Position::new(last_line_starting_with_from_statement as u32, 0),
+                Position::new(last_line_starting_with_from_statement as u32, 0),
+            );
+            commands.push(CommandInfo {
+                title: "Build and scan".to_string(),
+                command: SupportedCommands::ExecuteBuildAndScan.to_string(),
+                arguments: Some(vec![
+                    json!(uri),
+                    json!(last_line_starting_with_from_statement),
+                ]),
+                range,
+            });
+            commands.push(CommandInfo {
+                title: "Scan base image".to_string(),
+                command: SupportedCommands::ExecuteBaseImageScan.to_string(),
+                arguments: Some(vec![
+                    json!(uri),
+                    json!(last_line_starting_with_from_statement),
+                ]),
+                range,
+            });
+        }
+        commands
+    }
+}
+
 #[async_trait::async_trait]
 impl<C> LanguageServer for LSPServer<C>
 where
@@ -152,8 +241,6 @@ where
     }
 
     async fn code_action(&self, params: CodeActionParams) -> Result<Option<CodeActionResponse>> {
-        let mut code_actions = vec![];
-
         let Some(content) = self
             .query_executor
             .get_document_text(params.text_document.uri.as_str())
@@ -165,48 +252,23 @@ where
             )));
         };
 
-        let Some(last_line_starting_with_from_statement) = content
-            .lines()
-            .enumerate()
-            .filter(|(_, line)| line.trim_start().starts_with("FROM "))
-            .map(|(line_num, _)| line_num)
-            .last()
-        else {
-            return Ok(None);
-        };
-
-        let Ok(line_selected_as_usize) = usize::try_from(params.range.start.line) else {
-            return Err(Error::internal_error().with_message(format!(
-                "unable to parse u32 as usize: {}",
-                params.range.start.line
-            )));
-        };
-
-        if last_line_starting_with_from_statement == line_selected_as_usize {
-            code_actions.push(CodeActionOrCommand::Command(Command {
-                title: "Build and scan".to_string(),
-                command: SupportedCommands::ExecuteBuildAndScan.to_string(),
-                arguments: Some(vec![
-                    json!(params.text_document.uri),
-                    json!(line_selected_as_usize),
-                ]),
-            }));
-            code_actions.push(CodeActionOrCommand::Command(Command {
-                title: "Scan base image".to_string(),
-                command: SupportedCommands::ExecuteBaseImageScan.to_string(),
-                arguments: Some(vec![
-                    json!(params.text_document.uri),
-                    json!(line_selected_as_usize),
-                ]),
-            }));
-        }
+        let commands = self.generate_commands_for_uri(&params.text_document.uri, &content);
+        let code_actions: Vec<CodeActionOrCommand> = commands
+            .into_iter()
+            .filter(|cmd| cmd.range.start.line == params.range.start.line)
+            .map(|cmd| {
+                CodeActionOrCommand::Command(Command {
+                    title: cmd.title,
+                    command: cmd.command,
+                    arguments: cmd.arguments,
+                })
+            })
+            .collect();
 
         Ok(Some(code_actions))
     }
 
     async fn code_lens(&self, params: CodeLensParams) -> Result<Option<Vec<CodeLens>>> {
-        let mut code_lens = vec![];
-
         let Some(content) = self
             .query_executor
             .get_document_text(params.text_document.uri.as_str())
@@ -218,48 +280,21 @@ where
             )));
         };
 
-        let Some(last_line_starting_with_from_statement) = content
-            .lines()
-            .enumerate()
-            .filter(|(_, line)| line.trim_start().starts_with("FROM "))
-            .map(|(line_num, _)| line_num)
-            .last()
-        else {
-            return Ok(None);
-        };
-
-        code_lens.push(CodeLens {
-            range: Range::new(
-                Position::new(last_line_starting_with_from_statement as u32, 0),
-                Position::new(last_line_starting_with_from_statement as u32, 0),
-            ),
-            command: Some(Command {
-                title: "Build and scan".to_string(),
-                command: SupportedCommands::ExecuteBuildAndScan.to_string(),
-                arguments: Some(vec![
-                    json!(params.text_document.uri),
-                    json!(last_line_starting_with_from_statement),
-                ]),
-            }),
-            data: None,
-        });
-        code_lens.push(CodeLens {
-            range: Range::new(
-                Position::new(last_line_starting_with_from_statement as u32, 0),
-                Position::new(last_line_starting_with_from_statement as u32, 0),
-            ),
-            command: Some(Command {
-                title: "Scan base image".to_string(),
-                command: SupportedCommands::ExecuteBaseImageScan.to_string(),
-                arguments: Some(vec![
-                    json!(params.text_document.uri),
-                    json!(last_line_starting_with_from_statement),
-                ]),
-            }),
-            data: None,
-        });
+        let commands = self.generate_commands_for_uri(&params.text_document.uri, &content);
+        let code_lenses = commands
+            .into_iter()
+            .map(|cmd| CodeLens {
+                range: cmd.range,
+                command: Some(Command {
+                    title: cmd.title,
+                    command: cmd.command,
+                    arguments: cmd.arguments,
+                }),
+                data: None,
+            })
+            .collect();
 
-        Ok(Some(code_lens))
+        Ok(Some(code_lenses))
     }
 
     async fn execute_command(&self, params: ExecuteCommandParams) -> Result<Option<Value>> {

src/infra/mod.rs

@@ -7,5 +7,6 @@ mod sysdig_image_scanner_result;
 
 pub use sysdig_image_scanner::{SysdigAPIToken, SysdigImageScanner};
 pub mod lsp_logger;
+pub use compose_ast_parser::{ImageInstruction, parse_compose_file};
 pub use docker_image_builder::DockerImageBuilder;
 pub use dockerfile_ast_parser::{Instruction, parse_dockerfile};

tests/fixtures/compose.yaml

@@ -0,0 +1,14 @@
+# Test compose file with various YAML features
+services:
+  # Web service using a quoted key for the image
+  web:
+    "image": nginx:latest # Inline comment
+
+  # Database service using a literal block scalar for the image
+  db:
+    image: |
+      postgres:13
+
+  # Another service for good measure
+  api:
+    image: my-api:1.0

tests/fixtures/docker-compose.yml

@@ -0,0 +1,5 @@
+services:
+  web:
+    image: nginx:latest
+  db:
+    image: postgres:13

tests/general.rs

@@ -159,3 +159,111 @@ async fn when_the_client_asks_for_the_existing_code_lens_but_the_dockerfile_cont
         ]
     );
 }
+
+#[tokio::test]
+async fn when_the_client_asks_for_code_lens_in_a_compose_file_it_receives_them() {
+    let mut client = test::TestClient::new_initialized().await;
+    client
+        .open_file_with_contents(
+            "docker-compose.yml",
+            include_str!("fixtures/docker-compose.yml"),
+        )
+        .await;
+
+    let response = client
+        .request_available_code_lens_in_file("docker-compose.yml")
+        .await;
+
+    assert_eq!(
+        response.unwrap(),
+        vec![
+            CodeLens {
+                range: Range::new(Position::new(2, 11), Position::new(2, 23)),
+                command: Some(Command {
+                    title: "Scan base image".to_string(),
+                    command: "sysdig-lsp.execute-scan".to_string(),
+                    arguments: Some(vec![json!("file://docker-compose.yml/"), json!(2)])
+                }),
+                data: None
+            },
+            CodeLens {
+                range: Range::new(Position::new(4, 11), Position::new(4, 22)),
+                command: Some(Command {
+                    title: "Scan base image".to_string(),
+                    command: "sysdig-lsp.execute-scan".to_string(),
+                    arguments: Some(vec![json!("file://docker-compose.yml/"), json!(4)])
+                }),
+                data: None
+            }
+        ]
+    );
+}
+
+#[tokio::test]
+async fn when_the_client_asks_for_code_actions_in_a_compose_file_it_receives_them() {
+    let mut client = test::TestClient::new_initialized().await;
+    client
+        .open_file_with_contents(
+            "docker-compose.yml",
+            include_str!("fixtures/docker-compose.yml"),
+        )
+        .await;
+
+    let response = client
+        .request_available_actions_in_line("docker-compose.yml", 2)
+        .await;
+
+    assert_eq!(
+        response.unwrap(),
+        vec![CodeActionOrCommand::Command(Command {
+            title: "Scan base image".to_string(),
+            command: "sysdig-lsp.execute-scan".to_string(),
+            arguments: Some(vec![json!("file://docker-compose.yml/"), json!(2)])
+        })]
+    );
+}
+
+#[tokio::test]
+async fn when_the_client_asks_for_code_lens_in_a_complex_compose_yaml_file_it_receives_them() {
+    let mut client = test::TestClient::new_initialized().await;
+    client
+        .open_file_with_contents("compose.yaml", include_str!("fixtures/compose.yaml"))
+        .await;
+
+    let response = client
+        .request_available_code_lens_in_file("compose.yaml")
+        .await;
+
+    assert_eq!(
+        response.unwrap(),
+        vec![
+            CodeLens {
+                range: Range::new(Position::new(4, 13), Position::new(4, 25)),
+                command: Some(Command {
+                    title: "Scan base image".to_string(),
+                    command: "sysdig-lsp.execute-scan".to_string(),
+                    arguments: Some(vec![json!("file://compose.yaml/"), json!(4)])
+                }),
+                data: None
+            },
+            CodeLens {
+                range: Range::new(Position::new(9, 6), Position::new(9, 17)),
+                command: Some(Command {
+                    title: "Scan base image".to_string(),
+                    command: "sysdig-lsp.execute-scan".to_string(),
+                    arguments: Some(vec![json!("file://compose.yaml/"), json!(9)])
+                }),
+                data: None
+            },
+            CodeLens {
+                range: Range::new(Position::new(13, 11), Position::new(13, 21)),
+                command: Some(Command {
+                    title: "Scan base image".to_string(),
+                    command: "sysdig-lsp.execute-scan".to_string(),
+                    arguments: Some(vec![json!("file://compose.yaml/"), json!(13)])
+                }),
+                data: None
+            }
+        ]
+    );
+}