Add auto-complete for windows

src/cli/auto_complete.rs

@@ -1,16 +1,40 @@
 use std::{
-    fs::File,
-    io::{self, Cursor, Write},
-    path::PathBuf,
+    collections::HashSet, fs::File, io::{self, Cursor, Write}, path::PathBuf, sync::{Mutex, OnceLock}
 };
 
 use clap::Command;
 use clap_complete::{generate, Generator, Shell};
 
+
 use crate::cli::command;
 
 const AUTO_COMPLETE_FUNCTIONS_ZSH: &str = include_str!("autocomplete_functions/zshell.zsh");
+const AUTO_COMPLETE_FUNCTIONS_POWERSHELL: &str =
+    include_str!("autocomplete_functions/powershell.ps1");
+
+const SATORI_RUN_SUFFIX: &str = "satori;run;";
+
+static AUTO_COMPLETE_DATABASES_TOOLS: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
+static AUTO_COMPLETE_TOOLS: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
+
+
+pub fn add_database_tool_autocomplete(tool_name: String) {
+    let mut tools = AUTO_COMPLETE_DATABASES_TOOLS.get_or_init(|| Mutex::new(HashSet::new())).lock().unwrap();
+    tools.insert(tool_name);
+}
+
+pub fn add_tool_autocomplete(tool_name: String) {
+    let mut tools = AUTO_COMPLETE_TOOLS.get_or_init(|| Mutex::new(HashSet::new())).lock().unwrap();
+    tools.insert(tool_name);
+}
+
+fn get_database_tools() -> HashSet<String> {
+    AUTO_COMPLETE_DATABASES_TOOLS.get().unwrap().lock().unwrap().clone()
+}
 
