triggers.rb

require "erb"

module PostgreSQLTriggers
  def create_trigger(name, **options, &block)
    builder = TriggerBuilder.new(name, **options &block)
    builder.each_trigger do |trigger|
      reversible do |dir|
        dir.up do
          say "Create trigger #{trigger.name} on #{trigger.table}"
          trigger.up(self)
        end
        dir.down do
          say "Drop trigger #{builder.name} on #{trigger.table}"
          builder.down(self)
        end
      end
    end
  end

  def execute(sql, name = nil)
    message = sql.count("\n") > 0 ? "execute(<<~SQL)\e[1m\n#{sql}\e[0mSQL\n" : "execute(\e[1m#{sql}\e[0m)"
    say_with_time(message) do
      connection.execute(sql, name)
    end
  end

  class TriggerBuilder
    attr_reader :name, :options

    def initialize(name, **options, &block)
      @name    = name
      @options = options
      @groups  = {}
      yield(self) if block_given?
    end

    def each_trigger
      unless block_given?
        return to_enum(__method__) { @groups.size }
      end
      @groups.each_value { |k| yield(k) }
    end

    def define_trigger(*args, &block)
      store_trigger_in_group(__callee__, *args, &block)
    end

    alias :before     :define_trigger
    alias :after      :define_trigger
    alias :instead_of :define_trigger

    private :define_trigger

    def method_missing(meth, *args, &blk)
      if /^(?<cb>before|after|instead_of)_(?<event>delete|update|create|truncate)$/ =~ meth.to_s
        store_trigger_in_group(cb.to_sym, event.to_sym, *args, &blk)
      else
        super
      end
    end

    def respond_to_missing?(meth, include_private = false)
      /^(?:before|after|instead_of)_(?:delete|update|create|truncate)$/ === meth.to_s || super
    end

    private

    def store_trigger_in_group(callback, *args, on:, **options, &block)
      key = [on, callback]
      @groups[key] ||= TriggerObject.new(on, @name, callback, **options)
      @groups[key].ops.public_send(*args, &block)
    end
  end

  class TriggerObject
    WHENS = %i(before after instead_of).freeze
    NON_EMBED_METHODS = %w(up down call).freeze
    UnknownTriggerWhen = Class.new(ArgumentError)

    private_constat :WHENS, :NON_EMBED_METHODS

    attr_reader :name, :table, :ops, :extras
    delegate :events, to: :@ops

    def initialize(table, name, event, **options)
      unless WHENS.include?(event)
        raise UnknownTriggerWhen.new("#{event} is not present in allowed when condition: #{WHENS}")
      end

      @table, @name, @when, @ops = table, name, event, Ops.new
      @for_each = options.delete(:for_each).presence_in([:row, :statement]) || :row
      @extras = options.each_with_object("") do |key, bodies, extras|
        Array.wrap(bodies).each { |body| extras << "\nWHEN (#{body})" if key == :when }
      end
    end

    def up(adapter)
      adapter.execute([create_function_sql, create_trigger_sql].join("\n"))
    end

    def down(adapter)
      adapter.execute("DROP TRIGGER IF EXISTS #@name ON #@table")
      adapter.execute("DROP FUNCTION IF EXISTS #{function_name}")
    end

    def function_name
      @function_name ||= "process_#{@name}_on_#{@table}()"
    end

    def when_op
      @when.to_s.upcase.tr("_", " ")
    end

    private

    def create_function_sql
      inject_into_sql(<<~SQL)
      CREATE OR REPLACE FUNCTION __function_name__
      RETURNS TRIGGER AS $$
        BEGIN
          __ops__
          --
          RETURN NULL;
        END;
      $$ LANGUAGE plpgsql;
      SQL
    end

    def create_trigger_sql
      inject_into_sql(<<~SQL)
      CREATE TRIGGER __name__ #{when_op} #{events} ON #@table
        FOR EACH #{@for_each.upcase} #@extras
        EXECUTE PROCEDURE #{function_name};
      SQL
    end

    def inject_into_sql(sql)
      sql.gsub(/__([A-Za-z0-9][A-Za-z0-9_]+[A-Za-z0-9])__/) do |m|
        !NON_EMBED_METHODS.include?($1) ? public_send($1).to_s : m
      end
    end

    class Ops
      NAMES = %i(insert delete update).freeze
      private_constant :NAMES

      def initialize
        @ops = {}
      end

      def to_s
        @ops.each_value.map { |op| op.render(@ops.size == op.pos) } * ?\n
      end

      def events
        @ops.each_value.map(&:to_event) * " OR "
      end

      private

      def method_missing(meth, *args, &blk)
        NAMES.include?(meth) ? add_op(meth, *args, &blk) : super
      end

      def add_op(key, body = nil, options = nil) # :yields:
        body, options  = yield, body if block_given?
        @ops[key] = Body.new(@ops.size.next, key, body, options)
      end

      Body = Struct.new(:pos, :key, :body, :options) do
        attr_accessor :is_last
        erb = ERB.new(<<~ERB, nil, "%<>-")
          <%= pos == 1 ? :IF : :ELSIF %> (TG_OP = '<%= tg_op.upcase %>') THEN
          % body.each_line do |line|
            <%= line %>
          % end; if is_last
          END IF;
          % end
        ERB
        erb.location = __FILE__, __LINE__ + 1
        def_erb_method 'render(is_last)', erb

        def to_event
          case key
          when :update; [key.upcase, Array.wrap(of) * ', '].compact * ' OF '
          else          key.upcase
          end
        end
      end
    end
  end
end