Rack middleware to (theoretically) "isolate" certain requests in a read-only transaction

pg_isolation.rb

# "Isolates" Rack operations via Sequel/PostgreSQL into read-only transactions
#
# @note (see PGIsolation#shadow_tables!)
class PGIsolation
    attr_reader :app, :db_method, :shadow_tables, :setup, :teardown,
                :isolation_check

    # A template for creating multiple temporary tables in sequence
    #
    # Due to the +LIKE+ clause (and <tt>INCLUDING DEFAULTS</tt>), these tables
    # will be have *almost* identically to the tables they're replacing. There
    # is an additional caveat on this, see the {#call} method for more.
    SHADOW = 'CREATE TEMPORARY TABLE %s (LIKE %s INCLUDING DEFAULTS)'

    # Destroy all temporary tables.
    #
    # @note (see #teardown_isolation!)
    REVEAL = 'DISCARD TEMPORARY'

    # Create a new instance of +PGIsolation+
    #
    # This method is usually not used directly, even via +new+, but rather by
    # the +use+ method, as other rack middleware:
    #
    #     use PGIsolation { current_user.read_only? }
    #
    # @param [Object] app A Rack application to wrap
    # @param [Hash] options Optional modification parameters
    # @option options [Symbol] :db_method (:db) The method that the +app+
    #   instance exposes to access the application's database connection.
    # @option options [Array] :shadow_tables ([]) Tables which should have
    #   "shadow" versions created which can be written to, but their values
    #   eventually discarded
    # @option options [Array] :setup ([]) Raw SQL statements which should be
    #   run prior to beginning the transaction. See {#setup_isolation!} for
    #   more information on potential uses for this parameter.
    # @option options [Array] :teardown ([]) Raw SQL statements which should
    #   be run after the end of the transaction. See {#teardown_isolation!}
    #   for more information on potential uses for this parameter.
    # @yield A check to see whether a given call should be isolated. This is
    #   not invoked until the {#call} method, and is run in the context of the
    #   +app+ parameter (<i>i.e.</i>, you should have access to all helpers,
    #   models, etc.
    def initialize(app, options = {}, &isolation_check)
        @app = app

        @db_method     = options.fetch(:db_method,     :db)
        @shadow_tables = options.fetch(:shadow_tables, [])
        @setup         = options.fetch(:setup,         [])
        @teardown      = options.fetch(:teardown,      [])

        @isolation_check = isolation_check
    end

    # Fetch the Database as exposed by the application class
    #
    # This should be the end result of calling the method specified by
    # +db_method+ in the optional parameters for the {#initialize} method.
    # The default value is simply +:db+.
    #
    # @return [Sequel::Dataset] The database connection that should be
    #   isolated
    def db
        app.send(db_method)
    end

    # Process a single request (Middleware entry point)
    #
    # This is where the proper middleware behavior begins; every request is
    # wrapped in this method, and if the +isolation_check+ block returns
    # +true+, then the entire call will be wrapped in a <tt>READ ONLY</tt>
    # transaction. Additionally, any tables declared in the +shadow_tables+
    # option will have a <tt>TEMPORARY TABLE</tt> created in its place, so
    # that any real data in the table is masked, and any writes will process
    # like normal, but these values will be discarded after the request
    # returns.
    #
    # It is worth noting that any <tt>SERIAL</tt> types declared (<i>i.e.</i>,
    # auto-incrementing values) must currently be separately handled; the
    # default values are copied over, but <tt>READ ONLY</tt> transactions
    # cannot properly use them. You must pass any additional modifications to
    # the shadowed table to prevent errors, such as if you're expecting
    # exactly one write to an +users+ table:
    #
    #     shadow_tables = %w(users)
    #     setup         = ['ALTER TABLE users ALTER COLUMN id SET DEFAULT 1']
    #     use PGIsolation, {shadow_tables: shadow_tables, setup: setup} do
    #         should_be_isolated?
    #     end
    #
    # @note (see #teardown_isolation!)
    def call(env)
        response = nil

        if app.instance_eval(&isolation_check)
            setup_isolation!

            db.transaction(read_only: true, auto_savepoint: true) do
                db.after_rollback { teardown_isolation! }
                response = app.call(env)
                raise Sequel::Rollback
            end
        else
            response = app.call(env)
        end

        response
    end

    # Hide all real tables behind temporary tables
    #
    # Given an +Array+ of table names in the +shadow_tables+ option, any time
    # the +isolation_check+ option returns true, this method will "hide" all
    # those tables behind a <tt>TEMPORARY TABLE</tt> declaration.
    #
    # @note There is only one layer of "shadowing" possible; that is, you
    #   cannot hide tables which are already +TEMPORARY+.
    # @note (see #teardown_isolation!)
    def shadow_tables!
        quote_proxy = Sequel::Dataset.new(db)

        shadow_tables.each do |table|
            identifier = quote_proxy.quote_identifier_append('', table)
            db << (SHADOW % [identifier, identifier])
        end
    end

    # Run all additional isolation steps
    #
    # Assuming the +isolation_check+ passes, any raw SQL statements passed in
    # via the +setup+ arguments will be run in sequence here, after the tables
    # are shadowed, but before the transaction begins and the connection
    # becomes <tt>READ ONLY</tt>. In particular, this can be used to replace
    # +DEFAULT+ clauses which would otherwise call a function which would
    # qualify as a write (such as sequences):
    #
    #     setup = ['ALTER TABLE shadowed ALTER id SET DEFAULT 1']
    #     use PGIsolation, setup: setup { check_for_isolation }
    def setup_isolation!
        shadow_tables!
        setup.each { |query| db << query }
    end

    # Run any teardown steps, including discarding all temporary tables
    #
    # By default, this function will simply run the query in
    # {PGIsolation::REVEAL}, but if +teardown+ was given as an option during
    # the middleware creation, then all these statements will be run; this
    # will happen after the transaction closes, but before the temporary
    # tables are discarded. This can be used, for instance, to track what the
    # read-only user *attempted* to do:
    #
    #     shadow:   %w(users)
    #     teardown: ['INSERT INTO ro_user_audits SELECT * FROM users']
    #     use PGIsolation, shadow: shadow, teardown: teardown { user.ro }
    #
    # @note To discard all temporary tables safely without a call to +DROP+
    #   which could theoretically destroy the real table in certain edge
    #   cases, this extension calls <tt>DISCARD TEMPORARY</tt>, which also
    #   destroys any other +TEMPORARY+ tables, even those which are not
    #   created for use with this extension.
    def teardown_isolation!
        teardown.each { |query| db << query }
        db << REVEAL
    end
end