exwiw 1.1.4 → 1.1.6

data/lib/exwiw/query_ast.rb

@@ -26,6 +26,25 @@ module Exwiw
       end
     end
 
+    # The ID spaces whose values `query` filters by, anywhere in it.
+    def self.id_spaces(query)
+      where_clauses = query.where_clauses +
+        query.join_clauses.flat_map { |j| j.where_clauses + j.base_where_clauses }
+
+      where_clauses.flat_map do |where_clause|
+        next [] unless where_clause.is_a?(WhereClause)
+
+        own = where_clause.is_a?(ScopeWhereClause) ? [where_clause.id_space] : []
+        nested =
+          case where_clause.value
+          when SelectSubquery then id_spaces(where_clause.value.query)
+          when UnionSubquery then where_clause.value.queries.flat_map { |q| id_spaces(q) }
+          else []
+          end
+        own + nested
+      end.uniq
+    end
+
     class JoinClause
       # `where_clauses` is compiled against this join's join_table_name (the
       # joined-to table). `base_where_clauses`, on the other hand, is compiled
@@ -70,6 +89,12 @@ module Exwiw
       end
     end
 
+    # A scope-column filter. It compiles like any WhereClause; `id_space` records
+    # which ID space its values come from, so a query mixing spaces can be found.
+    ScopeWhereClause = Class.new(WhereClause) do
+      attr_accessor :id_space
+    end
+
     # Resolves a set of values on `where_column` to the rows' `select_column`
     # via a nested SELECT. Used as the `value` of a WhereClause whose operator
     # is `:in_subquery`, so a non primary-key `ids_field` can filter related

data/lib/exwiw/query_ast_builder.rb

@@ -12,25 +12,57 @@ module Exwiw
       new(table_name, table_by_name, dump_target, logger).scope_category
     end
 
-    # Scope-column mode is active when EITHER the named `--target-table` declares a
-    # per-table `scope_column` (the preferred trigger: the target is then scoped
-    # like any other table — its `--ids` are scope-column values, not primary
-    # keys), OR the deprecated `--scope-column` flag is set (a global column with no
-    # target). In both cases every table is filtered by a shared column instead of
-    # being anchored on one named target's primary key.
+    # Scope-column mode is active when the named `--target-table` declares a
+    # per-table `scope_column` (the target is then scoped like any other table,
+    # and its `--ids` are scope-column values, not primary keys); when no target is
+    # named but `--ids` are given and some table declares a `scope_column`; or when
+    # the deprecated `--scope-column` flag is set (a global column with no target).
+    # In each case every table is filtered by a shared column instead of being
+    # anchored on one named target's primary key.
     def self.scope_mode?(table_by_name, dump_target)
       return true unless dump_target.scope_column.nil?
-      return false if dump_target.table_name.nil?
+      if dump_target.table_name.nil?
+        return ids_given?(dump_target) && scope_column_declared?(table_by_name)
+      end
 
       target = table_by_name[dump_target.table_name]
       !!(target && target.respond_to?(:scope_column) && target.scope_column)
     end
 
+    def self.ids_given?(dump_target)
+      dump_target.ids.each_value.any?(&:any?)
+    end
+
+    # An ignore:true table is not extracted, so its declaration cannot anchor the
+    # scope.
+    def self.scope_column_declared?(table_by_name)
+      table_by_name.each_value.any? do |table|
+        !table.ignore && table.respond_to?(:scope_column) && !table.scope_column.nil?
+      end
+    end
+
+    # `--ids` without `--target-table` (and without the deprecated flag) selects
+    # scope-column mode, which needs a declared `scope_column` to start from.
+    def self.validate_ids_anchor!(table_by_name, dump_target)
+      return unless dump_target.table_name.nil? && dump_target.scope_column.nil?
+      return unless ids_given?(dump_target)
+      return if scope_column_declared?(table_by_name)
+
+      raise ArgumentError,
+            "--ids was given without --target-table, but no table in the schema declares " \
+            "`scope_column` (tables marked `ignore: true` are not counted). Either name the table " \
+            "the ids belong to with --target-table (single-target mode), or declare " \
+            "`scope_column: <column>` on the tables to filter by it (scope-column mode)."
+    end
+
     # Strict pre-flight: abort if any extractable table cannot be scoped (scope
     # mode), or declares a `batch_scope` its scoping shape cannot be sliced by
     # (both modes) — before any output is written. `tables` is the set of
     # dumpable configs (ignore:true tables are skipped — they are not extracted).
     def self.validate_scope!(tables, table_by_name, dump_target, logger)
