[WIP] Add file splitting to evm:db tasks

app/models/file_depot_ftp.rb

@@ -1,6 +1,8 @@
 require 'net/ftp'
 
 class FileDepotFtp < FileDepot
+  include ManageIQ::Util::FtpLib
+
   attr_accessor :ftp
 
   def self.uri_prefix
@@ -62,52 +64,12 @@ def with_connection(cred_hash = nil)
     end
   end
 
-  def connect(cred_hash = nil)
-    host = URI(uri).hostname
-
-    begin
-      _log.info("Connecting to #{self.class.name}: #{name} host: #{host}...")
-      @ftp         = Net::FTP.new(host)
-      @ftp.passive = true  # Use passive mode to avoid firewall issues see http://slacksite.com/other/ftp.html#passive
-      # @ftp.debug_mode = true if settings[:debug]  # TODO: add debug option
-      creds = cred_hash ? [cred_hash[:username], cred_hash[:password]] : login_credentials
-      @ftp.login(*creds)
-      _log.info("Connected to #{self.class.name}: #{name} host: #{host}")
-    rescue SocketError => err
-      _log.error("Failed to connect.  #{err.message}")
-      raise
-    rescue Net::FTPPermError => err
-      _log.error("Failed to login.  #{err.message}")
-      raise
-    else
-      @ftp
-    end
-  end
-
-  def file_exists?(file_or_directory)
-    !ftp.nlst(file_or_directory.to_s).empty?
-  rescue Net::FTPPermError
-    false
-  end
-
   def self.display_name(number = 1)
     n_('FTP', 'FTPs', number)
   end
 
   private
 
-  def create_directory_structure(directory_path)
-    pwd = ftp.pwd
-    directory_path.to_s.split('/').each do |directory|
-      unless ftp.nlst.include?(directory)
-        _log.info("creating #{directory}")
-        ftp.mkdir(directory)
-      end
-      ftp.chdir(directory)
-    end
-    ftp.chdir(pwd)
-  end
-
   def upload(source, destination)
     create_directory_structure(destination_path)
     _log.info("Uploading file: #{destination} to File Depot: #{name}...")

lib/evm_database_ops.rb

@@ -41,7 +41,7 @@ def self.validate_free_space(database_opts)
     end
   end
 
-  def self.backup(db_opts, connect_opts = {})
+  def self.backup(db_opts, connect_opts = {}, additional_opts = {})
     # db_opts:
     #   :dbname => 'vmdb_production',
     #   :username => 'root',
@@ -53,26 +53,24 @@ def self.backup(db_opts, connect_opts = {})
     #   :password => 'Zug-drep5s',
     #   :remote_file_name => "backup_1",     - Provide a base file name for the uploaded file
 
-    uri = with_mount_session(:backup, db_opts, connect_opts) do |database_opts, session, remote_file_uri|
+    uri = with_mount_session(:backup, db_opts, connect_opts) do |database_opts|
       validate_free_space(database_opts)
-      backup_result = PostgresAdmin.backup(database_opts)
-      session&.add(database_opts[:local_file], remote_file_uri)
-      backup_result
+      PostgresAdmin.backup(pg_admin_opts(database_opts, additional_opts))
     end
     _log.info("[#{merged_db_opts(db_opts)[:dbname]}] database has been backed up to file: [#{uri}]")
     uri
   end
 
-  def self.dump(db_opts, connect_opts = {})
+  def self.dump(db_opts, connect_opts = {}, additional_opts = {})
     # db_opts and connect_opts similar to .backup
 
-    uri = with_mount_session(:dump, db_opts, connect_opts) do |database_opts, _session, _remote_file_uri|
+    uri = with_mount_session(:dump, db_opts, connect_opts) do |database_opts|
       # For database dumps, this isn't going to be as accurate (since the dump
       # size will probably be larger than the calculated BD size), but it still
       # won't hurt to do as a generic way to get a rough idea if we have enough
       # disk space or the appliance for the task.
       validate_free_space(database_opts)
-      PostgresAdmin.backup(database_opts)
+      PostgresAdmin.backup_pg_dump(pg_admin_opts(database_opts, additional_opts))
     end
     _log.info("[#{merged_db_opts(db_opts)[:dbname]}] database has been dumped up to file: [#{uri}]")
     uri
@@ -89,10 +87,7 @@ def self.restore(db_opts, connect_opts = {})
     #   :username => 'samba_one',
     #   :password => 'Zug-drep5s',
 