+fn get_tools() -> HashSet<String> {
+    AUTO_COMPLETE_TOOLS.get().unwrap().lock().unwrap().clone()
+}
 pub fn auto_complete(shell: Shell, out: PathBuf) {
     let mut cmd = command::get();
     eprintln!("Generating completion file for {shell}...");
@@ -34,19 +58,19 @@ fn generate_autocomplete_buf<G: Generator>(
     generate(gen, cmd, cmd.get_name().to_string(), buf);
 }
 
-fn make_dynamic_completion_script(shell: Shell, auto_complete_static: &str) -> String {
+fn make_dynamic_completion_script(shell: Shell, auto_complete_script: &str) -> String {
     match shell {
         Shell::Bash => todo!(),
         Shell::Elvish => unimplemented!(),
         Shell::Fish => unimplemented!(),
-        Shell::PowerShell => todo!(),
-        Shell::Zsh => handle_zsh(auto_complete_static),
+        Shell::PowerShell => handle_power_shell(auto_complete_script),
+        Shell::Zsh => handle_zsh(auto_complete_script),
         _ => todo!(),
     }
 }
 
-fn handle_zsh(auto_complete_static: &str) -> String {
-    let mut auto_complete_static = auto_complete_static
+fn handle_zsh(auto_complete_script: &str) -> String {
+    let mut auto_complete_static = auto_complete_script
         .replace(
             "datastore_name -- The name as defined in Satori data portal:",
             "datastore_name -- The name as defined in Satori data portal:_datastores",
@@ -58,3 +82,67 @@ fn handle_zsh(auto_complete_static: &str) -> String {
     auto_complete_static.push_str(AUTO_COMPLETE_FUNCTIONS_ZSH);
     auto_complete_static
 }
+
+fn handle_power_shell(auto_complete_script: &str) -> String {
+    // We need to count also non bareword as elements, in order to support matching the datastore name which is single quoted
+    let mut auto_complete_string = auto_complete_script.replace("$element.StringConstantType -ne [StringConstantType]::BareWord -or", "");
+
+    // Add the Get-Datastore and Get-Databases functions
+    auto_complete_string.push_str(AUTO_COMPLETE_FUNCTIONS_POWERSHELL);
+
+    let mut moved_curr_complete = "".to_string();
+
+    for tool in get_tools() {
+        // will result in something like: satori;run;mongosh' { 
+        // We need to replace it to the datastore auto-complete
+        let full_command_line = format!("{SATORI_RUN_SUFFIX}{tool}' {{");
+
+        // Get the indexes of the current tool auto-complete and extract them
+        let (start_index_cur_compl, end_index_cur_compl) = get_pwsh_tool_completion_indexes(&auto_complete_string, &full_command_line);
+        let curr_complete = &auto_complete_string[start_index_cur_compl..end_index_cur_compl].to_string();
+
+        // Replace the current tool auto-complete with datastore auto-complete
+        auto_complete_string.replace_range(start_index_cur_compl..end_index_cur_compl, r#"
+            $completionResults += Get-CompletionDatastores
+            $completionResults
+        "#);
+
+        // Return the original auto-complete to a new section, also add the database auto-complete if needed
+        let original_auto_complete = if get_database_tools().contains(&tool)  {
+            format!(r#"
+            if($command.StartsWith('{SATORI_RUN_SUFFIX}{tool}')) {{
+                if ($commandElements.Count -eq 4) {{
+                    $datastoreName = $commandElements[3]
+                    $completions = Get-CompletionDatabases -DatastoreName $datastoreName
+                
+                }} else {{
+                    $completions=@({curr_complete})
+                }}
+            }}
+            "#)
+
+        } else {
+            format!(r#"
+            if($command.StartsWith('{SATORI_RUN_SUFFIX}{tool}')) {{
+                $completions=@({curr_complete})
+            }}
+            "#)
+        };
+        moved_curr_complete.push_str(&original_auto_complete);
+    }
+    let new_s = format!(r#"
+    if ($completions.Count -eq 0) {{
+        {moved_curr_complete}
+    
+    }}
+
+    $completions.Where"#);
+    auto_complete_string.replace("$completions.Where", &new_s)
+}
+
+fn get_pwsh_tool_completion_indexes(auto_complete_string: &str, full_command_name_with_suffix: &str) -> (usize, usize) {
+    let start_index = auto_complete_string.find(full_command_name_with_suffix).and_then(|start_index| Some(start_index + full_command_name_with_suffix.len())).unwrap();
+    let end_index = auto_complete_string[start_index..].find("break").and_then(|end_index| Some(start_index + end_index)).unwrap();
+
+    (start_index, end_index)
+}

src/cli/autocomplete_functions/powershell.ps1

@@ -0,0 +1,21 @@
+function Get-CompletionDatastores {
+    $additionalResults = @()
+    $datastoreResults = Invoke-Expression 'satori list --datastores'
+    $datastoreResults -split '\r?\n' | ForEach-Object {
+        $additionalResults += [CompletionResult]::new("'$_'", "'$_'", [CompletionResultType]::ParameterValue, "Datastore: $_")
+    }
+    return $additionalResults
+}
+
+function Get-CompletionDatabases {
+    param (
+        [string]$DatastoreName
+    )
+        $databases = @()
+        $listOfdatabases = Invoke-Expression "satori list --databases $datastoreName"
+        $listOfdatabases -split '\r?\n' | ForEach-Object {
+            $databases += [CompletionResult]::new($_, $_, [CompletionResultType]::ParameterValue, "Database: $_")
+        }
+        return $databases
+}
+

src/cli/command/run.rs

@@ -2,7 +2,7 @@ use std::path::PathBuf;
 
 use clap::{arg, value_parser, Arg, ArgAction, Command};
 
-use crate::helpers::{self, tools::CliArgs};
+use crate::{cli::auto_complete, helpers::{self, tools::CliArgs}};
 
 use super::common_args;
 
@@ -52,6 +52,12 @@ fn from_file() -> Vec<(&'static str, Vec<CliArgs>)> {
         .value
         .iter()
         .map(|tool| {
+            auto_complete::add_tool_autocomplete(tool.name.clone());
+            for cli_arg in tool.cli_args.iter() {
+                if cli_arg.name == "database" {
+                    auto_complete::add_database_tool_autocomplete(tool.name.clone());
+                }
+            }
             (
                 string_to_static_str(tool.name.clone()),
                 tool.cli_args.clone(),

src/cli/parsers/run/dynamic_tools.rs

@@ -2,8 +2,7 @@ use clap::ArgMatches;
 
 use crate::{
     cli::{
-        parsers::{self, common::build_login_common_args, run::common},
-        Flow,
+        parsers::{self, common::build_login_common_args, run::common}, Flow
     },
     helpers::tools::{self, Tool},
     run::{DynamicTool, Run},