Merge branch '245323-arel-support-for-materialized-cte' into 'master'

Support AS MATERIALIZED in PG12 [RUN AS-IF-FOSS]

See merge request gitlab-org/gitlab!56976

app/services/issue_rebalancing_service.rb

@@ -62,7 +62,7 @@ class IssueRebalancingService
 
   def run_update_query(values, query_name)
     Issue.connection.exec_query(<<~SQL, query_name)
-      WITH cte(cte_id, new_pos) AS (
+      WITH cte(cte_id, new_pos) AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
        SELECT *
        FROM (VALUES #{values}) as t (id, pos)
       )

changelogs/unreleased/245323-arel-support-for-materialized-cte.yml

+---
+title: Add support for the MATERIALIZED keyword when using WITH (CTE) queries in PostgreSQL 12
+merge_request: 56976
+author:
+type: other

config/initializers/postgres_cte_as_materialized.rb

+# frozen_string_literal: true
+
+# This patch adds support for AS MATERIALIZED in Arel, see Gitlab::Database::AsWithMaterialized for more info
+module Arel
+  module Visitors
+    class Arel::Visitors::PostgreSQL
+      def visit_Gitlab_Database_AsWithMaterialized(obj, collector) # rubocop:disable Naming/MethodName
+        collector = visit obj.left, collector
+        collector << " AS#{obj.expr} "
+        visit obj.right, collector
+      end
+    end
+  end
+end

config/initializers/postgresql_cte.rb

@@ -121,6 +121,8 @@ module ActiveRecord
           end
         when Arel::Nodes::As
           with_value
+        when Gitlab::Database::AsWithMaterialized
+          with_value
         end
       end
 

db/migrate/20200609212701_add_incident_settings_to_all_existing_projects.rb

@@ -8,7 +8,7 @@ class AddIncidentSettingsToAllExistingProjects < ActiveRecord::Migration[6.0]
     # to preserve behavior for existing projects that
     # are using the create issue functionality with the default setting of true
     query = <<-SQL
-      WITH project_ids AS (
+      WITH project_ids AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported}(
         SELECT DISTINCT issues.project_id AS id
         FROM issues
         LEFT OUTER JOIN project_incident_management_settings

db/post_migrate/20191105094625_set_report_type_for_vulnerabilities.rb

@@ -7,7 +7,7 @@ class SetReportTypeForVulnerabilities < ActiveRecord::Migration[5.2]
     # set report_type based on vulnerability_occurrences from which the vulnerabilities were promoted,
     # that is, first vulnerability_occurrences among those having the same vulnerability_id
     execute <<~SQL
-    WITH first_findings_for_vulnerabilities AS (
+    WITH first_findings_for_vulnerabilities AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
       SELECT MIN(id) AS id, vulnerability_id
       FROM vulnerability_occurrences
       WHERE vulnerability_id IS NOT NULL

db/post_migrate/20191114173624_set_resolved_state_on_vulnerabilities.rb

@@ -6,7 +6,7 @@ class SetResolvedStateOnVulnerabilities < ActiveRecord::Migration[5.2]
   def up
     execute <<~SQL
     -- selecting IDs for all non-orphan Findings that either have no feedback or it's a non-dismissal feedback
-    WITH resolved_vulnerability_ids AS (
+    WITH resolved_vulnerability_ids AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
       SELECT DISTINCT vulnerability_id AS id
       FROM vulnerability_occurrences
       LEFT JOIN vulnerability_feedback ON vulnerability_feedback.project_fingerprint = ENCODE(vulnerability_occurrences.project_fingerprint::bytea, 'HEX')

db/post_migrate/20200305082754_remove_duplicate_labels_from_project.rb

@@ -55,7 +55,7 @@ class RemoveDuplicateLabelsFromProject < ActiveRecord::Migration[6.0]
     # project_id title template description type color
 
     duplicate_labels = ApplicationRecord.connection.execute(<<-SQL.squish)
-WITH data AS (
+WITH data AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
   SELECT labels.*,
   row_number() OVER (PARTITION BY labels.project_id, labels.title, labels.template, labels.description, labels.type, labels.color ORDER BY labels.id) AS row_number,
   #{CREATE} AS restore_action
@@ -83,7 +83,7 @@ WITH data AS (
     # then add `_duplicate#{ID}`
 
     soft_duplicates = ApplicationRecord.connection.execute(<<-SQL.squish)
-WITH data AS (
+WITH data AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
   SELECT
      *,
      substring(title from 1 for 245 - length(id::text)) || '_duplicate' || id::text as new_title,
@@ -108,7 +108,7 @@ WHERE labels.id IN (#{soft_duplicates.map { |dup| dup["id"] }.join(", ")});
   def restore_renamed_labels(start_id, stop_id)
     # the backup label IDs are not incremental, they are copied directly from the Labels table
     ApplicationRecord.connection.execute(<<-SQL.squish)
-WITH backups AS (
+WITH backups AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
   SELECT id, title
   FROM backup_labels
   WHERE project_id BETWEEN #{start_id} AND #{stop_id} AND

db/post_migrate/20200716234259_remove_duplicate_labels_from_group.rb

@@ -59,7 +59,7 @@ class RemoveDuplicateLabelsFromGroup < ActiveRecord::Migration[6.0]
     # group_id title template description type color
 
     duplicate_labels = ApplicationRecord.connection.execute(<<-SQL.squish)
-WITH data AS (
+WITH data AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
   SELECT labels.*,
   row_number() OVER (PARTITION BY labels.group_id, labels.title, labels.template, labels.description, labels.type, labels.color ORDER BY labels.id) AS row_number,
   #{CREATE} AS restore_action
@@ -87,7 +87,7 @@ WITH data AS (
     # then add `_duplicate#{ID}`
 
     soft_duplicates = ApplicationRecord.connection.execute(<<-SQL.squish)
-WITH data AS (
+WITH data AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
   SELECT
      *,
      substring(title from 1 for 245 - length(id::text)) || '_duplicate' || id::text as new_title,
@@ -112,7 +112,7 @@ WHERE labels.id IN (#{soft_duplicates.map { |dup| dup["id"] }.join(", ")});
   def restore_renamed_labels(start_id, stop_id)
     # the backup label IDs are not incremental, they are copied directly from the Labels table
     ApplicationRecord.connection.execute(<<-SQL.squish)
-WITH backups AS (
+WITH backups AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
   SELECT id, title
   FROM backup_labels
   WHERE id BETWEEN #{start_id} AND #{stop_id}

db/post_migrate/20200809221641_migrate_license_management_artifacts_to_license_scanning.rb

@@ -26,7 +26,7 @@ class MigrateLicenseManagementArtifactsToLicenseScanning < ActiveRecord::Migrati
       min, max = relation.pluck('MIN(job_id)', 'MAX(job_id)').flatten
 
       ActiveRecord::Base.connection.execute <<~SQL
-        WITH ci_job_artifacts_with_row_number as (
+        WITH ci_job_artifacts_with_row_number as #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
           SELECT job_id, id, ROW_NUMBER() OVER (PARTITION BY job_id ORDER BY id ASC) as row_number
           FROM ci_job_artifacts
           WHERE (file_type = #{LICENSE_SCANNING_FILE_TYPE} OR file_type = #{LICENSE_MANAGEMENT_FILE_TYPE})

db/post_migrate/20200831065705_ensure_target_project_id_is_filled.rb

@@ -32,7 +32,7 @@ class EnsureTargetProjectIdIsFilled < ActiveRecord::Migration[6.0]
       )
 
       MergeRequestMetrics.connection.execute <<-SQL
-        WITH target_project_id_and_metrics_id as (
+        WITH target_project_id_and_metrics_id as #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
           #{query_for_cte.to_sql}
         )
         UPDATE #{MergeRequestMetrics.connection.quote_table_name(MergeRequestMetrics.table_name)}

db/post_migrate/20210105030125_cleanup_projects_with_bad_has_external_wiki_data.rb

@@ -45,7 +45,7 @@ class CleanupProjectsWithBadHasExternalWikiData < ActiveRecord::Migration[6.0]
         .merge(Project.where(has_external_wiki: false).where(pending_delete: false).where(archived: false))
 
       execute(<<~SQL)
-      WITH project_ids_to_update (id) AS (
+      WITH project_ids_to_update (id) AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
         #{scope_with_projects.to_sql}
       )
       UPDATE projects SET has_external_wiki = true WHERE id IN (SELECT id FROM project_ids_to_update)
@@ -75,7 +75,7 @@ class CleanupProjectsWithBadHasExternalWikiData < ActiveRecord::Migration[6.0]
     Project.where(index_where).each_batch(of: BATCH_SIZE) do |relation|
       relation_with_exists_query = relation.where('NOT EXISTS (?)', services_sub_query)
       execute(<<~SQL)
-      WITH project_ids_to_update (id) AS (
+      WITH project_ids_to_update (id) AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
         #{relation_with_exists_query.select(:id).to_sql}
       )
       UPDATE projects SET has_external_wiki = false WHERE id IN (SELECT id FROM project_ids_to_update)

db/post_migrate/20210210221006_cleanup_projects_with_bad_has_external_issue_tracker_data.rb

@@ -44,7 +44,7 @@ class CleanupProjectsWithBadHasExternalIssueTrackerData < ActiveRecord::Migratio
         .merge(Project.where(has_external_issue_tracker: false).where(pending_delete: false))
 
       execute(<<~SQL)
-      WITH project_ids_to_update (id) AS (
+      WITH project_ids_to_update (id) AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
         #{scope_with_projects.to_sql}
       )
       UPDATE projects SET has_external_issue_tracker = true WHERE id IN (SELECT id FROM project_ids_to_update)
@@ -71,7 +71,7 @@ class CleanupProjectsWithBadHasExternalIssueTrackerData < ActiveRecord::Migratio
     Project.where(index_where).each_batch(of: BATCH_SIZE) do |relation|
       relation_with_exists_query = relation.where('NOT EXISTS (?)', services_sub_query)
       execute(<<~SQL)
-      WITH project_ids_to_update (id) AS (
+      WITH project_ids_to_update (id) AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
         #{relation_with_exists_query.select(:id).to_sql}
       )
       UPDATE projects SET has_external_issue_tracker = false WHERE id IN (SELECT id FROM project_ids_to_update)

ee/db/geo/migrate/20170906174622_remove_duplicates_from_project_registry.rb

@@ -7,7 +7,7 @@ class RemoveDuplicatesFromProjectRegistry < ActiveRecord::Migration[4.2]
 
   def up
     execute <<-SQL
-      WITH good_rows AS (
+      WITH good_rows AS #{::Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
         SELECT project_id, MAX(id) as max_id
         FROM project_registry
         GROUP BY project_id

ee/lib/ee/gitlab/background_migration/fix_orphan_promoted_issues.rb

@@ -11,14 +11,14 @@ module EE
         override :perform
         def perform(note_id)
           ActiveRecord::Base.connection.execute <<~SQL
-            WITH promotion_notes AS (
+            WITH promotion_notes AS #{::Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
               SELECT noteable_id, note as promotion_note, projects.namespace_id as epic_group_id FROM notes
               INNER JOIN projects ON notes.project_id = projects.id
               WHERE notes.noteable_type = 'Issue'
               AND notes.system IS TRUE
               AND notes.note like 'promoted to epic%'
               AND notes.id = #{Integer(note_id)}
-            ), promoted_epics AS (
+            ), promoted_epics AS #{::Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
               SELECT epics.id as promoted_epic_id, promotion_notes.noteable_id as issue_id FROM epics
               INNER JOIN promotion_notes on epics.group_id = promotion_notes.epic_group_id
               WHERE concat('promoted to epic &', epics.iid) = promotion_notes.promotion_note

lib/gitlab/background_migration/backfill_design_internal_ids.rb

@@ -97,13 +97,13 @@ module Gitlab
 
         ActiveRecord::Base.connection.execute <<~SQL
           WITH
-          starting_iids(project_id, iid) as (
+          starting_iids(project_id, iid) as #{Gitlab::Database::AsWithMaterialized.materialized_if_supported}(
             SELECT project_id, MAX(COALESCE(iid, 0))
             FROM #{table}
             WHERE project_id BETWEEN #{start_id} AND #{end_id}
             GROUP BY project_id
           ),
-          with_calculated_iid(id, iid) as (
+          with_calculated_iid(id, iid) as #{Gitlab::Database::AsWithMaterialized.materialized_if_supported}(
             SELECT design.id,
                    init.iid + ROW_NUMBER() OVER (PARTITION BY design.project_id ORDER BY design.id ASC)
               FROM #{table} as design, starting_iids as init

lib/gitlab/background_migration/backfill_project_updated_at_after_repository_storage_move.rb

@@ -8,7 +8,7 @@ module Gitlab
         updated_repository_storages = Projects::RepositoryStorageMove.select("project_id, MAX(updated_at) as updated_at").where(project_id: project_ids).group(:project_id)
 
         Project.connection.execute <<-SQL
-          WITH repository_storage_cte as (
+          WITH repository_storage_cte as #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
             #{updated_repository_storages.to_sql}
           )
           UPDATE projects

lib/gitlab/background_migration/copy_merge_request_target_project_to_merge_request_metrics.rb

@@ -8,7 +8,7 @@ module Gitlab
 
       def perform(start_id, stop_id)
         ActiveRecord::Base.connection.execute <<~SQL
-          WITH merge_requests_batch AS (
+          WITH merge_requests_batch AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
             SELECT id, target_project_id
             FROM merge_requests WHERE id BETWEEN #{Integer(start_id)} AND #{Integer(stop_id)}
           )

lib/gitlab/background_migration/fix_projects_without_project_feature.rb

@@ -22,7 +22,7 @@ module Gitlab
 
       def sql(from_id, to_id)
         <<~SQL
-          WITH created_records AS (
+          WITH created_records AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
             INSERT INTO project_features (
               project_id,
               merge_requests_access_level,

lib/gitlab/background_migration/fix_projects_without_prometheus_service.rb

@@ -136,7 +136,7 @@ module Gitlab
       # there is no uniq constraint on project_id and type pair, which prevents us from using ON CONFLICT
       def create_sql(from_id, to_id)
         <<~SQL
-          WITH created_records AS (
+          WITH created_records AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
             INSERT INTO services (project_id, #{DEFAULTS.keys.map { |key| %("#{key}")}.join(',')}, created_at, updated_at)
             #{select_insert_values_sql(from_id, to_id)}
             RETURNING *
@@ -149,7 +149,7 @@ module Gitlab
       # there is no uniq constraint on project_id and type pair, which prevents us from using ON CONFLICT
       def update_sql(from_id, to_id)
         <<~SQL
-          WITH updated_records AS (
+          WITH updated_records AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
             UPDATE services SET active = TRUE
             WHERE services.project_id BETWEEN #{Integer(from_id)} AND #{Integer(to_id)} AND services.properties = '{}' AND services.type = '#{Migratable::PrometheusService.type}'
             AND #{group_cluster_condition(from_id, to_id)} AND services.active = FALSE

lib/gitlab/background_migration/fix_user_namespace_names.rb

@@ -14,7 +14,7 @@ module Gitlab
 
       def fix_namespace_names(from_id, to_id)
         ActiveRecord::Base.connection.execute <<~UPDATE_NAMESPACES
-          WITH namespaces_to_update AS (
+          WITH namespaces_to_update AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
               SELECT
                   namespaces.id,
                   users.name AS correct_name
@@ -39,7 +39,7 @@ module Gitlab
 
       def fix_namespace_route_names(from_id, to_id)
         ActiveRecord::Base.connection.execute <<~ROUTES_UPDATE
-          WITH routes_to_update AS (
+          WITH routes_to_update AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
               SELECT
                   routes.id,
                   users.name AS correct_name

lib/gitlab/background_migration/fix_user_project_route_names.rb

@@ -8,7 +8,7 @@ module Gitlab
     class FixUserProjectRouteNames
       def perform(from_id, to_id)
         ActiveRecord::Base.connection.execute <<~ROUTES_UPDATE
-          WITH routes_to_update AS (
+          WITH routes_to_update AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
             SELECT
                 routes.id,
                 users.name || ' / ' || projects.name AS correct_name

lib/gitlab/background_migration/populate_has_vulnerabilities.rb

@@ -8,21 +8,23 @@ module Gitlab
       class ProjectSetting < ActiveRecord::Base # rubocop:disable Style/Documentation
         self.table_name = 'project_settings'
 
-        UPSERT_SQL = <<~SQL
-          WITH upsert_data (project_id, has_vulnerabilities, created_at, updated_at) AS (
-            SELECT projects.id, true, current_timestamp, current_timestamp FROM projects WHERE projects.id IN (%{project_ids})
-          )
-          INSERT INTO project_settings
-          (project_id, has_vulnerabilities, created_at, updated_at)
-          (SELECT * FROM upsert_data)
-          ON CONFLICT (project_id)
-          DO UPDATE SET
-            has_vulnerabilities = true,
-            updated_at = EXCLUDED.updated_at
-        SQL
-
         def self.upsert_for(project_ids)
-          connection.execute(UPSERT_SQL % { project_ids: project_ids.join(', ') })
+          connection.execute(upsert_sql % { project_ids: project_ids.join(', ') })
+        end
+
+        def self.upsert_sql
+          <<~SQL
+            WITH upsert_data (project_id, has_vulnerabilities, created_at, updated_at) AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
+              SELECT projects.id, true, current_timestamp, current_timestamp FROM projects WHERE projects.id IN (%{project_ids})
+            )
+            INSERT INTO project_settings
+            (project_id, has_vulnerabilities, created_at, updated_at)
+            (SELECT * FROM upsert_data)
+            ON CONFLICT (project_id)
+            DO UPDATE SET
+              has_vulnerabilities = true,
+              updated_at = EXCLUDED.updated_at
+          SQL
         end
       end
 

lib/gitlab/background_migration/wrongfully_confirmed_email_unconfirmer.rb

@@ -57,7 +57,7 @@ module Gitlab
 
       def update_email_records(start_id, stop_id)
         EmailModel.connection.execute <<-SQL
-          WITH md5_strings as (
+          WITH md5_strings as #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (
             #{email_query_for_update(start_id, stop_id).to_sql}
           )
           UPDATE #{EmailModel.connection.quote_table_name(EmailModel.table_name)}

lib/gitlab/database/as_with_materialized.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    # This class is a special Arel node which allows optionally define the `MATERIALIZED` keyword for CTE and Recursive CTE queries.
+    class AsWithMaterialized < Arel::Nodes::Binary
+      extend Gitlab::Utils::StrongMemoize
+
+      MATERIALIZED = Arel.sql(' MATERIALIZED')
+      EMPTY_STRING = Arel.sql('')
+      attr_reader :expr
+
+      def initialize(left, right, materialized: true)
+        @expr = if materialized && self.class.materialized_supported?
+                  MATERIALIZED
+                else
+                  EMPTY_STRING
+                end
+
+        super(left, right)
+      end
+
+      # Note: to be deleted after the minimum PG version is set to 12.0
+      def self.materialized_supported?
+        strong_memoize(:materialized_supported) do
+          Gitlab::Database.version.match?(/^1[2-9]\./) # version 12.x and above
+        end
+      end
+
+      # Note: to be deleted after the minimum PG version is set to 12.0
+      def self.materialized_if_supported
+        materialized_supported? ? 'MATERIALIZED' : ''
+      end
+    end
+  end
+end

lib/gitlab/database/bulk_update.rb

@@ -130,7 +130,7 @@ module Gitlab
 
         def sql
           <<~SQL
-            WITH cte(#{list_of(cte_columns)}) AS (VALUES #{list_of(values)})
+            WITH cte(#{list_of(cte_columns)}) AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (VALUES #{list_of(values)})
             UPDATE #{table_name} SET #{list_of(updates)} FROM cte WHERE cte_id = id
           SQL
         end

lib/gitlab/database/postgres_hll/batch_distinct_counter.rb

@@ -41,19 +41,6 @@ module Gitlab
         BUCKET_ID_MASK = (Buckets::TOTAL_BUCKETS - ZERO_OFFSET).to_s(2)
         BIT_31_MASK = "B'0#{'1' * 31}'"
         BIT_32_NORMALIZED_BUCKET_ID_MASK = "B'#{'0' * (32 - BUCKET_ID_MASK.size)}#{BUCKET_ID_MASK}'"
-        # @example source_query
-        #   SELECT CAST(('X' || md5(CAST(%{column} as text))) as bit(32)) attr_hash_32_bits
-        #   FROM %{relation}
-        #   WHERE %{pkey} >= %{batch_start}
-        #   AND %{pkey} < %{batch_end}
-        #   AND %{column} IS NOT NULL
-        BUCKETED_DATA_SQL = <<~SQL
-          WITH hashed_attributes AS (%{source_query})
-          SELECT (attr_hash_32_bits & #{BIT_32_NORMALIZED_BUCKET_ID_MASK})::int AS bucket_num,
-            (31 - floor(log(2, min((attr_hash_32_bits & #{BIT_31_MASK})::int))))::int as bucket_hash
-          FROM hashed_attributes
-          GROUP BY 1
-        SQL
 
         WRONG_CONFIGURATION_ERROR = Class.new(ActiveRecord::StatementInvalid)
 
@@ -103,7 +90,7 @@ module Gitlab
         def hll_buckets_for_batch(start, finish)
           @relation
             .connection
-            .execute(BUCKETED_DATA_SQL % { source_query: source_query(start, finish) })
+            .execute(bucketed_data_sql % { source_query: source_query(start, finish) })
             .map(&:values)
             .to_h
         end
@@ -139,6 +126,22 @@ module Gitlab
         def actual_finish(finish)
           finish || @relation.unscope(:group, :having).maximum(@relation.primary_key) || 0
         end
+
+        # @example source_query
+        #   SELECT CAST(('X' || md5(CAST(%{column} as text))) as bit(32)) attr_hash_32_bits
+        #   FROM %{relation}
+        #   WHERE %{pkey} >= %{batch_start}
+        #   AND %{pkey} < %{batch_end}
+        #   AND %{column} IS NOT NULL
+        def bucketed_data_sql
+          <<~SQL
+            WITH hashed_attributes AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported} (%{source_query})
+            SELECT (attr_hash_32_bits & #{BIT_32_NORMALIZED_BUCKET_ID_MASK})::int AS bucket_num,
+              (31 - floor(log(2, min((attr_hash_32_bits & #{BIT_31_MASK})::int))))::int as bucket_hash
+            FROM hashed_attributes
+            GROUP BY 1
+          SQL
+        end
       end
     end
   end

lib/gitlab/sql/cte.rb

@@ -15,20 +15,27 @@ module Gitlab
     #     Namespace
     #       with(cte.to_arel).
     #       from(cte.alias_to(ns))
+    #
+    # To skip materialization of the CTE query by passing materialized: false
+    # More context: https://www.postgresql.org/docs/12/queries-with.html
+    #
+    # cte = CTE.new(:my_cte_name, materialized: false)
+    #
     class CTE
       attr_reader :table, :query
 
       # name - The name of the CTE as a String or Symbol.
-      def initialize(name, query)
+      def initialize(name, query, materialized: true)
         @table = Arel::Table.new(name)
         @query = query
+        @materialized = materialized
       end
 
       # Returns the Arel relation for this CTE.
       def to_arel
         sql = Arel::Nodes::SqlLiteral.new("(#{query.to_sql})")
 
-        Arel::Nodes::As.new(table, sql)
+        Gitlab::Database::AsWithMaterialized.new(table, sql, materialized: @materialized)
       end
 
       # Returns an "AS" statement that aliases the CTE name as the given table

spec/lib/gitlab/sql/cte_spec.rb

@@ -41,4 +41,15 @@ RSpec.describe Gitlab::SQL::CTE do
       expect(relation.to_a).to eq(User.where(id: user.id).to_a)
     end
   end
+
+  it_behaves_like 'CTE with MATERIALIZED keyword examples' do
+    let(:expected_query_block_with_materialized) { 'WITH "some_cte" AS MATERIALIZED (' }
+    let(:expected_query_block_without_materialized) { 'WITH "some_cte" AS (' }
+
+    let(:query) do
+      cte = described_class.new(:some_cte, User.active, **options)
+
+      User.with(cte.to_arel).to_sql
+    end
+  end
 end

spec/lib/gitlab/sql/recursive_cte_spec.rb

@@ -57,4 +57,17 @@ RSpec.describe Gitlab::SQL::RecursiveCTE do
       expect(relation.to_a).to eq(User.where(id: user.id).to_a)
     end
   end
+
+  it_behaves_like 'CTE with MATERIALIZED keyword examples' do
+    # MATERIALIZED keyword is not needed for recursive queries
+    let(:expected_query_block_with_materialized) { 'WITH RECURSIVE "some_cte" AS (' }
+    let(:expected_query_block_without_materialized) { 'WITH RECURSIVE "some_cte" AS (' }
+
+    let(:query) do
+      recursive_cte = described_class.new(:some_cte)
+      recursive_cte << User.active
+
+      User.with.recursive(recursive_cte.to_arel).to_sql
+    end
+  end
 end

spec/models/project_spec.rb

@@ -4109,7 +4109,7 @@ RSpec.describe Project, factory_default: :keep do
     subject { described_class.wrap_with_cte(projects) }
 
     it 'wrapped query matches original' do
-      expect(subject.to_sql).to match(/^WITH "projects_cte" AS/)
+      expect(subject.to_sql).to match(/^WITH "projects_cte" AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported}/)
       expect(subject).to match_array(projects)
     end
   end

spec/support/shared_examples/lib/gitlab/database/cte_materialized_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'CTE with MATERIALIZED keyword examples' do
+  describe 'adding MATERIALIZE to the CTE' do
+    let(:options) { {} }
+
+    before do
+      # Clear the cached value before the test
+      Gitlab::Database::AsWithMaterialized.clear_memoization(:materialized_supported)
+    end
+
+    context 'when PG version is <12' do
+      it 'does not add MATERIALIZE keyword' do
+        allow(Gitlab::Database).to receive(:version).and_return('11.1')
+
+        expect(query).to include(expected_query_block_without_materialized)
+      end
+    end
+
+    context 'when PG version is >=12' do
+      it 'adds MATERIALIZE keyword' do
+        allow(Gitlab::Database).to receive(:version).and_return('12.1')
+
+        expect(query).to include(expected_query_block_with_materialized)
+      end
+
+      context 'when version is higher than 12' do
+        it 'adds MATERIALIZE keyword' do
+          allow(Gitlab::Database).to receive(:version).and_return('15.1')
+
+          expect(query).to include(expected_query_block_with_materialized)
+        end
+      end
+
+      context 'when materialized is disabled' do
+        let(:options) { { materialized: false } }
+
+        it 'does not add MATERIALIZE keyword' do
+          expect(query).to include(expected_query_block_without_materialized)
+        end
+      end
+    end
+  end
+end