Extend middleware to give access to the response of the SQL query

lib/arel/enhance/node.rb

@@ -59,20 +59,39 @@ def add(path_node, node)
         node.path = path.append(path_node)
         node.parent = self
         node.root_node = root_node
-        @children[path_node.value] = node
+        @children[path_node.value.to_s] = node
       end
 
       def to_sql(engine = Table.engine)
         return nil if children.empty?
 
-        target_object = object.is_a?(Arel::TreeManager) ? object.ast : object
-        collector = Arel::Collectors::SQLString.new
-        collector = engine.connection.visitor.accept target_object, collector
-        collector.value
+        if object.respond_to?(:to_sql)
+          object.to_sql(engine)
+        else
+          collector = Arel::Collectors::SQLString.new
+          collector = engine.connection.visitor.accept object, collector
+          collector.value
+        end
+      end
+
+      def to_sql_and_binds(engine = Table.engine)
+        object.to_sql_and_binds(engine)
+      end
+
+      def method_missing(name, *args, &block)
+        child = @children[name.to_s]
+        return super if child.nil?
+
+        child
+      end
+
+      def respond_to_missing?(method, include_private = false)
+        child = @children[name.to_s]
+        child.present? || super
       end
 
       def [](key)
-        @children.fetch(key)
+        @children.fetch(key.to_s)
       end
 
       def child_at_path(path_items)

lib/arel/extensions.rb

@@ -112,6 +112,7 @@
 require 'arel/extensions/active_model_attribute_with_cast_value'
 require 'arel/extensions/exists'
 require 'arel/extensions/bind_param'
+require 'arel/extensions/node'
 
 module Arel
   module Extensions

lib/arel/extensions/node.rb

@@ -0,0 +1,10 @@
+module Arel
+  module Nodes
+    class Node
+      def to_sql_and_binds(engine = Arel::Table.engine)
+        collector = engine.connection.send(:collector)
+        engine.connection.visitor.accept(self, collector).value
+      end
+    end
+  end
+end

lib/arel/extensions/tree_manager.rb

@@ -6,5 +6,10 @@ def each(&block)
 
       ::Arel::Visitors::DepthFirst.new(block).accept ast
     end
+
+    def to_sql_and_binds(engine = Arel::Table.engine)
+      collector = engine.connection.send(:collector)
+      engine.connection.visitor.accept(@ast, collector).value
+    end
   end
 end

lib/arel/middleware.rb

@@ -1,6 +1,7 @@
 require 'active_record'
 require_relative './middleware/railtie'
 require_relative './middleware/chain'
+require_relative './middleware/executor'
 require_relative './middleware/postgresql_adapter'
 
 module Arel

lib/arel/middleware/chain.rb

@@ -2,29 +2,36 @@ module Arel
   module Middleware
     class Chain
       attr_reader :executing_middleware
+      attr_reader :executor
 
       def initialize(internal_middleware = [], internal_context = {})
         @internal_middleware = internal_middleware
         @internal_context = internal_context
+        @executor = Arel::Middleware::Executor.new(internal_middleware)
         @executing_middleware = false
       end
 
-      def execute(sql, binds = [])
-        return sql if internal_middleware.length.zero?
+      def execute(sql, binds = [], &execute_sql)
+        return execute_sql.call(sql, binds) if internal_middleware.length.zero?
 
         check_middleware_recursion(sql)
-        @executing_middleware = true
 
-        result = Arel.sql_to_arel(sql, binds: binds)
         updated_context = context.merge(original_sql: sql)
-
-        internal_middleware.each do |middleware_item|
-          result = result.map do |arel|
-            middleware_item.call(arel, updated_context.dup)
-          end
+        enhanced_arel = Arel.enhance(Arel.sql_to_arel(sql, binds: binds))
+
+        result = executor.run(enhanced_arel, updated_context, execute_sql)
+
+        # TODO: pass this type in from the postgres adapter
+        case result
+        when PG::Result
+          result
+        when Array
+          result
+        else
+          raise "Datatype returned from middleware `#{result.class}` should be a SQL result"
         end
