Enable caching on connections

lib/train/errors.rb

@@ -21,6 +21,9 @@ class ClientError < ::StandardError; end
   # in the transport layer.
   class TransportError < ::StandardError; end
 
-  # Exception for when no platform can be detected
+  # Exception for when no platform can be detected.
   class PlatformDetectionFailed < ::StandardError; end
+
+  # Exception for when a invalid cache type is passed.
+  class UnknownCacheType < ::StandardError; end
 end

lib/train/plugins/base_connection.rb

@@ -15,18 +15,42 @@ class Train::Plugins::Transport
   class BaseConnection
     include Train::Extras
 
-    # Provide access to the files cache.
-    attr_reader :files
-
     # Create a new Connection instance.
     #
     # @param options [Hash] connection options
     # @yield [self] yields itself for block-style invocation
     def initialize(options = nil)
       @options = options || {}
       @logger = @options.delete(:logger) || Logger.new(STDOUT)
-      @files = {}
       Train::Platforms::Detect::Specifications::OS.load
+
+      # default caching options
+      @cache_enabled = {
+        file: true,
+        command: false,
+      }
+
+      @cache = {}
+      @cache_enabled.each_key do |type|
+        clear_cache(type)
+      end
+    end
+
+    def cache_enabled?(type)
+      @cache_enabled[type.to_sym]
+    end
+
+    # Enable caching types for Train. Currently we support
+    # :file and :command types
+    def enable_cache(type)
+      fail Train::UnknownCacheType, "#{type} is not a valid cache type" unless @cache_enabled.keys.include?(type.to_sym)
+      @cache_enabled[type.to_sym] = true
+    end
+
+    def disable_cache(type)
+      fail Train::UnknownCacheType, "#{type} is not a valid cache type" unless @cache_enabled.keys.include?(type.to_sym)
+      @cache_enabled[type.to_sym] = false
+      clear_cache(type.to_sym)
     end
 
     # Closes the session connection, if it is still active.
@@ -36,14 +60,14 @@ def close
 
     def to_json
       {
-        'files' => Hash[@files.map { |x, y| [x, y.to_json] }],
+        'files' => Hash[@cache[:file].map { |x, y| [x, y.to_json] }],
       }
     end
 
     def load_json(j)
       require 'train/transports/mock'
       j['files'].each do |path, jf|
-        @files[path] = Train::Transports::Mock::Connection::File.from_json(jf)
+        @cache[:file][path] = Train::Transports::Mock::Connection::File.from_json(jf)
       end
     end
 
@@ -52,39 +76,37 @@ def local?
       false
     end
 
-    # Execute a command using this connection.
-    #
-    # @param command [String] command string to execute
-    # @return [CommandResult] contains the result of running the command
-    def run_command(_command)
-      fail Train::ClientError, "#{self.class} does not implement #run_command()"
-    end
-
     # Get information on the operating system which this transport connects to.
     #
     # @return [Platform] system information
     def platform
       @platform ||= Train::Platforms::Detect.scan(self)
     end
-
     # we need to keep os as a method for backwards compatibility with inspec
     alias os platform
 
-    # Interact with files on the target. Read, write, and get metadata
-    # from files via the transport.
-    #
-    # @param [String] path which is being inspected
-    # @return [FileCommon] file object that allows for interaction
-    def file(_path, *_args)
-      fail Train::ClientError, "#{self.class} does not implement #file(...)"
+    # This is the main command call for all connections. This will call the private
+    # run_command_via_connection on the connection with optional caching
+    def run_command(cmd)
+      return run_command_via_connection(cmd) unless cache_enabled?(:command)
+
+      @cache[:command][cmd] ||= run_command_via_connection(cmd)
+    end
+
+    # This is the main file call for all connections. This will call the private
+    # file_via_connection on the connection with optional caching
+    def file(path, *args)
+      return file_via_connection(path, *args) unless cache_enabled?(:file)
+
+      @cache[:file][path] ||= file_via_connection(path, *args)
     end
 
     # Builds a LoginCommand which can be used to open an interactive
     # session on the remote host.
     #
     # @return [LoginCommand] array of command line tokens
     def login_command
-      fail Train::ClientError, "#{self.class} does not implement #login_command()"
+      fail NotImplementedError, "#{self.class} does not implement #login_command()"
     end
 
     # Block and return only when the remote host is prepared and ready to
@@ -98,6 +120,27 @@ def wait_until_ready
 
     private
 
+    # Execute a command using this connection.
+    #
+    # @param command [String] command string to execute
+    # @return [CommandResult] contains the result of running the command
+    def run_command_via_connection(_command)
+      fail NotImplementedError, "#{self.class} does not implement #run_command_via_connection()"
+    end
+
+    # Interact with files on the target. Read, write, and get metadata
+    # from files via the transport.
+    #
+    # @param [String] path which is being inspected
+    # @return [FileCommon] file object that allows for interaction
+    def file_via_connection(_path, *_args)
+      fail NotImplementedError, "#{self.class} does not implement #file_via_connection(...)"
+    end
+
+    def clear_cache(type)
+      @cache[type.to_sym] = {}
+    end
+
     # @return [Logger] logger for reporting information
     # @api private
     attr_reader :logger

