Add querying and additional context to the enhanced Arel AST (#106)

* Add contextual information to transformer nodes

* Sort rubocop

* Fixed lint issues

* Fix comparison for select,update,insert and delete manager

* Added .query method to Node

* Added AddchemaToTable transformer

* Add support for schemas in function calls

* Add missing attributes to the SelectCore dot visitor

* Added TODO to Arel middleware

.rubocop.yml

@@ -3,18 +3,12 @@ AllCops:
     - 'gemfiles/arel_gems.gemfile'
     - 'gemfiles/default.gemfile'
 
-Style/FrozenStringLiteralComment:
-  Enabled: false
-
 Bundler/OrderedGems:
   Enabled: false
 
 Gemspec/OrderedDependencies:
   Enabled: false
 
-Style/MultilineBlockChain:
-  Enabled: false
-
 Metrics/LineLength:
   Enabled: true
   Max: 100
@@ -31,12 +25,21 @@ Metrics/BlockLength:
 Style/Documentation:
   Enabled: false
 
+Style/MultilineBlockChain:
+  Enabled: false
+
+Style/TrailingCommaInHashLiteral:
+  EnforcedStyleForMultiline: comma
+
 Style/TrailingCommaInArrayLiteral:
   EnforcedStyleForMultiline: comma
 
 Style/TrailingCommaInArguments:
   EnforcedStyleForMultiline: comma
 
+Style/FrozenStringLiteralComment:
+  Enabled: false
+
 Layout/MultilineMethodCallIndentation:
   Enabled: true
   EnforcedStyle: indented

lib/arel/extensions/delete_manager.rb

@@ -4,7 +4,7 @@
 module Arel
   class DeleteManager < Arel::TreeManager
     def ==(other)
-      @ast == other.ast && @ctx == other.ctx
+      other.is_a?(self.class) && @ast == other.ast && @ctx == other.ctx
     end
 
     protected

lib/arel/extensions/delete_statement.rb

@@ -3,7 +3,7 @@
 
 module Arel
   module Nodes
-    # https://www.postgresql.org/docs/9.5/sql-insert.html
+    # https://www.postgresql.org/docs/10/sql-delete.html
     class DeleteStatement
       module DeleteStatementExtension
         attr_accessor :using

lib/arel/extensions/function.rb

@@ -10,6 +10,8 @@ module FunctionExtension
         attr_accessor :filter
         attr_accessor :within_group
         attr_accessor :variardic
+        # postgres only: https://www.postgresql.org/docs/10/ddl-schemas.html
+        attr_accessor :schema_name
 
         def initialize(expr, aliaz = nil)
           super
@@ -31,6 +33,7 @@ class ToSql
       # rubocop:disable Metrics/CyclomaticComplexity
       # rubocop:disable Metrics/AbcSize
       def aggregate(name, o, collector)
+        collector << "#{o.schema_name}." if o.schema_name
         collector << "#{name}("
         collector << 'DISTINCT ' if o.distinct
         collector << 'VARIADIC ' if o.variardic

lib/arel/extensions/insert_manager.rb

@@ -4,7 +4,7 @@
 module Arel
   class InsertManager < Arel::TreeManager
     def ==(other)
-      @ast == other.ast
+      other.is_a?(self.class) && @ast == other.ast
     end
   end
 

lib/arel/extensions/select_manager.rb

@@ -4,7 +4,7 @@
 module Arel
   class SelectManager
     def ==(other)
-      @ast == other.ast && @ctx == other.ctx
+      other.is_a?(self.class) && @ast == other.ast && @ctx == other.ctx
     end
 
     protected

lib/arel/extensions/select_statement.rb

@@ -32,7 +32,10 @@ module SelectStatementExtension
         def visit_Arel_Nodes_SelectStatement(o)
           super
 
+          visit_edge o, 'lock'
+          visit_edge o, 'with'
           visit_edge o, 'union'
+          visit_edge o, 'values_lists'
         end
       end
 

lib/arel/extensions/table.rb

@@ -7,7 +7,7 @@ class Table
     module TableExtension
       # postgres only: https://www.postgresql.org/docs/9.5/sql-select.html
       attr_accessor :only
-      # postgres only: https://www.postgresql.org/docs/9.5/ddl-schemas.html
+      # postgres only: https://www.postgresql.org/docs/10/ddl-schemas.html
       attr_accessor :schema_name
       # postgres only: https://www.postgresql.org/docs/9.1/catalog-pg-class.html
       attr_accessor :relpersistence

lib/arel/extensions/update_manager.rb

@@ -4,7 +4,7 @@
 module Arel
   class UpdateManager < Arel::TreeManager
     def ==(other)
-      @ast == other.ast && @ctx == other.ctx
+      other.is_a?(self.class) && @ast == other.ast && @ctx == other.ctx
     end
 
     protected

lib/arel/sql_to_arel/pg_query_visitor.rb

@@ -427,11 +427,20 @@ def visit_FuncCall(
                  [Arel::Nodes::Overlaps.new(start1, end1, start2, end2)]
 
                else
-                 if function_names.length > 1
-                   boom "Don't know how to handle function names `#{function_names}`"
-                 end
+                 case function_names.length
+                 when 2
+                   if function_names.first == PG_CATALOG
+                     boom "Missing postgres function `#{function_names.last}`"
+                   end
 
-                 Arel::Nodes::NamedFunction.new(function_names.first, args)
+                   func = Arel::Nodes::NamedFunction.new(function_names.last, args)
+                   func.schema_name = function_names.first
+                   func
+                 when 1
+                   Arel::Nodes::NamedFunction.new(function_names.first, args)
+                 else
+                   boom "Don't know how to handle function names length `#{function_names.length}`"
+                 end
                end
 
         func.distinct = (agg_distinct.nil? ? false : true) unless func.is_a?(::Array)
@@ -540,12 +549,12 @@ def visit_LockingClause(strength:, wait_policy:)
           1 => 'FOR KEY SHARE',
           2 => 'FOR SHARE',
           3 => 'FOR NO KEY UPDATE',
-          4 => 'FOR UPDATE'
+          4 => 'FOR UPDATE',
         }.fetch(strength)
         wait_policy_clause = {
           0 => '',
           1 => ' SKIP LOCKED',
-          2 => ' NOWAIT'
+          2 => ' NOWAIT',
         }.fetch(wait_policy)
 
         Arel::Nodes::Lock.new Arel.sql("#{strength_clause}#{wait_policy_clause}")

lib/arel/sql_to_arel/pg_query_visitor/frame_options.rb

@@ -56,7 +56,7 @@ def arel(frame_options, start_offset, end_offset)
             'FRAMEOPTION_START_VALUE_PRECEDING' => 0x00400,
             'FRAMEOPTION_END_VALUE_PRECEDING' => 0x00800,
             'FRAMEOPTION_START_VALUE_FOLLOWING' => 0x01000,
-            'FRAMEOPTION_END_VALUE_FOLLOWING' => 0x02000
+            'FRAMEOPTION_END_VALUE_FOLLOWING' => 0x02000,
           }.freeze
 
           def biggest_detractable_number(number, candidates)