+      validate_ids_anchor!(table_by_name, dump_target)
+      validate_named_id_spaces_allowed!(tables, table_by_name, dump_target)
+
       # Unscopable is reported before a bad batch_scope shape — it is the more
       # fundamental problem.
       if scope_mode?(table_by_name, dump_target)
@@ -57,6 +89,9 @@ module Exwiw
                 "add a belongs_to path to a table that carries the scope column, mark it " \
                 "`scope_exempt: true` to export it in full, or set `ignore: true` to skip it."
         end
+
+        validate_id_space_values!(tables, table_by_name, dump_target, logger)
+        validate_single_id_space_per_table!(tables, table_by_name, dump_target, logger)
       end
 
       tables.reject(&:ignore).each do |table|
@@ -66,12 +101,85 @@ module Exwiw
       end
     end
 
+    # Only scope-column mode on the SQL adapters can filter by a named ID space.
+    def self.validate_named_id_spaces_allowed!(tables, table_by_name, dump_target)
+      named = dump_target.named_id_spaces
+      return if named.empty?
+
+      given = named.map { |space| "'#{space}'" }.join(', ')
+      # The CLI rejects this up front; this covers callers that build a Runner directly.
+      if tables.any? { |table| table.is_a?(MongodbCollectionConfig) }
+        raise ArgumentError,
+              "--ids were given for the ID space(s) #{given}, but the mongodb adapter supports " \
+              "only the default ID space (--ids without a `<id_space>=` prefix)."
+      end
+
+      return if scope_mode?(table_by_name, dump_target)
+
+      raise ArgumentError,
+            "--ids were given for the ID space(s) #{given}, but --target-table " \
+            "'#{dump_target.table_name}' declares no scope_column, so the run is single-target mode, " \
+            "which uses only the default ID space. Name a table that declares a scope_column as " \
+            "--target-table, or drop the named --ids."
+    end
+
+    # Every ID space a scope_column belongs to must be given values,
+    # and every ID space given values must be used by some table.
+    def self.validate_id_space_values!(tables, table_by_name, dump_target, logger)
+      tables_by_space =
+        tables.reject(&:ignore).each_with_object({}) do |table, hash|
+          space = new(table.name, table_by_name, dump_target, logger).declared_id_space
+          (hash[space] ||= []) << table.name if space
+        end
+
+      missing = tables_by_space.keys.select { |space| dump_target.ids_for(space).empty? }.sort
+      if missing.any?
+        details = missing.map do |space|
+          flag = space == DEFAULT_ID_SPACE ? "--ids=<ids>" : "--ids=#{space}=<ids>"
+          "'#{space}' (#{tables_by_space[space].sort.join(', ')}; pass #{flag})"
+        end
+        raise ArgumentError,
+              "scope-column mode: no ids were given for the ID space(s) that these tables are " \
+              "filtered by: #{details.join('; ')}."
+      end
+
+      unused = dump_target.ids.keys.reject { |space| tables_by_space.key?(space) }.sort
+      if unused.any?
+        raise ArgumentError,
+              "scope-column mode: ids were given for the ID space(s) #{unused.map { |space| "'#{space}'" }.join(', ')}, " \
+              "but no table declares a scope_column in them (a table without `id_space` is in " \
+              "'#{DEFAULT_ID_SPACE}'). Check the ID space name, or declare `id_space` on the tables it scopes."
+      end
+    end
+
+    # A table is filtered by the values of one ID space only. Without this check
+    # the belongs_to walk would silently settle on the nearest scoped table and
+    # ignore a route into another space.
+    def self.validate_single_id_space_per_table!(tables, table_by_name, dump_target, logger)
+      return if dump_target.ids.size < 2
+
+      mixed =
+        tables.reject(&:ignore).filter_map do |table|
+          spaces = new(table.name, table_by_name, dump_target, logger).reached_id_spaces
+          "#{table.name} (#{spaces.join(', ')})" if spaces.size > 1
+        end
+      return if mixed.empty?
+
+      raise ArgumentError,
+            "scope-column mode: #{mixed.size} table(s) reach more than one ID space: #{mixed.sort.join(', ')}. " \
+            "A table must be scoped by the values of a single ID space; set `ignore: true` on the " \
+            "belongs_to (or remove the reverse_scope arm) that leads into the other space, or declare " \
+            "`scope_column:` on the table to filter it directly."
+    end
+
     attr_reader :table_name, :table_by_name, :dump_target
 
     # The multiple constrained referencers (comma-joined) that made the
     # automatic reverse detection step aside, when it did.
     attr_reader :ambiguous_referencers
 