lib/train/transports/docker.rb

@@ -69,18 +69,27 @@ def close
       # nothing to do at the moment
     end
 
-    def file(path)
-      @files[path] ||=\
-        if os.aix?
-          Train::File::Remote::Aix.new(self, path)
-        elsif os.solaris?
-          Train::File::Remote::Unix.new(self, path)
-        else
-          Train::File::Remote::Linux.new(self, path)
-        end
+    def uri
+      if @container.nil?
+        "docker://#{@id}"
+      else
+        "docker://#{@container.id}"
+      end
+    end
+
+    private
+
+    def file_via_connection(path)
+      if os.aix?
+        Train::File::Remote::Aix.new(self, path)
+      elsif os.solaris?
+        Train::File::Remote::Unix.new(self, path)
+      else
+        Train::File::Remote::Linux.new(self, path)
+      end
     end
 
-    def run_command(cmd)
+    def run_command_via_connection(cmd)
       cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil?
       stdout, stderr, exit_status = @container.exec(
         [
@@ -93,13 +102,5 @@ def run_command(cmd)
       # @TODO: differentiate any other error
       raise
     end
-
-    def uri
-      if @container.nil?
-        "docker://#{@id}"
-      else
-        "docker://#{@container.id}"
-      end
-    end
   end
 end

lib/train/transports/local.rb

@@ -23,34 +23,35 @@ def initialize(options)
         @cmd_wrapper = CommandWrapper.load(self, options)
       end
 
-      def run_command(cmd)
-        cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil?
-        res = Mixlib::ShellOut.new(cmd)
-        res.run_command
-        CommandResult.new(res.stdout, res.stderr, res.exitstatus)
-      rescue Errno::ENOENT => _
-        CommandResult.new('', '', 1)
+      def login_command
+        nil # none, open your shell
+      end
+
+      def uri
+        'local://'
       end
 
       def local?
         true
       end
 
-      def file(path)
-        @files[path] ||= \
-          if os.windows?
-            Train::File::Local::Windows.new(self, path)
-          else
-            Train::File::Local::Unix.new(self, path)
-          end
-      end
+      private
 
-      def login_command
-        nil # none, open your shell
+      def run_command_via_connection(cmd)
+        cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil?
+        res = Mixlib::ShellOut.new(cmd)
+        res.run_command
+        CommandResult.new(res.stdout, res.stderr, res.exitstatus)
+      rescue Errno::ENOENT => _
+        CommandResult.new('', '', 1)
       end
 
-      def uri
-        'local://'
+      def file_via_connection(path)
+        if os.windows?
+          Train::File::Local::Windows.new(self, path)
+        else
+          Train::File::Local::Unix.new(self, path)
+        end
       end
     end
   end

lib/train/transports/mock.rb

@@ -57,13 +57,13 @@ def trace_calls
 
 class Train::Transports::Mock
   class Connection < BaseConnection
-    attr_accessor :files, :commands
     attr_reader :os
 
     def initialize(conf = nil)
       super(conf)
       mock_os
-      @commands = {}
+      enable_cache(:file)
+      enable_cache(:command)
     end
 
     def uri
@@ -90,8 +90,24 @@ def mock_os_hierarchy(plat)
       end
     end
 
+    def commands=(commands)
+      @cache[:command] = commands
+    end
+
+    def commands
+      @cache[:command]
+    end
+
+    def files=(files)
+      @cache[:file] = files
+    end
+
+    def files
+      @cache[:file]
+    end
+
     def mock_command(cmd, stdout = nil, stderr = nil, exit_status = 0)
-      @commands[cmd] = Command.new(stdout || '', stderr || '', exit_status)
+      @cache[:command][cmd] = Command.new(stdout || '', stderr || '', exit_status)
     end
 
     def command_not_found(cmd)
@@ -104,24 +120,25 @@ def command_not_found(cmd)
       mock_command(cmd, nil, nil, 1)
     end
 
-    def run_command(cmd)
-      @commands[cmd] ||
-        @commands[Digest::SHA256.hexdigest cmd.to_s] ||
-        command_not_found(cmd)
-    end
-
     def file_not_found(path)
       STDERR.puts('File not mocked: '+path.to_s) if @options[:verbose]
       File.new(self, path)
     end
 
-    def file(path)
-      @files[path] ||= file_not_found(path)
-    end
-
     def to_s
       'Mock Connection'
     end
+
+    private
+
+    def run_command_via_connection(cmd)
+      @cache[:command][Digest::SHA256.hexdigest cmd.to_s] ||
+        command_not_found(cmd)
+    end
+
+    def file_via_connection(path)
+      file_not_found(path)
+    end
   end
 end