client.rb in pgslice-0.4.4

- old
+ new

@@ -1,35 +1,36 @@
 module PgSlice
   class Client < Thor
     check_unknown_options!
 
-    class_option :url
-    class_option :dry_run, type: :boolean, default: false
+    class_option :url, desc: "Database URL"
+    class_option :dry_run, type: :boolean, default: false, desc: "Print statements without executing"
 
     map %w[--version -v] => :version
 
     def self.exit_on_failure?
       true
     end
 
     SQL_FORMAT = {
       day: "YYYYMMDD",
-      month: "YYYYMM"
+      month: "YYYYMM",
+      year: "YYYY"
     }
 
     def initialize(*args)
       $client = self
       $stdout.sync = true
       $stderr.sync = true
       super
     end
 
     desc "prep TABLE [COLUMN] [PERIOD]", "Create an intermediate table for partitioning"
-    option :partition, type: :boolean, default: true
-    option :trigger_based, type: :boolean, default: false
+    option :partition, type: :boolean, default: true, desc: "Partition the table"
+    option :trigger_based, type: :boolean, default: false, desc: "Use trigger-based partitioning"
     def prep(table, column=nil, period=nil)
-      table = Table.new(qualify_table(table))
+      table = qualify_table(table)
       intermediate_table = table.intermediate_table
       trigger_name = table.trigger_name
 
       unless options[:partition]
         abort "Usage: \"pgslice prep TABLE --no-partition\"" if column || period
@@ -100,11 +101,11 @@
       run_queries(queries)
     end
 
     desc "unprep TABLE", "Undo prep"
     def unprep(table)
-      table = Table.new(qualify_table(table))
+      table = qualify_table(table)
       intermediate_table = table.intermediate_table
       trigger_name = table.trigger_name
 
       abort "Table not found: #{intermediate_table}" unless intermediate_table.exists?
 
@@ -114,15 +115,15 @@
       ]
       run_queries(queries)
     end
 
     desc "add_partitions TABLE", "Add partitions"
-    option :intermediate, type: :boolean, default: false
-    option :past, type: :numeric, default: 0
-    option :future, type: :numeric, default: 0
+    option :intermediate, type: :boolean, default: false, desc: "Add to intermediate table"
+    option :past, type: :numeric, default: 0, desc: "Number of past partitions to add"
+    option :future, type: :numeric, default: 0, desc: "Number of future partitions to add"
     def add_partitions(table)
-      original_table = Table.new(qualify_table(table))
+      original_table = qualify_table(table)
       table = options[:intermediate] ? original_table.intermediate_table : original_table
       trigger_name = original_table.trigger_name
 
       abort "Table not found: #{table}" unless table.exists?
 
@@ -243,21 +244,21 @@
 
       run_queries(queries) if queries.any?
     end
 
     desc "fill TABLE", "Fill the partitions in batches"
-    option :batch_size, type: :numeric, default: 10000
-    option :swapped, type: :boolean, default: false
-    option :source_table
-    option :dest_table
-    option :start
-    option :where
-    option :sleep, type: :numeric
+    option :batch_size, type: :numeric, default: 10000, desc: "Batch size"
+    option :swapped, type: :boolean, default: false, desc: "Use swapped table"
+    option :source_table, desc: "Source table"
+    option :dest_table, desc: "Destination table"
+    option :start, type: :numeric, desc: "Primary key to start"
+    option :where, desc: "Conditions to filter"
+    option :sleep, type: :numeric, desc: "Seconds to sleep between batches"
     def fill(table)
-      table = Table.new(qualify_table(table))
-      source_table = Table.new(options[:source_table]) if options[:source_table]
-      dest_table = Table.new(options[:dest_table]) if options[:dest_table]
+      table = qualify_table(table)
+      source_table = qualify_table(options[:source_table]) if options[:source_table]
+      dest_table = qualify_table(options[:dest_table]) if options[:dest_table]
 
       if options[:swapped]
         source_table ||= table.retired_table
         dest_table ||= table
       else
@@ -343,13 +344,13 @@
         end
       end
     end
 
     desc "swap TABLE", "Swap the intermediate table with the original table"
-    option :lock_timeout, default: "5s"
+    option :lock_timeout, default: "5s", desc: "Lock timeout"
     def swap(table)
-      table = Table.new(qualify_table(table))
+      table = qualify_table(table)
       intermediate_table = table.intermediate_table
       retired_table = table.retired_table
 
       abort "Table not found: #{table}" unless table.exists?
       abort "Table not found: #{intermediate_table}" unless intermediate_table.exists?
@@ -369,11 +370,11 @@
       run_queries(queries)
     end
 
     desc "unswap TABLE", "Undo swap"
     def unswap(table)
-      table = Table.new(qualify_table(table))
+      table = qualify_table(table)
       intermediate_table = table.intermediate_table
       retired_table = table.retired_table
 
       abort "Table not found: #{table}" unless table.exists?
       abort "Table not found: #{retired_table}" unless retired_table.exists?
@@ -390,13 +391,13 @@
 
       run_queries(queries)
     end
 
     desc "analyze TABLE", "Analyze tables"
-    option :swapped, type: :boolean, default: false
+    option :swapped, type: :boolean, default: false, desc: "Use swapped table"
     def analyze(table)
-      table = Table.new(qualify_table(table))
+      table = qualify_table(table)
       parent_table = options[:swapped] ? table : table.intermediate_table
 
       existing_tables = table.existing_partitions
       analyze_list = existing_tables + [parent_table]
       run_queries_without_transaction(analyze_list.map { |t| "ANALYZE VERBOSE #{quote_table(t)};" })
@@ -500,32 +501,38 @@
 
     def name_format(period)
       case period.to_sym
       when :day
         "%Y%m%d"
-      else
+      when :month
         "%Y%m"
+      else
+        "%Y"
       end
     end
 
     def round_date(date, period)
       date = date.to_date
       case period.to_sym
       when :day
         date
-      else
+      when :month
         Date.new(date.year, date.month)
+      else
+        Date.new(date.year)
       end
     end
 
     def advance_date(date, period, count = 1)
       date = date.to_date
       case period.to_sym
       when :day
         date.next_day(count)
-      else
+      when :month
         date.next_month(count)
+      else
+        date.next_year(count)
       end
     end
 
     def quote_ident(value)
       PG::Connection.quote_ident(value)
@@ -538,11 +545,11 @@
     def quote_no_schema(table)
       quote_ident(table.to_s.split(".", 2)[-1])
     end
 
     def qualify_table(table)
-      table.to_s.include?(".") ? table : [schema, table].join(".")
+      Table.new(table.to_s.include?(".") ? table : [schema, table].join("."))
     end
 
     def settings_from_trigger(original_table, table)
       trigger_name = original_table.trigger_name
 
@@ -553,10 +560,10 @@
         field, period, cast = comment["comment"].split(",").map { |v| v.split(":").last } rescue [nil, nil, nil]
       end
 
       unless period
         needs_comment = true
-        function_def = execute("select pg_get_functiondef(oid) from pg_proc where proname = $1", [trigger_name])[0]
+        function_def = execute("SELECT pg_get_functiondef(oid) FROM pg_proc WHERE proname = $1", [trigger_name])[0]
         return [] unless function_def
         function_def = function_def["pg_get_functiondef"]
         sql_format = SQL_FORMAT.find { |_, f| function_def.include?("'#{f}'") }
         return [] unless sql_format
         period = sql_format[0]