-    uri = with_mount_session(:restore, db_opts, connect_opts) do |database_opts, session, remote_file_uri|
-      if session && !File.exist?(database_opts[:local_file])
-        database_opts[:local_file] = session.download(database_opts[:local_file], remote_file_uri)
-      end
+    uri = with_mount_session(:restore, db_opts, connect_opts) do |database_opts|
       prepare_for_restore(database_opts[:local_file])
 
       # remove all the connections before we restore; AR will reconnect on the next query
@@ -124,12 +119,27 @@ def self.restore(db_opts, connect_opts = {})
       db_opts[:local_file] = session.uri_to_local_path(uri)
     end
 
-    block_result = yield(db_opts, session, uri) if block_given?
+    download_from_mount_if_needed(action, session, uri, database_opts)
+    block_result = yield(db_opts) if block_given?
+    upload_to_mount_if_needed(action, session, uri, database_opts)
     uri || block_result
   ensure
     session.disconnect if session
   end
 
+  private_class_method def self.download_from_mount_if_needed(action, session, uri, db_opts)
+    if action == :restore && session.kind_of?(MiqS3Session) && !File.exist?(db_opts[:local_file])
+      db_opts[:local_file] = session.download(db_opts[:local_file], uri)
+    end
+  end
+
+  private_class_method def self.upload_to_mount_if_needed(action, session, uri, db_opts)
+    if action == :backup && session.kind_of?(MiqS3Session)
+      session.add(database_opts[:local_file], uri)
+      FileUtils.rm_rf database_opts[:local_file] if false # TODO: consider doing this
+    end
+  end
+
   private_class_method def self.prepare_for_restore(filename)
     backup_type = validate_backup_file_type(filename)
 
@@ -195,4 +205,19 @@ def self.backup_file_name(action = :backup)
     "#{action == :backup ? BACKUP_TMP_FILE : DUMP_TMP_FILE}_#{time_suffix}"
   end
   private_class_method :backup_file_name
+
+  FILE_SPLITTER = File.expand_path("../manageiq/util/file_splitter.rb", __FILE__).freeze
+  def self.pg_admin_opts(db_opts, additional_opts)
+    if additional_opts[:byte_count]
+      split_opts = [
+        additional_opts.slice(:byte_count),
+        "-",
+        db_opts[:local_file]
+      ]
+      db_opts.merge(:local_file => "-", :pipe => [[FILE_SPLITTER, :params => split_opts]])
+    else
+      db_opts
+    end
+  end
+  private_class_method :pg_admin_opts
 end

lib/manageiq/util/file_splitter.rb