+    attr_reader :narrowed_by_automatic_reverse
+
     def initialize(table_name, table_by_name, dump_target, logger, allow_reverse: true, allow_declared_reverse: true, forward_path: [], reverse_path: [], deep_chain_warned: nil, batch_ids: nil)
       @table_name = table_name
       @table_by_name = table_by_name
@@ -103,6 +211,7 @@ module Exwiw
       # Deep-chain warnings already emitted, keyed by chain path and shared (the
       # same Hash) across every recursive build under one top-level run.
       @deep_chain_warned = deep_chain_warned || {}
+      @narrowed_by_automatic_reverse = false
     end
 
     def run
@@ -110,8 +219,15 @@ module Exwiw
 
       return build_scoped(table) if scope_mode?
 
-      where_clauses = build_where_clauses(table, dump_target)
-      join_clauses = build_join_clauses(table, table_by_name, dump_target)
+      arm_queries = target_polymorphic_arm_queries(table)
+      if arm_queries
+        where_clauses = [pk_union_clause(table, arm_queries)]
+        where_clauses.push(table.filter) if table.filter
+        join_clauses = []
+      else
+        where_clauses = build_where_clauses(table, dump_target)
+        join_clauses = build_join_clauses(table, table_by_name, dump_target)
+      end
 
       # Reverse / "referenced_by" extraction. A table with no belongs_to path to
       # the dump target produces no where/join clauses and would otherwise dump
@@ -126,7 +242,10 @@ module Exwiw
          table.name != dump_target.table_name &&
          where_clauses.empty? && join_clauses.empty?
         reverse_clause = build_referenced_by_clause(table)
-        where_clauses.push(reverse_clause) if reverse_clause
+        if reverse_clause
+          where_clauses.push(reverse_clause)
+          @narrowed_by_automatic_reverse = !explicit_reverse_scope?(table)
+        end
       end
 
       # Forward cascade. A satellite of a reverse_scope'd (or referenced-by-scoped)
@@ -141,7 +260,7 @@ module Exwiw
         if parent_clause
           where_clauses.push(parent_clause)
         elsif @allow_reverse && @forward_path.empty? && !scope_exempt?(table) &&