-
-        result.to_sql
+      rescue ::PgQuery::ParseError
+        execute_sql.call(sql, binds)
       ensure
         @executing_middleware = false
       end
@@ -98,21 +105,23 @@ def maybe_execute_block(new_chain, &block)
       end
 
       def check_middleware_recursion(sql)
-        return unless executing_middleware
-
-        message = <<~ERROR
-          Middleware is being called from within middleware, aborting execution
-          to prevent endless recursion. You can do the following if you want to execute SQL
-          inside middleware:
-
-            - Set middleware context before entering the middleware
-            - Use `Arel.middleware.none { ... }` to temporarily disable middleware
-
-          SQL that triggered the error:
-          #{sql}
-        ERROR
-
-        raise message
+        if executing_middleware
+          message = <<~ERROR
+            Middleware is being called from within middleware, aborting execution
+            to prevent endless recursion. You can do the following if you want to execute SQL
+            inside middleware:
+
+              - Set middleware context before entering the middleware
+              - Use `Arel.middleware.none { ... }` to temporarily disable middleware
+
+            SQL that triggered the error:
+            #{sql}
+          ERROR
+
+          raise message
+        else
+          @executing_middleware = true
+        end
       end
     end
   end

lib/arel/middleware/executor.rb

@@ -0,0 +1,61 @@
+module Arel
+  module Middleware
+    class Executor
+      attr_reader :middleware
+
+      attr_accessor :index
+      attr_reader :context
+      attr_reader :final_block
+
+      def initialize(middleware)
+        @middleware = middleware
+      end
+
+      def run(enhanced_arel, context, final_block)
+        @index = 0
+        @context = context
+        @final_block = final_block
+
+        call(enhanced_arel)
+      ensure
+        @index = 0
+        @context = nil
+        @final_block = nil
+      end
+
+      def call(next_arel)
+        check_type next_arel
+
+        current_middleware = middleware[index]
+
+        return execute_sql(next_arel) if current_middleware.nil?
+
+        self.index += 1
+
+        case current_middleware.method(:call).arity
+        when 2
+          current_middleware.call(next_arel, self)
+        else
+          current_middleware.call(next_arel, self, context.dup)
+        end
+      end
+
+      private
+
+      def execute_sql(next_arel)
+        sql, binds = next_arel.to_sql_and_binds
+        final_block.call(sql, binds)
+      end
+
+      def connection
+        Arel::Table.engine.connection
+      end
+
+      def check_type(next_arel)
+        return if next_arel.is_a?(Arel::Enhance::Node)
+
+        raise "Only `Arel::Enhance::Node` is valid for middleware, passed `#{next_arel.class}`"
+      end
+    end
+  end
+end

lib/arel/middleware/postgresql_adapter.rb

@@ -8,23 +8,27 @@ def initialize(*args)
       end
 
       def execute(sql, name = nil)
-        sql = Arel::Middleware.current_chain.execute(sql)
-        super(sql, name)
+        Arel::Middleware.current_chain.execute(sql) do |processed_sql|
+          super(processed_sql, name)
+        end
       end
 
-      def exec_no_cache(sql, name, binds)
-        sql = Arel::Middleware.current_chain.execute(sql, binds)
-        super(sql, name, binds)
+      def query(sql, name = nil)
+        Arel::Middleware.current_chain.execute(sql) do |processed_sql|
+          super(processed_sql, name)
+        end
       end
 
-      def exec_cache(sql, name, binds)
-        sql = Arel::Middleware.current_chain.execute(sql, binds)
-        super(sql, name, binds)
+      def exec_no_cache(sql, name, binds)
+        Arel::Middleware.current_chain.execute(sql, binds) do |processed_sql, processed_binds|
+          super(processed_sql, name, processed_binds)
+        end
       end
 
