Keep scope conditions for associations

If an association defines a scope with arity != 0
then the scope is not evaluated when preloading
that relation. This fixes it.

Changelog: fixed

gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/macro_reflection.rb

@@ -5,13 +5,38 @@ module GitlabPatches
     module Partitioning
       module Reflection
         module MacroReflection
+          NO_OWNER = Struct.new(:partition_id).new(1..100_000)
+
+          # We override the method to allow eager loading of partitioned records
+          #
+          # For eager loading the owner is always nil and we supply a benign
+          # object that allows the scope to be evaluated with a query like
+          # `where partition_id between 1 and 100000`
+          # and transforms it into
+          # `where partition_id is not null`
+          # to ensure that no partition is left out by the query.
+          # This is safe because `partition_id` columns are defined as `not null`
+          #
           def scope_for(relation, owner = nil)
             if scope.arity == 1 && owner.nil? && options.key?(:partition_foreign_key)
-              relation
+              relation = relation.instance_exec(NO_OWNER, &scope)
+              if relation_includes_partition_id_condition?(relation)
+                relation.rewhere(relation.table[:partition_id].not_eq(nil))
+              else
+                relation
+              end
             else
               super
             end
           end
+
+          def relation_includes_partition_id_condition?(relation)
+            relation
+              .where_clause
+              .extract_attributes
+              .map(&:name)
+              .include?('partition_id')
+          end
         end
       end
     end

gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/joins_spec.rb

@@ -9,7 +9,8 @@
     join_statement = <<~SQL.squish
       SELECT \"pipelines\".*
       FROM \"pipelines\"
-      INNER JOIN \"jobs\" ON \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
+      INNER JOIN \"jobs\" ON \"jobs\".\"partition_id\" IS NOT NULL
+      AND \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
       AND \"jobs\".\"partition_id\" = \"pipelines\".\"partition_id\"
       WHERE \"pipelines\".\"partition_id\" = #{pipeline.partition_id}
     SQL
@@ -25,9 +26,11 @@
     join_statement = <<~SQL.squish
       SELECT \"pipelines\".*
       FROM \"pipelines\"
-      INNER JOIN \"jobs\" ON \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
+      INNER JOIN \"jobs\" ON \"jobs\".\"partition_id\" IS NOT NULL
+      AND \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
       AND \"jobs\".\"partition_id\" = \"pipelines\".\"partition_id\"
-      INNER JOIN \"metadata\" ON \"metadata\".\"job_id\" = \"jobs\".\"id\"
+      INNER JOIN \"metadata\" ON \"metadata\".\"partition_id\" IS NOT NULL
+      AND \"metadata\".\"job_id\" = \"jobs\".\"id\"
       AND \"metadata\".\"partition_id\" = \"jobs\".\"partition_id\"
       WHERE \"pipelines\".\"partition_id\" = #{pipeline.partition_id}
     SQL

gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/preloads_spec.rb

@@ -216,9 +216,11 @@
         AS t2_r2, \"metadata\".\"test_flag\"
         AS t2_r3
         FROM \"pipelines\"
-        LEFT OUTER JOIN \"jobs\" ON \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
+        LEFT OUTER JOIN \"jobs\" ON \"jobs\".\"partition_id\" IS NOT NULL
+        AND \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
         AND \"jobs\".\"partition_id\" = \"pipelines\".\"partition_id\"
-        LEFT OUTER JOIN \"metadata\" ON \"metadata\".\"job_id\" = \"jobs\".\"id\"
+        LEFT OUTER JOIN \"metadata\" ON \"metadata\".\"partition_id\" IS NOT NULL
+        AND \"metadata\".\"job_id\" = \"jobs\".\"id\"
         AND \"metadata\".\"partition_id\" = \"jobs\".\"partition_id\"
         WHERE \"pipelines\".\"project_id\" = 1
         AND \"pipelines\".\"id\"
@@ -237,5 +239,75 @@
 
       expect(result).to include(preload_statement)
     end