-              scopable_parent_candidates(table).size > 1
+              (scopable_parent_candidates(table).size > 1 || polymorphic_parent_groups(table).size > 1)
           @logger.warn(
             "  #{table.name} belongs_to multiple scopable parents; the cascade cannot " \
             "pick one unambiguously, so it is dumped in full. If this is intended, set " \
@@ -184,63 +303,60 @@ module Exwiw
       # the path is 1, it's impossible case
       return [] if path_tables.size < 2
 
-      join_clauses = []
-
-      path_tables.each_cons(2) do |from_table_name, to_table_name|
+      path_tables.each_cons(2).map do |from_table_name, to_table_name|
         from_table = table_by_name[from_table_name]
         to_table = table_by_name[to_table_name]
+        build_join_clause(from_table, to_table, from_table.belongs_to(to_table_name))
+      end
+    end
 
-        relation = from_table.belongs_to(to_table_name)
+    private def build_join_clause(from_table, to_table, relation)
+      join_clause = QueryAst::JoinClause.new(
+        base_table_name: from_table.name,
+        foreign_key: relation.foreign_key,
+        join_table_name: to_table.name,
+        primary_key: to_table.primary_key,
+        where_clauses: [],
+        base_where_clauses: []
+      )
 
-        join_clause = QueryAst::JoinClause.new(
-          base_table_name: from_table.name,
-          foreign_key: relation.foreign_key,
-          join_table_name: to_table.name,
-          primary_key: to_table.primary_key,
-          where_clauses: [],
-          base_where_clauses: []
+      # When this hop itself is a polymorphic belongs_to (e.g. comments
+      # polymorphically belongs_to posts as commentable), the type column
+      # (foreign_type) lives on the source table (from_table = base_table_name).
+      # The foreign key alone is not enough — a value like reviewable_id=1 can
+      # collide with rows of another model — so add the type condition to
+      # base_where_clauses to narrow down the source table.
+      if relation.polymorphic?
+        join_clause.base_where_clauses.push QueryAst::WhereClause.new(
+          column_name: relation.foreign_type,
+          operator: :eq,
+          value: [relation.type_value]
         )
-
-        # When this hop itself is a polymorphic belongs_to (e.g. comments
-        # polymorphically belongs_to posts as commentable), the type column
-        # (foreign_type) lives on the source table (from_table = base_table_name).
-        # The foreign key alone is not enough — a value like reviewable_id=1 can
-        # collide with rows of another model — so add the type condition to
-        # base_where_clauses to narrow down the source table.
-        if relation.polymorphic?
-          join_clause.base_where_clauses.push QueryAst::WhereClause.new(
-            column_name: relation.foreign_type,
+      end
+      relation_to_dump_target = to_table.belongs_to(dump_target.table_name)
+      if relation_to_dump_target
+        join_clause.where_clauses.push dump_target_fk_clause(relation_to_dump_target.foreign_key)
+
+        # When the intermediate table polymorphically belongs_to the dump
+        # target, also add the type column (foreign_type) to the join
+        # condition. The type column lives on to_table (= join_table_name), so
+        # it rides on the existing mechanism where a JoinClause's where_clauses
+        # are compiled against join_table_name.
+        if relation_to_dump_target.polymorphic?
+          join_clause.where_clauses.push QueryAst::WhereClause.new(
+            column_name: relation_to_dump_target.foreign_type,
             operator: :eq,
-            value: [relation.type_value]
+            value: [relation_to_dump_target.type_value]
           )
         end
-        relation_to_dump_target = to_table.belongs_to(dump_target.table_name)
-        if relation_to_dump_target
-          join_clause.where_clauses.push dump_target_fk_clause(relation_to_dump_target.foreign_key)
-
-          # When the intermediate table polymorphically belongs_to the dump
-          # target, also add the type column (foreign_type) to the join
-          # condition. The type column lives on to_table (= join_table_name), so
-          # it rides on the existing mechanism where a JoinClause's where_clauses
-          # are compiled against join_table_name.
-          if relation_to_dump_target.polymorphic?
-            join_clause.where_clauses.push QueryAst::WhereClause.new(
-              column_name: relation_to_dump_target.foreign_type,
-              operator: :eq,
-              value: [relation_to_dump_target.type_value]
-            )
-          end
-        end
-
-        # Add filter from intermediate table to join clause
-        if to_table.filter
-          join_clause.where_clauses.push to_table.filter
-        end
+      end
 
-        join_clauses.push(join_clause)
+      # Add filter from intermediate table to join clause
+      if to_table.filter
+        join_clause.where_clauses.push to_table.filter
       end
 
-      join_clauses
+      join_clause
     end
 
     # Builds a `pk IN (SELECT child.fk FROM <child extraction query>)` clause
@@ -419,6 +535,14 @@ module Exwiw
     private def build_belongs_to_scoped_clause(table)
       candidates = scopable_parent_candidates(table)
 
+      # Plain parents take precedence over polymorphic ones: polymorphic parents
+      # are consulted only when no plain parent is scopable. Multiple independent
+      # polymorphic associations are as ambiguous as multiple plain parents.
+      if candidates.empty?
+        groups = polymorphic_parent_groups(table)
+        return groups.size == 1 ? pk_union_clause(table, groups.first) : nil
+      end
+
       # Only the unambiguous single-parent case. Multiple scopable parents would
       # need their subqueries combined (not supported); fall back to unscopable.
       if candidates.size != 1
@@ -446,6 +570,21 @@ module Exwiw
       )
     end
 
+    private def polymorphic_parent_groups(table)
+      @polymorphic_parent_groups ||= {}
+      return @polymorphic_parent_groups[table.name] if @polymorphic_parent_groups.key?(table.name)
+
+      @polymorphic_parent_groups[table.name] =
+        if table.primary_key.nil?
+          []
+        else
+          table.belongs_tos.select(&:polymorphic?).group_by(&:foreign_type).filter_map do |_, relations|
+            arm_queries = relations.filter_map { |relation| scoped_target_arm_query(table, relation) }
+            arm_queries unless arm_queries.empty?
+          end
+        end
+    end
+
     # The scopable belongs_to parents of `table`: each non-polymorphic parent
     # whose own extraction query comes out constrained, paired with the relation
     # and that query. Shared by build_belongs_to_scoped_clause (which requires
@@ -497,7 +636,7 @@ module Exwiw
         clauses.push Exwiw::QueryAst::WhereClause.new(
           column_name: dump_target.ids_field || table.primary_key,
           operator: :eq,
-          value: dump_target.ids
+          value: dump_target.default_ids
         )
 
         return clauses
@@ -539,7 +678,7 @@ module Exwiw
         return Exwiw::QueryAst::WhereClause.new(
           column_name: foreign_key,
           operator: :eq,
-          value: dump_target.ids
+          value: dump_target.default_ids
         )
       end
 
@@ -551,11 +690,58 @@ module Exwiw
           table_name: target.name,
           select_column: target.primary_key,
           where_column: dump_target.ids_field,
-          where_values: dump_target.ids
+          where_values: dump_target.default_ids
         )
       )
     end
 
