From 9177bd8a48e2813d1c9e80567fd554c8574b40cc Mon Sep 17 00:00:00 2001
From: Shawn Neal <sneal@sneal.net>
Date: Tue, 22 Apr 2014 11:29:22 -0700
Subject: [PATCH] Added WQL support to Vagrant communicator execute

- Removed duplication for WQL specific commands
---
 plugins/communicators/winrm/communicator.rb | 27 +++++++++------------
 plugins/communicators/winrm/shell.rb        | 19 +++------------
 2 files changed, 16 insertions(+), 30 deletions(-)

diff --git a/plugins/communicators/winrm/communicator.rb b/plugins/communicators/winrm/communicator.rb
index c04bfbeb3..850b0eba7 100644
--- a/plugins/communicators/winrm/communicator.rb
+++ b/plugins/communicators/winrm/communicator.rb
@@ -57,11 +57,19 @@ module VagrantPlugins
           :command     => command,
           :shell       => :powershell
         }.merge(opts || {})
-        exit_status = do_execute(command, opts[:shell], &block)
-        if opts[:error_check] && exit_status != 0
-          raise_execution_error(opts, exit_status)
+
+        if opts[:shell] == :powershell
+          script  = File.expand_path("../scripts/command_alias.ps1", __FILE__)
+          script  = File.read(script)
+          command = script << "\r\n" << command << "\r\nexit $LASTEXITCODE"
         end
-        exit_status
+
+        output = shell.send(opts[:shell], command, &block)
+
+        return output if opts[:shell] == :wql
+        exitcode = output[:exitcode]
+        raise_execution_error(opts, exitcode) if opts[:error_check] && exitcode != 0
+        exitcode
       end
       alias_method :sudo, :execute
 
@@ -103,17 +111,6 @@ module VagrantPlugins
         )
       end
 
-      def do_execute(command, shell_type, &block)
-        if shell_type == :cmd
-          return shell.cmd(command, &block)[:exitcode]
-        end
-
-        script  = File.expand_path("../scripts/command_alias.ps1", __FILE__)
-        script  = File.read(script)
-        command = script << "\r\n" << command << "\r\nexit $LASTEXITCODE"
-        shell.powershell(command, &block)[:exitcode]
-      end
-
       def raise_execution_error(opts, exit_code)
         # The error classes expect the translation key to be _key, but that makes for an ugly
         # configuration parameter, so we set it here from `error_key`
diff --git a/plugins/communicators/winrm/shell.rb b/plugins/communicators/winrm/shell.rb
index e1bbb5da0..56fc4b992 100644
--- a/plugins/communicators/winrm/shell.rb
+++ b/plugins/communicators/winrm/shell.rb
@@ -57,8 +57,8 @@ module VagrantPlugins
         execute_shell(command, :cmd, &block)
       end
 
-      def wql(query)
-        execute_wql(query)
+      def wql(query, &block)
+        execute_shell(query, :wql, &block)
       end
 
       def upload(from, to)
@@ -92,7 +92,7 @@ module VagrantPlugins
       protected
 
       def execute_shell(command, shell=:powershell, &block)
-        raise Errors::InvalidShell, shell: shell unless shell == :cmd || shell == :powershell
+        raise Errors::WinRMInvalidShell, :shell => shell unless [:powershell, :cmd, :wql].include?(shell)
 
         begin
           execute_shell_with_retry(command, shell, &block)
@@ -108,22 +108,11 @@ module VagrantPlugins
             block.call(:stdout, out) if block_given? && out
             block.call(:stderr, err) if block_given? && err
           end
-          @logger.debug("Exit status: #{output[:exitcode].inspect}")
+          @logger.debug("Output: #{output.inspect}")
           return output
         end
       end
 
-      def execute_wql(query)
-        retryable(:tries => @max_tries, :on => @@exceptions_to_retry_on, :sleep => 10) do
-          @logger.debug("#executing wql: #{query}")
-          output = session.wql(query)
-          @logger.debug("wql result: #{output.inspect}")
-          return output
-        end
-      rescue => e
-        raise_winrm_exception(e, :wql, query)
-      end
-
       def raise_winrm_exception(winrm_exception, shell, command)
         # If the error is a 401, we can return a more specific error message
         if winrm_exception.message.include?("401")