+
+    it 'keeps join conditions from scope' do
+      preload_statement = <<~SQL.squish
+        SELECT \"pipelines\".\"id\"
+        AS t0_r0, \"pipelines\".\"project_id\"
+        AS t0_r1, \"pipelines\".\"partition_id\"
+        AS t0_r2, \"jobs\".\"id\"
+        AS t1_r0, \"jobs\".\"pipeline_id\"
+        AS t1_r1, \"jobs\".\"partition_id\"
+        AS t1_r2, \"jobs\".\"name\"
+        AS t1_r3, \"metadata\".\"id\"
+        AS t2_r0, \"metadata\".\"job_id\"
+        AS t2_r1, \"metadata\".\"partition_id\"
+        AS t2_r2, \"metadata\".\"test_flag\"
+        AS t2_r3
+        FROM \"pipelines\"
+        LEFT OUTER JOIN \"jobs\" ON \"jobs\".\"partition_id\" IS NOT NULL
+        AND \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
+        AND \"jobs\".\"partition_id\" = \"pipelines\".\"partition_id\"
+        LEFT OUTER JOIN \"metadata\" ON \"metadata\".\"test_flag\" = 1
+        AND \"metadata\".\"partition_id\" IS NOT NULL
+        AND \"metadata\".\"job_id\" = \"jobs\".\"id\"
+        AND \"metadata\".\"partition_id\" = \"jobs\".\"partition_id\"
+        WHERE \"pipelines\".\"project_id\" = 1
+        AND \"pipelines\".\"id\"
+        IN (#{pipeline.id}, #{other_pipeline.id})
+        AND \"pipelines\".\"partition_id\" = 100
+      SQL
+
+      result = QueryRecorder.log do
+        project
+          .pipelines
+          .includes(jobs: :test_metadata)
+          .references(:jobs, :test_metadata)
+          .where(id: [1, 2], partition_id: 100)
+          .to_a
+      end
+
+      expect(result).to include(preload_statement)
+    end
+
+    it 'does rewhere the partition_id condition when missing' do
+      preload_statement = <<~SQL.squish
+        SELECT \"pipelines\".\"id\"
+        AS t0_r0, \"pipelines\".\"project_id\"
+        AS t0_r1, \"pipelines\".\"partition_id\"
+        AS t0_r2, \"jobs\".\"id\"
+        AS t1_r0, \"jobs\".\"pipeline_id\"
+        AS t1_r1, \"jobs\".\"partition_id\"
+        AS t1_r2, \"jobs\".\"name\"
+        AS t1_r3 FROM \"pipelines\"
+        LEFT OUTER JOIN \"jobs\" ON \"jobs\".\"pipeline_id\" = NULL
+        AND \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
+        AND \"jobs\".\"partition_id\" = \"pipelines\".\"partition_id\"
+        WHERE \"pipelines\".\"project_id\" = 1
+        AND \"pipelines\".\"id\" IN (1, 2)
+        AND \"pipelines\".\"partition_id\" = 100
+      SQL
+
+      result = QueryRecorder.log do
+        project
+          .pipelines
+          .includes(:unpartitioned_jobs)
+          .references(:unpartitioned_jobs)
+          .where(id: [1, 2], partition_id: 100)
+          .to_a
+      end
+
+      expect(result).to include(preload_statement)
+    end
   end
 end

gems/activerecord-gitlab/spec/support/models.rb

@@ -22,6 +22,12 @@ class Pipeline < PartitionedRecord
     ->(pipeline) { where(partition_id: pipeline.partition_id) },
     partition_foreign_key: :partition_id,
     dependent: :destroy
+
+  has_many :unpartitioned_jobs,
+    ->(pipeline) { where(pipeline: pipeline).order(id: :desc) },
+    partition_foreign_key: :partition_id,
+    dependent: :destroy,
+    class_name: 'Job'
 end
 
 class Job < PartitionedRecord
@@ -38,6 +44,13 @@ class Job < PartitionedRecord
     inverse_of: :job,
     autosave: true
 
+  has_one :test_metadata,
+    ->(build) { where(partition_id: build.partition_id, test_flag: true) },
+    foreign_key: :job_id,
+    partition_foreign_key: :partition_id,
+    inverse_of: :job,
+    class_name: 'Metadata'
+
   accepts_nested_attributes_for :metadata
 end