+    # Single-target counterpart of #scoped_arms.
+    private def target_polymorphic_arm_queries(table)
+      return nil if table.name == dump_target.table_name
+
+      first_hop =
+        if table.belongs_to(dump_target.table_name)
+          dump_target.table_name
+        else
+          find_path_to_dump_target(table, table_by_name, dump_target)[1]
+        end
+      return nil if first_hop.nil?
+
+      sibling_arms = polymorphic_sibling_arms(table, first_hop)
+      return nil if sibling_arms.nil?
+
+      arm_queries = sibling_arms.filter_map { |relation| target_arm_query(table, relation) }
+      @logger.debug(
+        "  #{table.name} reaches the dump target through #{arm_queries.size} of " \
+        "#{sibling_arms.size} polymorphic '#{sibling_arms.first.foreign_type}' arm(s)."
+      )
+      # A lone arm is always the one the shortest path leaves through, which the
+      # single-path build already emits as a plain JOIN.
+      arm_queries.size >= 2 ? arm_queries : nil
+    end
+
+    private def target_arm_query(table, relation)
+      if relation.table_name == dump_target.table_name
+        return arm_pk_query(table).tap do |query|
+          query.where dump_target_fk_clause(relation.foreign_key)
+          query.where polymorphic_type_clause(relation)
+        end
+      end
+
+      target = table_by_name[relation.table_name]
+      return nil if target.nil?
+
+      path = find_path_to_dump_target(target, table_by_name, dump_target)
+      if path.any? && !path.include?(table.name)
+        return arm_pk_query(table).tap do |query|
+          query.join(build_join_clause(table, target, relation))
+          build_join_clauses(target, table_by_name, dump_target).each { |join_clause| query.join(join_clause) }
+        end
+      end
+
+      scoped_target_arm_query(table, relation)
+    end
+
     private def find_path_to_dump_target(table, table_by_name, dump_target)
       return [] if table.name == dump_target.table_name
 
@@ -592,7 +778,9 @@ module Exwiw
     # ------------------------------------------------------------------
 
     private def scope_mode?