-      def query(sql, name = nil)
-        sql = Arel::Middleware.current_chain.execute(sql)
-        super(sql, name)
+      def exec_cache(sql, name, binds)
+        Arel::Middleware.current_chain.execute(sql, binds) do |processed_sql, processed_binds|
+          super(processed_sql, name, processed_binds)
+        end
       end
     end
   end

lib/arel/sql_to_arel/pg_query_visitor.rb

@@ -26,6 +26,9 @@ def accept(sql, binds = [])
         @sql = sql
 
         Result.new visit(object, :top)
+      rescue ::PgQuery::ParseError => e
+        new_error = ::PgQuery::ParseError.new(e.message, __FILE__, __LINE__, -1)
+        raise new_error, e.message, e.backtrace
       rescue ::StandardError => e
         raise e.class, e.message, e.backtrace if e.is_a?(Arel::SqlToArel::Error)
 
@@ -592,7 +595,7 @@ def visit_OnConflictClause(action:, infer: nil, target_list: nil, where_clause:
         conflict
       end
 
-      def visit_ParamRef(number:)
+      def visit_ParamRef(number: nil)
         value = (binds[number - 1] unless binds.empty?)
 
         Arel::Nodes::BindParam.new(value)

lib/arel/sql_to_arel/result.rb

@@ -1,8 +1,25 @@
 module Arel
   module SqlToArel
     class Result < Array
-      def to_sql
-        map(&:to_sql).join('; ')
+      def to_sql(engine = Arel::Table.engine)
+        sql, _binds = to_sql_and_binds(engine)
+        sql
+      end
+
+      def to_sql_and_binds(engine = Arel::Table.engine)
+        sql_collection = []
+        binds_collection = []
+
+        each do |item|
+          sql, binds = item.to_sql_and_binds(engine)
+          sql_collection << sql
+          binds_collection.concat(binds)
+        end
+
+        [
+          sql_collection.join('; '),
+          binds_collection,
+        ]
       end
 
       def map(&block)

lib/arel/transformer/prefix_schema_name.rb

@@ -4,27 +4,24 @@ class PrefixSchemaName
       PG_CATALOG = 'pg_catalog'.freeze
       DEFAULT_SCHEMA_PRIORITY = ['public', PG_CATALOG].freeze
 
-      attr_reader :connection
       attr_reader :object_mapping
       attr_reader :schema_priority
 
       def initialize(
-        connection,
         schema_priority = DEFAULT_SCHEMA_PRIORITY,
         override_object_mapping = {}
       )
-        @connection = connection
         @schema_priority = schema_priority
         @object_mapping = database_object_mapping.merge(override_object_mapping)
       end
 
-      # https://github.com/mvgijssel/arel_toolkit/issues/110
-      def call(arel, _context)
+      def call(arel, next_middleware)
         tree = Arel.enhance(arel)
         update_arel_tables(tree)
         update_typecasts(tree)
         update_functions(tree)
-        tree.object
+
+        next_middleware.call tree.object
       end
 
       private
@@ -177,6 +174,10 @@ def database_functions
           'FROM pg_proc INNER JOIN pg_namespace ON pg_proc.pronamespace = pg_namespace.oid',
         )
       end
+
+      def connection
+        Arel::Table.engine.connection
+      end
     end
   end
 end

lib/arel/transformer/remove_active_record_info.rb

@@ -2,7 +2,7 @@ module Arel
   module Transformer
     class RemoveActiveRecordInfo
       class << self
-        def call(arel, _context)
+        def call(arel, next_middleware)
           tree = Arel.enhance(arel)
 
           tree.query(class: Arel::Table).each do |node|
@@ -15,7 +15,7 @@ def call(arel, _context)
             )
           end
 
-          tree.object
+          next_middleware.call tree.object
         end
 
         private