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

@@ -77,21 +77,24 @@
         sql_format = SQL_FORMAT[period.to_sym]
         queries << <<-SQL
 CREATE FUNCTION #{trigger_name}()
     RETURNS trigger AS $$
     BEGIN
-        EXECUTE 'INSERT INTO #{table}_' || to_char(NEW.#{column}, '#{sql_format}') || ' VALUES ($1.*)' USING NEW;
-        RETURN NULL;
+        RAISE EXCEPTION 'Date out of range. Create partitions first.';
     END;
     $$ LANGUAGE plpgsql;
         SQL
 
         queries << <<-SQL
 CREATE TRIGGER #{trigger_name}
     BEFORE INSERT ON #{intermediate_table}
     FOR EACH ROW EXECUTE PROCEDURE #{trigger_name}();
       SQL
+
+        queries << <<-SQL
+COMMENT ON TRIGGER #{trigger_name} ON #{intermediate_table} is 'column:#{column},period:#{period}';
+SQL
       end
 
       run_queries(queries)
     end
 
@@ -128,18 +131,29 @@
       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 = []
 
-      period, field = settings_from_table(original_table, table)
+      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
+
+      unless period
+        period, field = settings_from_table(original_table, table)
+        queries << "COMMENT ON TRIGGER #{trigger_name} ON #{table} is 'column:#{field},period:#{period}';"
+      end
       abort "Could not read settings" unless period
-      today = round_date(Date.today, 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)
 
         partition_name = "#{original_table}_#{day.strftime(name_format(period))}"
         next if table_exists?(partition_name)
+        added_partitions << partition_name
 
         queries << <<-SQL
 CREATE TABLE #{partition_name}
     (CHECK (#{field} >= #{sql_date(day)} AND #{field} < #{sql_date(advance_date(day, period, 1))}))
     INHERITS (#{table});
@@ -148,9 +162,51 @@
         queries << "ALTER TABLE #{partition_name} ADD PRIMARY KEY (#{primary_key});" if primary_key
 
         index_defs.each do |index_def|
           queries << index_def.sub(" ON #{original_table} USING ", " ON #{partition_name} USING ").sub(/ INDEX .+ ON /, " INDEX ON ") + ";"
         end
+      end
+
+      # update trigger based on existing partitions
+      current_defs = []
+      future_defs = []
+      past_defs = []
+      name_format = self.name_format(period)
+      existing_tables = self.existing_tables(like: "#{original_table}_%").select { |t| /#{Regexp.escape("#{original_table}_")}(\d{4,6})/.match(t) }
+      existing_tables = (existing_tables + added_partitions).uniq.sort
+
+      existing_tables.each do |table|
+        day = DateTime.strptime(table.split("_").last, name_format)
+        partition_name = "#{original_table}_#{day.strftime(name_format(period))}"
+
+        sql = "(NEW.#{field} >= #{sql_date(day)} AND NEW.#{field} < #{sql_date(advance_date(day, period, 1))}) THEN
+            INSERT INTO #{partition_name} VALUES (NEW.*);"
+
+        if day.to_date < today
+          past_defs << sql
+        elsif advance_date(day, period, 1) < today
+          current_defs << sql
+        else
+          future_defs << sql
+        end
+      end
+
+      # order by current period, future periods asc, past periods desc
+      trigger_defs = current_defs + future_defs + past_defs.reverse
+
+      if trigger_defs.any?
+        queries << <<-SQL
+CREATE OR REPLACE FUNCTION #{trigger_name}()
+    RETURNS trigger AS $$
+    BEGIN
+        IF #{trigger_defs.join("\n        ELSIF ")}
+        ELSE
+            RAISE EXCEPTION 'Date out of range. Ensure partitions are created.';
+        END IF;
+        RETURN NULL;
+    END;
+    $$ LANGUAGE plpgsql;
+        SQL
       end
 
       run_queries(queries) if queries.any?
     end