-      self.class.scope_mode?(table_by_name, dump_target)
+      return @scope_mode if defined?(@scope_mode)
+
+      @scope_mode = self.class.scope_mode?(table_by_name, dump_target)
     end
 
     # Classifier used by validate_scope! and mirrored by build_scoped below.
@@ -607,6 +795,27 @@ module Exwiw
       :unscopable
     end
 
+    # The ID space this table's scope column belongs to, or nil when it has none.
+    # A scope_exempt table is dumped in full, but tables reaching it through
+    # belongs_to are still filtered by its scope column, so its ID space counts.
+    def declared_id_space
+      table = table_by_name.fetch(table_name)
+      return nil unless directly_scoped?(table)
+
+      id_space_of(table)
+    end
+
+    # The ID spaces a table that is not directly scoped can be scoped through:
+    # those of the scoped tables its belongs_to walk can stop at, and those its
+    # built query filters by (polymorphic arms, reverse_scope, referenced-by and
+    # the parent cascade).
+    def reached_id_spaces
+      table = table_by_name.fetch(table_name)
+      return [] if scope_exempt?(table) || directly_scoped?(table)
+
+      (belongs_to_reachable_id_spaces(table) | QueryAst.id_spaces(run)).sort
+    end
+
     # True when this table may still attempt the forward "scope via a scoped
     # belongs_to parent" rescue: it is not already on the forward-resolution
     # path, so descending into its parent cannot revisit it (a belongs_to cycle).
@@ -711,15 +920,43 @@ module Exwiw
       table.columns.any? { |c| c.name == column }
     end
 
+    private def id_space_of(table)
+      table.id_space || DEFAULT_ID_SPACE
+    end
+
     private def scope_where_clause(table)
       batch_clause = batch_ids_clause(table)
       return batch_clause if batch_clause
 
-      Exwiw::QueryAst::WhereClause.new(
+      Exwiw::QueryAst::ScopeWhereClause.new(
         column_name: resolved_scope_column(table),
         operator: :eq,
-        value: dump_target.ids
-      )
+        value: dump_target.ids_for(id_space_of(table))
+      ).tap { |clause| clause.id_space = id_space_of(table) }
+    end
+
+    # The ID spaces of every directly scoped table the belongs_to walk of
+    # #find_path_to_scoped could stop at, not only the nearest one.
+    private def belongs_to_reachable_id_spaces(table)
+      visited = Set[table.name]
+      queue = [table]
+      spaces = Set.new
+
+      until queue.empty?
+        current = queue.shift
+        current.belongs_tos.each do |relation|
+          next_table = table_by_name[relation.table_name]
+          next if next_table.nil? || !visited.add?(next_table.name)
+
+          if directly_scoped?(next_table)
+            spaces << id_space_of(next_table)
+          else
+            queue.push(next_table)
+          end
+        end
+      end
+
+      spaces
     end
 
     # This batch's ids, in place of the batch table's scope filter. nil when the
@@ -935,15 +1172,33 @@ module Exwiw
       path = find_path_to_scoped(table, first_relation: relation)
       return ScopedArm.new(relation: relation, path: path) if path.size >= 2
 
+      target_query = scoped_target_query(table, relation, allow_automatic_reverse: true)
+      target_query && ScopedArm.new(relation: relation, target_query: target_query)
+    end
+
+    private def scoped_target_arm_query(table, relation)
+      # A parent kept only by the automatic referenced_by exists to keep a child's
+      # foreign key valid. Single-target mode extracts only what sits under the
+      # dump target, so such a parent must not pull in the rows that point at it.
+      target_query = scoped_target_query(table, relation, allow_automatic_reverse: scope_mode?)
+      target_query && target_ids_arm_query(table, relation, target_query)
+    end
+
+    private def scoped_target_query(table, relation, allow_automatic_reverse:)
       target = table_by_name[relation.table_name]
       return nil if target.nil? || target.primary_key.nil?
       # Descending into a table already being resolved would close a cycle.
       return nil if target.name == table.name || @forward_path.include?(target.name)
 
