Add querying and additional context to the enhanced Arel AST

.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_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/sql_to_arel/pg_query_visitor.rb

@@ -540,12 +540,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/context_enhancer/arel_table.rb

@@ -0,0 +1,64 @@
+module Arel
+  module Transformer
+    module ContextEnhancer
+      class ArelTable
+        # Context we can add:
+        # - this attribute is using in a projection, group, order, where
+        # - this node belongs to this "group", like SELECT statement class
+        #   - helps for subscriptions: "How are table A and B related?" -> "Give me references to table A" -> "Are there references here which point to B?"
+        # - this Arel::Table node is a way of brining more data into the SQL query (FROM/JOIN) versus something else
+        # - this Arel::Attributes::Attribute belongs to this reference (allows to re-link unqualified columns)
+        # -
+        def self.call(node)
+          context = node.context.merge!(range_variable: false, column_reference: false)
+
+          # Using Arel::Table as SELECT ... FROM <table>
+          if node.parent.object.is_a?(Arel::Nodes::JoinSource)
+            context[:range_variable] = true
+          # Using Arel::Table as SELECT ... FROM [<table>]
+          elsif node.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 node.parent.object.is_a?(Arel::Nodes::Join)
+            context[:range_variable] = true
+
+          # Using Arel::Table as an attribute SELECT <table>.id ...
+          elsif node.parent.object.is_a?(Arel::Attributes::Attribute)
+            context[:column_reference] = true
+
+          # Using Arel::Table in an INSERT INTO <table>
+          elsif node.parent.object.is_a?(Arel::Nodes::InsertStatement)
+            context[:range_variable] = true
+
+          # Using Arel::Table in an UPDATE <table> ...
+          elsif node.parent.object.is_a?(Arel::Nodes::UpdateStatement)
+            context[:range_variable] = true
+          elsif node.parent.object.is_a?(Array) && node.parent.parent.object.is_a?(Arel::Nodes::UpdateStatement)
+            # Arel::Table in UPDATE ... FROM [<table>]
+            context[:range_variable] = true
+
+          # Using Arel::Table in an DELETE FROM <table>
+          elsif node.parent.object.is_a?(Arel::Nodes::DeleteStatement)
+            context[:range_variable] = true
+          elsif node.parent.object.is_a?(Array) && node.parent.parent.object.is_a?(Arel::Nodes::DeleteStatement)
+            # Arel::Table in DELETE ... USING [<table>]
+            context[:range_variable] = true
+
+          # Using Arel::Table as an "alias" for WITH <table> AS (SELECT 1) SELECT 1
+          elsif node.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 node.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 node.parent.object.is_a?(Arel::Nodes::Into)
+            context[:alias] = true
+
+          else
+            raise "Unknown AST location for table #{node.inspect}, #{sql}"
+          end
+        end
+      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,15 @@ 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
+
       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/visitor.rb

@@ -1,13 +1,25 @@
+require_relative './context_enhancer/arel_table'
+
 module Arel
   module Transformer
     # rubocop:disable Naming/MethodName
     class Visitor < Arel::Visitors::Dot
-      def accept(object)
+      DEFAULT_CONTEXT_ENHANCERS = {
+        Arel::Table => Arel::Transformer::ContextEnhancer::ArelTable,
+      }.freeze
+
+      attr_reader :context_enhancers
+
+      def accept(object, context_enhancers = DEFAULT_CONTEXT_ENHANCERS)
+        @context_enhancers = context_enhancers
+
         root_node = Arel::Transformer::Node.new(object)
         accept_with_root(object, root_node)
       end
 
-      def accept_with_root(object, root_node)
+      def accept_with_root(object, root_node, context_enhancers = DEFAULT_CONTEXT_ENHANCERS)
+        @context_enhancers = context_enhancers
+
         with_node(root_node) do
           visit object
         end
@@ -42,6 +54,8 @@ def process_node(arel_node, path_node)
         node = Arel::Transformer::Node.new(arel_node)
         current_node.add(path_node, node)
 
+        update_context(node)
+
         with_node node do
           visit arel_node
         end
@@ -54,6 +68,13 @@ def visit(object)
       def current_node
         @node_stack.last
       end
+
+      def update_context(node)
+        enhancer = context_enhancers[node.object.class]
+        return if enhancer.nil?
+
+        enhancer.call(node)
+      end
     end
     # rubocop:enable Naming/MethodName
   end

spec/arel/sql_to_arel_spec.rb

@@ -215,6 +215,8 @@
         pg_node: 'PgQuery::RANGE_SUBSELECT'
   visit 'select', '1 FROM "public"."table_is_range_var" "alias", ONLY "b"',
         pg_node: 'PgQuery::RANGE_VAR'
+  # TODO: this one breaks
+  # visit 'select', '1 FROM "b", (SELECT 1) "a", "c"', pg_node: 'PgQuery::RANGE_VAR'
   visit 'select', '1', pg_node: 'PgQuery::RAW_STMT'
   visit 'sql', 'REFRESH MATERIALIZED VIEW view WITH NO DATA',
         pg_node: 'PgQuery::REFRESH_MAT_VIEW_STMT',

spec/arel/transformer_spec.rb

@@ -154,4 +154,16 @@
 
     expect(original_arel).to be_identical_arel(new_arel)
   end
+
+  it 'adds additional context to Arel::Table nodes' do
+    result = Arel.sql_to_arel('SELECT posts.id FROM posts')
+    original_arel = result.first
+
+    tree = Arel.transformer(original_arel)
+    from_table_node = tree.child_at_path(['ast', 'cores', 0, 'source', 'left'])
+    projection_table_node = tree.child_at_path(['ast', 'cores', 0, 'projections', 0, 'relation'])
+
+    expect(from_table_node.context).to eq(range_variable: true, column_reference: false)
+    expect(projection_table_node.context).to eq(range_variable: false, column_reference: true)
+  end
 end