split_pgdump 0.2.0 → 0.3.0

data/bin/split_pgdump

@@ -7,12 +7,15 @@ require 'strscan'
 $debug = false
 
 class CWorker
-  attr_accessor :rules_file, :output_file, :sorter, :rules
+  attr_accessor :rules_file, :output_file, :sorter, :rules, :num_sorters
+  attr_accessor :could_fork
   def initialize
     @rules_file = 'split.rules'
     @output_file = 'dump.sql'
     @sorter = `which sort`.chomp
     @rules = []
+    @num_sorters = 0
+    @could_fork = true
   end
 
   def tables_dir
@@ -49,6 +52,7 @@ class CWorker
       @table = Table.new(tables_dir, @schema, table_name, columns)
       @tables << @table
       puts "Start to write table #{table_name}" if $debug
+      @start_time = Time.now
       @state = :table
     else
       if line =~ /^SET search_path = ([^,]+)/
@@ -62,6 +66,7 @@ class CWorker
     if line =~ /^\\\.[\r\n]/
       @table.flush_all
       @table.copy_lines{|l| out.puts l}
+      puts "Table #{@table.table} copied in #{Time.now - @start_time}s" if $debug
       @table = nil
       @state = :schema
     else
@@ -86,7 +91,35 @@ class CWorker
       end
     end
 
-    @tables.each{|table| table.finish_all}
+    @start_time = Time.now
+    sort_and_finish
+    puts "Finished in #{Time.now - @start_time}s #{Process.pid}" if $debug
+  end
+
+  def sort_and_finish
+    files = []
+    for table in @tables
+      for one_file in table.files.values
+        files << [table.sort_args, one_file]
+      end
+    end
+    if @num_sorters > 1
+      files.each_slice(@num_sorters) do |one_files|
+        cmd = one_files.map{|sort_args, one_file| 
+          one_file.sort_args(sort_args).map{|a|"'#{a}'"}.join(' ') 
+        }
+        cmd = cmd.map{|c| "{ #{c} & }"}  if @could_fork
+        cmd = cmd.join(' ; ')
+        cmd += ' ; wait '  if @could_fork
+        system cmd
+        one_files.each{|sort_args, one_file| one_file.write_finish}
+      end
+    else
+      files.each do |sort_args, one_file|
+        one_file.sort sort_args
+        one_file.write_finish
+      end
+    end
   end
 end
 
@@ -152,7 +185,7 @@ end
 
 class Table
   class NoColumn < StandardError; end
-  ONE_FILE_CACHE_SIZE = 128 * 1024
+  ONE_FILE_CACHE_SIZE = 256 * 1024
   TOTAL_CACHE_SIZE = 5 * 1024 * 1024
   class OneFile
     attr_reader :file_name, :cache_size
@@ -165,19 +198,17 @@ class Table
     def add_line(line)
       @cache_lines << line
       @cache_size += line.size
-      flush if @cache_size > ONE_FILE_CACHE_SIZE
     end
 
-    def flush
+    def flush(&block)
+      @cache_size = 0
       dir = File.dirname(@file_name)
       unless File.directory?(dir)
         FileUtils.mkdir_p(dir)
       end
-      File.open(@file_name, 'a') do |f|
-        @cache_lines.each{|l| f.write(l)}
-      end
+      content = @cache_lines.join
+      File.open(@file_name, 'a'){|f| f.write(content)}
       @cache_lines.clear
-      @cache_size = 0
     end
 
     def write_finish
@@ -186,7 +217,7 @@ class Table
       end
     end
 
-    def sort(sort_line = [])
+    def sort_args(sort_line = [])
       args = [Worker.sorter]
       if sort_line && !sort_line.empty?
         args.concat sort_line
@@ -195,11 +226,15 @@ class Table
       end
       args.push '-o', @file_name, @file_name
       puts args.join(' ')  if $debug
-      system *args
+      args
+    end
+
+    def sort(sort_line = [])
+      system *sort_args(sort_line)
     end
   end
 
-  attr_reader :name, :columns, :files, :sort_line
+  attr_reader :table, :columns, :files, :sort_line, :sort_args
   def initialize(dir, schema, name, columns)
     @dir = dir
     @table = name
@@ -243,7 +278,7 @@ class Table
 
     eval <<-"EOF"
       def self.file_name(values)
-        name = %{#{split_string}}.gsub(/\\.\\.|\\s|\\?|\\*/, '_')
+        name = %{#{split_string}}.gsub(/\\.\\.|\\s|\\?|\\*|'|"/, '_')
         "\#{table_schema}/\#{name}.dat"
       end
     EOF
@@ -268,9 +303,12 @@ class Table
     values = line.chomp.split("\t")
     fname = file_name(values)
     one_file = @files[fname] ||= OneFile.new(@dir, fname)
-    @total_cache_size -= one_file.cache_size
     one_file.add_line(line)
-    @total_cache_size += one_file.cache_size
+    @total_cache_size += line.size
+    if one_file.cache_size > ONE_FILE_CACHE_SIZE
+      @total_cache_size -= one_file.cache_size
+      one_file.flush
+    end
     flush_all if @total_cache_size > TOTAL_CACHE_SIZE
   end
 
@@ -288,13 +326,6 @@ class Table
       to_enum(:copy_lines)
     end
   end
-
-  def finish_all
-    @files.each do |name, one_file|
-      one_file.sort(@sort_args)
-      one_file.write_finish
-    end
-  end
 end
 
 opts = OptionParser.new do |opts|
@@ -316,9 +347,16 @@ effectivly transmitted using rsync, repacking by 7z and other.
           "Table content will be storred in FILE-tables directory") do |v|
     Worker.output_file = v
   end
-  opts.on("-s", "--sort=SORT_BIN", "sort executable compatible with gnu coreutils sort") do |v|
+  opts.on("-s", "--sort=SORT_BIN", "sort executable compatible with gnu coreutils sort (default `which sort`)") do |v|
     Worker.sorter = v
   end
+  opts.on("-n", "--sorters=NUM", Integer, "number of sorters started in a shell",
+                                          "(1 or 0 for safe command line) (default 0)") do |n|
+    Worker.num_sorters = n
+  end
+  opts.on("--no-shell-fork", "could not use shell & for parrallel execution of sorters") do 
+    Worker.could_fork = true
+  end
   opts.on("-d", "--debug", "debug"){|v| $debug = true}
   opts.on_tail("-h", "--help", "this message"){|v| puts opts; exit}
 

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: split_pgdump
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.3.0
   prerelease: 
 platform: ruby
 authors: