lib/pgslice.rb in pgslice-0.2.0 vs lib/pgslice.rb in pgslice-0.2.1

@@ -77,11 +77,11 @@
         sql_format = SQL_FORMAT[period.to_sym]
         queries << <<-SQL
 CREATE FUNCTION #{trigger_name}()
     RETURNS trigger AS $$
     BEGIN
-        RAISE EXCEPTION 'Date out of range. Create partitions first.';
+        RAISE EXCEPTION 'Create partitions first.';
     END;
     $$ LANGUAGE plpgsql;
         SQL
 
         queries << <<-SQL
@@ -131,20 +131,17 @@
       index_defs = execute("select pg_get_indexdef(indexrelid) from pg_index where indrelid = $1::regclass AND indisprimary = 'f'", [original_table]).map { |r| r["pg_get_indexdef"] }
       primary_key = self.primary_key(table)
 
       queries = []
 
-      comment = execute("SELECT obj_description(oid, 'pg_trigger') AS comment FROM pg_trigger WHERE tgname = $1 AND tgrelid = $2::regclass", [trigger_name, table]).first
-      if comment
-        field, period = comment["comment"].split(",").map { |v| v.split(":").last } rescue [nil, nil]
-      end
+      period, field, needs_comment = settings_from_trigger(original_table, table)
+      abort "Could not read settings" unless period
 
-      unless period
-        period, field = settings_from_table(original_table, table)
+      if needs_comment
         queries << "COMMENT ON TRIGGER #{trigger_name} ON #{table} is 'column:#{field},period:#{period}';"
       end
-      abort "Could not read settings" unless period
+
       # today = utc date
       today = round_date(DateTime.now.new_offset(0).to_date, period)
       added_partitions = []
       range.each do |n|
         day = advance_date(today, period, n)
@@ -227,11 +224,11 @@
       end
 
       abort "Table not found: #{source_table}" unless table_exists?(source_table)
       abort "Table not found: #{dest_table}" unless table_exists?(dest_table)
 
-      period, field = settings_from_table(table, dest_table)
+      period, field, needs_comment = settings_from_trigger(table, dest_table)
 
       if period
         name_format = self.name_format(period)
 
         existing_tables = self.existing_tables(like: "#{table}_%").select { |t| /#{Regexp.escape("#{table}_")}(\d{4,6})/.match(t) }.sort
@@ -256,20 +253,23 @@
           max_dest_id = min_source_id - 1 if min_source_id
         end
       end
 
       starting_id = max_dest_id
-      fields = columns(source_table).join(", ")
+      fields = columns(source_table).map { |c| PG::Connection.quote_ident(c) }.join(", ")
       batch_size = options[:batch_size]
 
       i = 1
       batch_count = ((max_source_id - starting_id) / batch_size.to_f).ceil
       while starting_id < max_source_id
         where = "#{primary_key} > #{starting_id} AND #{primary_key} <= #{starting_id + batch_size}"
         if period
           where << " AND #{field} >= #{sql_date(starting_time)} AND #{field} < #{sql_date(ending_time)}"
         end
+        if options[:where]
+          where << " AND #{options[:where]}"
+        end
 
         query = <<-SQL
 /* #{i} of #{batch_count} */
 INSERT INTO #{dest_table} (#{fields})
     SELECT #{fields} FROM #{source_table}
@@ -346,10 +346,11 @@
         o.boolean "--dry-run", default: false
         o.boolean "--no-partition", default: false
         o.integer "--start"
         o.string "--url"
         o.string "--source-table"
+        o.string "--where"
         o.on "-v", "--version", "print the version" do
           log PgSlice::VERSION
           @exit = true
         end
       end
@@ -525,18 +526,29 @@
       else
         date.next_month(count)
       end
     end
 
-    def settings_from_table(original_table, table)
+    def settings_from_trigger(original_table, table)
       trigger_name = self.trigger_name(original_table)
-      function_def = execute("select pg_get_functiondef(oid) from pg_proc where proname = $1", [trigger_name])[0]
-      return [nil, nil] unless function_def
-      function_def = function_def["pg_get_functiondef"]
-      sql_format = SQL_FORMAT.find { |_, f| function_def.include?("'#{f}'") }
-      return [nil, nil] unless sql_format
-      period = sql_format[0]
-      field = /to_char\(NEW\.(\w+),/.match(function_def)[1]
-      [period, field]
+
+      needs_comment = false
+      comment = execute("SELECT obj_description(oid, 'pg_trigger') AS comment FROM pg_trigger WHERE tgname = $1 AND tgrelid = $2::regclass", [trigger_name, table])[0]
+      if comment
+        field, period = comment["comment"].split(",").map { |v| v.split(":").last } rescue [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]
+        return [nil, nil] unless function_def
+        function_def = function_def["pg_get_functiondef"]
+        sql_format = SQL_FORMAT.find { |_, f| function_def.include?("'#{f}'") }
+        return [nil, nil] unless sql_format
+        period = sql_format[0]
+        field = /to_char\(NEW\.(\w+),/.match(function_def)[1]
+      end
+
+      [period, field, needs_comment]
     end
   end
 end