-      target_query = self.class.run(
+      # Built the same way as the target's own extraction, so the arm never keeps
+      # a row pointing at a target row the dump leaves out.
+      builder = self.class.new(
         target.name, table_by_name, dump_target, @logger,
-        allow_reverse: true, forward_path: @forward_path + [table.name], reverse_path: @reverse_path, deep_chain_warned: @deep_chain_warned
+        forward_path: @forward_path + [table.name], reverse_path: @reverse_path, deep_chain_warned: @deep_chain_warned
       )
+      target_query = builder.run
+      return nil if !allow_automatic_reverse && builder.narrowed_by_automatic_reverse
+
       # An unconstrained target selects every id, i.e. does not scope the arm at
       # all; dropping the arm is the safe outcome.
       return nil unless target_query.where_clauses.any? || target_query.join_clauses.any?
@@ -964,7 +1219,7 @@ module Exwiw
         return nil
       end
 
-      ScopedArm.new(relation: relation, target_query: target_query)
+      target_query
     end
 
     # The polymorphic belongs_to arms sharing the `foreign_type` of the relation
@@ -1019,42 +1274,57 @@ module Exwiw
     # (`replace_with` / `raw_sql`) configured on it cannot corrupt the id
     # comparison — the same guard the reverse-scope projections use.
     private def polymorphic_arms_clause(table, arms)
-      pk_column = TableColumn.from_symbol_keys(name: table.primary_key)
-
       queries = arms.map do |arm|
-        query = QueryAst::Select.new
-        query.from(table.name)
-        query.select([pk_column])
-
         if arm.path
-          build_scoped_join_clauses(arm.path, first_relation: arm.relation).each { |jc| query.join(jc) }
+          arm_pk_query(table).tap do |query|
+            build_scoped_join_clauses(arm.path, first_relation: arm.relation).each { |jc| query.join(jc) }
+          end
         else
-          # The arm's target is scoped without a join path of its own, so probe
-          # its id set instead of joining up to a scoped ancestor. The type
-          # column still has to be constrained: the foreign key alone cannot tell
-          # this arm's rows from another arm's (record_id=1 may be any type).
-          target = table_by_name.fetch(arm.relation.table_name)
-          query.where QueryAst::WhereClause.new(
-            column_name: arm.relation.foreign_key,
-            operator: :in_subquery,
-            value: QueryAst::SelectSubquery.new(
-              query: project_query_to(arm.target_query, target.primary_key)
-            )
-          )
-          query.where QueryAst::WhereClause.new(
-            column_name: arm.relation.foreign_type,
-            operator: :eq,
-            value: [arm.relation.type_value]
-          )
+          target_ids_arm_query(table, arm.relation, arm.target_query)
         end
-
-        query
       end
 
+      pk_union_clause(table, queries)
+    end
+
+    private def pk_union_clause(table, arm_queries)
       QueryAst::WhereClause.new(
         column_name: table.primary_key,
         operator: :in_subquery,
-        value: QueryAst::UnionSubquery.new(queries: queries)
+        value: QueryAst::UnionSubquery.new(queries: arm_queries)
+      )
+    end
+
+    private def arm_pk_query(table)
+      QueryAst::Select.new.tap do |query|
+        query.from(table.name)
+        query.select([TableColumn.from_symbol_keys(name: table.primary_key)])
+      end
+    end
+
+    # The arm's target is scoped without a join path of its own, so probe its id
+    # set instead of joining up to a scoped ancestor. The type column still has
+    # to be constrained: the foreign key alone cannot tell this arm's rows from
+    # another arm's (record_id=1 may be any type).
+    private def target_ids_arm_query(table, relation, target_query)
+      target = table_by_name.fetch(relation.table_name)
+      arm_pk_query(table).tap do |query|
+        query.where QueryAst::WhereClause.new(
+          column_name: relation.foreign_key,
+          operator: :in_subquery,
+          value: QueryAst::SelectSubquery.new(
+            query: project_query_to(target_query, target.primary_key)
+          )
+        )
+        query.where polymorphic_type_clause(relation)
+      end
+    end
+
+    private def polymorphic_type_clause(relation)
+      QueryAst::WhereClause.new(
+        column_name: relation.foreign_type,
+        operator: :eq,
+        value: [relation.type_value]
       )
     end