lib/arel/transformer.rb

@@ -1,8 +1,11 @@
 require_relative './transformer/node'
 require_relative './transformer/path'
 require_relative './transformer/path_node'
+require_relative './transformer/query'
 require_relative './transformer/visitor'
 
+require_relative './transformer/add_schema_to_table'
+
 module Arel
   module Transformer
   end

lib/arel/transformer/add_schema_to_table.rb

@@ -0,0 +1,26 @@
+module Arel
+  module Transformer
+    class AddSchemaToTable
+      attr_reader :schema_name
+
+      def initialize(schema_name)
+        @schema_name = schema_name
+      end
+
+      # https://github.com/mvgijssel/arel_toolkit/issues/110
+      def call(arel, _context)
+        tree = Arel.transformer(arel)
+
+        tree.query(
+          class: Arel::Table,
+          schema_name: nil,
+          context: { range_variable: true },
+        ).each do |node|
+          node['schema_name'].replace(schema_name)
+        end
+
+        tree.object
+      end
+    end
+  end
+end

lib/arel/transformer/context_enhancer/arel_table.rb

@@ -0,0 +1,75 @@
+module Arel
+  module Transformer
+    module ContextEnhancer
+      class ArelTable
+        # rubocop:disable Metrics/PerceivedComplexity
+        # rubocop:disable Metrics/CyclomaticComplexity
+        # rubocop:disable Metrics/AbcSize
+        def self.call(node)
+          context = node.context.merge!(range_variable: false, column_reference: false)
+          parent_object = node.parent.object
+
+          # Using Arel::Table as SELECT ... FROM <table>
+          if parent_object.is_a?(Arel::Nodes::JoinSource)
+            context[:range_variable] = true
+
+          # Using Arel::Table as SELECT ... FROM [<table>]
+          elsif parent_object.is_a?(Array) &&
+                node.parent.parent.object.is_a?(Arel::Nodes::JoinSource)
+            context[:range_variable] = true
+
+          # Using Arel::Table as SELECT ... INNER JOIN <table> ON TRUE
+          elsif parent_object.is_a?(Arel::Nodes::Join)
+            context[:range_variable] = true
+
+          # Using Arel::Table as an attribute SELECT <table>.id ...
+          elsif parent_object.is_a?(Arel::Attributes::Attribute)
+            context[:column_reference] = true
+
+          # Using Arel::Table in an INSERT INTO <table>
+          elsif parent_object.is_a?(Arel::Nodes::InsertStatement)
+            context[:range_variable] = true
+
+          # Using Arel::Table in an UPDATE <table> ...
+          elsif parent_object.is_a?(Arel::Nodes::UpdateStatement)
+            context[:range_variable] = true
+
+          # Arel::Table in UPDATE ... FROM [<table>]
+          elsif parent_object.is_a?(Array) &&
+                node.parent.parent.object.is_a?(Arel::Nodes::UpdateStatement)
+            context[:range_variable] = true
+
+          # Using Arel::Table in an DELETE FROM <table>
+          elsif parent_object.is_a?(Arel::Nodes::DeleteStatement)
+            context[:range_variable] = true
+
+          # Arel::Table in DELETE ... USING [<table>]
+          elsif parent_object.is_a?(Array) &&
+                node.parent.parent.object.is_a?(Arel::Nodes::DeleteStatement)
+            context[:range_variable] = true
+
+          # Using Arel::Table as an "alias" for WITH <table> AS (SELECT 1) SELECT 1
+          elsif parent_object.is_a?(Arel::Nodes::As) &&
+                node.parent.parent.parent.object.is_a?(Arel::Nodes::With)
+            context[:alias] = true
+
+          # Using Arel::Table as an "alias" for WITH RECURSIVE <table> AS (SELECT 1) SELECT 1
+          elsif parent_object.is_a?(Arel::Nodes::As) &&
+                node.parent.parent.parent.object.is_a?(Arel::Nodes::WithRecursive)
+            context[:alias] = true
+
+          # Using Arel::Table as an "alias" for SELECT INTO <table> ...
+          elsif parent_object.is_a?(Arel::Nodes::Into)
+            context[:alias] = true
+
+          else
+            raise "Unknown AST location for table #{node.inspect}, #{node.root_node.to_sql}"
+          end
+        end
+        # rubocop:enable Metrics/PerceivedComplexity
+        # rubocop:enable Metrics/CyclomaticComplexity
+        # rubocop:enable Metrics/AbcSize
+      end
+    end
+  end
+end