@@ -0,0 +1,204 @@
+#!/usr/bin/env ruby
+#
+# Yes, you guessed it... this is basically the `split` command...
+#
+# The intent of this is to allow for greater flexibility and utility than is
+# provided by the BSD/GNU variants however, specifically allowing to pipe from
+# a pg_dump directly to a upload target (whether it be a mounted volume or
+# something like a FTP endpoint, though the former is basically already
+# supported with vanilla `split`).
+#
+# FTP, specifically, will be supported natively since ruby has Net::FTP support
+# built in (feature WIP).  This also allows for this to work correctly cross
+# platform, without having to be concerned with differences in `split`
+# functionality.
+
+require_relative 'ftp_lib'
+require 'optparse'
+
+module ManageIQ
+  module Util
+    class FileSplitter
+      include ManageIQ::Util::FtpLib
+
+      KILOBYTE = 1024
+      MEGABYTE = KILOBYTE * 1024
+      GIGABYTE = MEGABYTE * 1024
+
+      BYTE_HASH = {
+        "k" => KILOBYTE,
+        "m" => MEGABYTE,
+        "g" => GIGABYTE
+      }.freeze
+
+      attr_accessor :input_file, :byte_count
+
+      class << self
+        attr_writer :instance_logger
+
+        # Don't log by default, but allow this to work with FtpLib logging.
+        def instance_logger
+          @instance_logger ||= Logger.new(File::NULL)
+        end
+      end
+
+      def self.run(options = nil)
+        options ||= parse_argv
+        new(options).split
+      end
+
+      def self.parse_argv
+        options = {}
+        OptionParser.new do |opt|
+          opt.on("-b", "--byte-count=BYTES", "Number of bytes for each split") do |bytes|
+            options[:byte_count] = parse_byte_value(bytes)
+          end
+          opt.on("--ftp-host=HOST", "Host of the FTP server") do |host|
+            options[:ftp_host] = host
+          end
+          opt.on("--ftp-dir=DIR", "Dir on the FTP server to save files") do |dir|
+            options[:ftp_dir] = dir
+          end
+          opt.on("-v", "--verbose", "Turn on logging") do
+            options[:verbose] = logging
+          end
+        end.parse!
+
+        input_file, file_pattern = determine_input_file_and_file_pattern
+
+        options[:input_file]     = input_file
+        options[:input_filename] = file_pattern
+
+        options
+      end
+
+      def initialize(options = {})
+        @input_file     = options[:input_file] || ARGF
+        @input_filename = options[:input_filename]
+        @byte_count     = options[:byte_count] || (10 * MEGABYTE)
+        @position       = 0
+
+        setup_logging(options)
+        setup_ftp(options)
+      end
+
+      def split
+        until input_file.eof?
+          if ftp
+            split_ftp
+          else
+            split_local
+          end
+          @position += byte_count
+        end
+      ensure
+        input_file.close
+        ftp.close if ftp
+      end
+
+      private
+
+      def setup_logging(options)
+        self.class.instance_logger = Logger.new(STDOUT) if options[:verbose]
+      end
+
+      def setup_ftp(options)
+        if options[:ftp_host]
+          @uri      = options[:ftp_host]
+          @ftp_user = options[:ftp_user] || ENV["FTP_USERNAME"] || "anonymous"
+          @ftp_pass = options[:ftp_pass] || ENV["FTP_PASSWORD"]
+          @ftp      = connect
+
+          @input_filename = File.join(options[:ftp_dir] || "", File.basename(input_filename))
+        end
+      end
+
+      def login_credentials
+        [@ftp_user, @ftp_pass]
+      end
+
+      def split_local
+        File.open(next_split_filename, "w") do |split_file|
+          split_file << input_file.read(byte_count)
+        end
+      end
+
+      # Specific version of Net::FTP#storbinary that doesn't use an existing local
+      # file, and only uploads a specific size from the input_file
+      FTP_CHUNKSIZE = ::Net::FTP::DEFAULT_BLOCKSIZE
+      def split_ftp
+        ftp_mkdir_p
+        ftp.synchronize do
+          ftp.send(:with_binary, true) do
+            conn         = ftp.send(:transfercmd, "STOR #{next_split_filename}")
+            buf_left     = byte_count
+            while buf_left.positive?
+              cur_readsize = buf_left - FTP_CHUNKSIZE >= 0 ? FTP_CHUNKSIZE : buf_left
+              buf = input_file.read(cur_readsize)
+              break if buf == nil # rubocop:disable Style/NilComparison (from original)
+              conn.write(buf)
+              buf_left -= FTP_CHUNKSIZE
+            end
+            conn.close
+            ftp.send(:voidresp)
+          end
+        end
+      rescue Errno::EPIPE
+        # EPIPE, in this case, means that the data connection was unexpectedly
+        # terminated.  Rather than just raising EPIPE to the caller, check the
+        # response on the control connection.  If getresp doesn't raise a more
+        # appropriate exception, re-raise the original exception.
+        getresp
+        raise
+      end
+
+      def ftp_mkdir_p
+        dir_path = File.dirname(input_filename)[1..-1].split('/') - ftp.pwd[1..-1].split("/")
+        create_directory_structure(dir_path.join('/'))
+      end
+
+      def input_filename
+        @input_filename ||= File.expand_path(input_file.path)
+      end
+
+      def next_split_filename
+        "#{input_filename}.#{'%05d' % (@position / byte_count + 1)}"
+      end
+
+      def self.parse_byte_value(bytes)
+        match = bytes.match(/^(?<BYTE_NUM>\d+)(?<BYTE_QUALIFIER>K|M|G)?$/i)
+        raise ArgumentError, "Invalid byte-count", [] if match.nil?
+
+        bytes = match[:BYTE_NUM].to_i
+        if match[:BYTE_QUALIFIER]
+          bytes *= BYTE_HASH[match[:BYTE_QUALIFIER].downcase]
+        end
+        bytes
+      end
+      private_class_method :parse_byte_value
+
+      def self.determine_input_file_and_file_pattern
+        input_file   = ARGV.shift
+        file_pattern = nil
+
+        case input_file
+        when "-"
+          input_file = nil
+        else
+          if input_file && File.exist?(input_file)
+            input_file = File.open(input_file)
+          else
+            file_pattern, input_file = input_file, nil
+          end
+        end
+        file_pattern ||= ARGV.shift
+        raise ArgumentError, "must pass a file pattern if piping from STDIN" if file_pattern.nil? && input_file.nil?
+
+        [input_file, file_pattern]
+      end
+      private_class_method :determine_input_file_and_file_pattern
+    end
+  end
+end
+
+ManageIQ::Util::FileSplitter.run if $PROGRAM_NAME == __FILE__