Merge pull request #32816 from vespa-engine/theodorkl/runYQLQuery

LSP Refactor yql querying
vespa-engine · Nov 8, 2024 · 1696093 · 1696093
2 parents 49b9232 + 7117650
commit 1696093
Showing 1 changed file with 58 additions and 49 deletions.
diff --git a/...server/src/main/java/ai/vespa/schemals/lsp/common/command/commandtypes/RunVespaQuery.java b/...server/src/main/java/ai/vespa/schemals/lsp/common/command/commandtypes/RunVespaQuery.java
@@ -8,6 +8,7 @@
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 
 import org.apache.commons.math3.analysis.function.Pow;
 import org.eclipse.lsp4j.ApplyWorkspaceEditParams;
@@ -53,31 +54,34 @@ public boolean setArguments(List<Object> arguments) {
     public Object execute(EventExecuteCommandContext context) {
         context.logger.info("Running Vespa query: " + queryCommand);
 
-        QueryResult result = runVespaQuery(queryCommand, context.logger);
-
-        if (!result.success()) {
-            context.messageHandler.sendMessage(MessageType.Error, "Failed to run query:\n" + result.result());
-            return null;
-        }
+        runVespaQuery(queryCommand, context.logger).thenAccept(result -> {
+            if (!result.success()) {
+                if (result.result().toLowerCase().contains("command not found")) {
+                    context.messageHandler.sendMessage(MessageType.Error, "Could not find vespa CLI. Make sure vespa CLI is installed and added to path. Download vespa CLI here: https://docs.vespa.ai/en/vespa-cli.html");
+                    return;
+                }
+                context.messageHandler.sendMessage(MessageType.Error, "Failed to run query:\n" + result.result());
+                return;
+            }
 
-        String response = result.result();
+            String response = result.result();
 
-        String targetFileURI = getTargetFileURI(sourceFileURI);
+            String targetFileURI = getTargetFileURI(sourceFileURI);
 
-        CreateFile newFile = new CreateFile(targetFileURI);
+            CreateFile newFile = new CreateFile(targetFileURI);
 
-        TextEdit textEdit = new TextEdit(new Range(new Position(0, 0), new Position(0, 0)), response);
+            TextEdit textEdit = new TextEdit(new Range(new Position(0, 0), new Position(0, 0)), response);
 
-        WorkspaceEdit wsEdit = new WorkspaceEditBuilder()
-            .addResourceOperation(newFile)
-            .addTextEdit(targetFileURI, textEdit)
-            .build();
+            WorkspaceEdit wsEdit = new WorkspaceEditBuilder()
+                .addResourceOperation(newFile)
+                .addTextEdit(targetFileURI, textEdit)
+                .build();
 
-        context.messageHandler.applyEdit(new ApplyWorkspaceEditParams(wsEdit)).thenRun(() -> {
-            context.messageHandler.showDocument(targetFileURI);
+            context.messageHandler.applyEdit(new ApplyWorkspaceEditParams(wsEdit)).thenRun(() -> {
+                context.messageHandler.showDocument(targetFileURI);
+            });
         });
 
-
         return null;
     }
 
@@ -97,50 +101,55 @@ private static String getTargetFileURI(String sourceFileURI) {
 
     private record QueryResult(boolean success, String result) {};
 
-    private QueryResult runVespaQuery(String query, ClientLogger logger) {
+    private CompletableFuture<QueryResult> runVespaQuery(String query, ClientLogger logger) {
 
         boolean isWindows = System.getProperty("os.name").toLowerCase().startsWith("windows");
 
         ProcessBuilder builder = new ProcessBuilder();
 
+        String queryEscaped = query.replace("\"", "\\\"");
+        String vespaCommand = String.format("vespa query \"%s\"", queryEscaped);
+
         if (isWindows) {
-            builder.command(String.format("cmd.exe /c %s", query)); // TODO: Fix this for window
+            builder.command("cmd.exe", "/c", vespaCommand); // TODO: Test this on windows
         } else {
-            builder.command("/usr/local/bin/vespa", "query", query);
+            builder.command("/bin/sh", "-c", vespaCommand);
         }
 
-        try {
+        return CompletableFuture.supplyAsync(() -> {
+            try {
 
-            Process process = builder.start();
-
-            BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
-
-            String line;
-            StringBuilder output = new StringBuilder();
-            while ((line = reader.readLine()) != null) {
-                output.append(line).append("\n");
-            }
-
-            int exitCode = process.waitFor();
+                Process process = builder.start();
 
-            if (exitCode == 0) {
-                return new QueryResult(true, output.toString());
-            }
+                BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
 
-            BufferedReader errorReader = new BufferedReader(new InputStreamReader(process.getErrorStream()));
-            StringBuilder error = new StringBuilder();
-            while ((line = errorReader.readLine()) != null) {
-                error.append(line).append("\n");
-            }
+                String line;
+                StringBuilder output = new StringBuilder();
+                while ((line = reader.readLine()) != null) {
+                    output.append(line).append("\n");
+                }
 
-            return new QueryResult(false, error.toString());
-
-        } catch (InterruptedException e) {
-            return new QueryResult(false, "Program interrupted");
-        } catch (IOException e) {
-            logger.error(e.getMessage());
-            return new QueryResult(false, "IOException occurred.");
-        }
+                int exitCode = process.waitFor();
+
+                if (exitCode == 0) {
+                    return new QueryResult(true, output.toString());
+                }
+
+                BufferedReader errorReader = new BufferedReader(new InputStreamReader(process.getErrorStream()));
+                StringBuilder error = new StringBuilder();
+                while ((line = errorReader.readLine()) != null) {
+                    error.append(line).append("\n");
+                }
+
+                return new QueryResult(false, error.toString());
+
+            } catch (InterruptedException e) {
+                return new QueryResult(false, "Program interrupted");
+            } catch (IOException e) {
+                logger.error(e.getMessage());
+                return new QueryResult(false, "IOException occurred.");
+            }
+        });
     }
-    
+
 }