lib/arel/transformer/node.rb

@@ -7,6 +7,7 @@ class Node
       attr_reader :fields
       attr_reader :children
       attr_reader :root_node
+      attr_reader :context
 
       def initialize(object)
         @object = object
@@ -15,6 +16,7 @@ def initialize(object)
         @fields = []
         @children = {}
         @dirty = false
+        @context = {}
       end
 
       def inspect
@@ -63,7 +65,7 @@ def add(path_node, node)
       def to_sql(engine = Table.engine)
         return nil if children.empty?
 
-        target_object = object.is_a?(Arel::SelectManager) ? object.ast : object
+        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
@@ -73,6 +75,19 @@ def [](key)
         @children.fetch(key)
       end
 
+      def child_at_path(path_items)
+        selected_node = self
+        path_items.each do |path_item|
+          selected_node = selected_node[path_item]
+          return nil if selected_node.nil?
+        end
+        selected_node
+      end
+
+      def query(**kwargs)
+        Arel::Transformer::Query.call(self, kwargs)
+      end
+
       protected
 
       attr_writer :path

lib/arel/transformer/path_node.rb

@@ -16,7 +16,7 @@ def arguments?
       def inspect
         case value
         when String
-          "\"#{value}\""
+          "'#{value}'"
         else
           value.inspect
         end

lib/arel/transformer/query.rb

@@ -0,0 +1,36 @@
+module Arel
+  module Transformer
+    class Query
+      def self.call(node, kwargs)
+        node_attributes = %i[context parent]
+        node_args = kwargs.slice(*node_attributes)
+        object_args = kwargs.except(*node_attributes)
+
+        node.each.select do |child_node|
+          next unless matches?(child_node, node_args)
+
+          matches?(child_node.object, object_args)
+        end
+      end
+
+      def self.matches?(object, test)
+        case test
+        when Hash
+          case object
+          when Hash
+            test <= object
+          else
+            test.all? do |test_key, test_value|
+              next false unless object.respond_to?(test_key)
+
+              object_attribute_value = object.public_send(test_key)
+              matches? object_attribute_value, test_value
+            end
+          end
+        else
+          object == test
+        end
+      end
+    end
+  end
+end