From e97c23d486f3442a4b1d200ccf8de19e7303e78f Mon Sep 17 00:00:00 2001
From: Ben Sheldon <bensheldon@gmail.com>
Date: Tue, 15 Dec 2020 19:10:13 -0800
Subject: [PATCH] Ensure advisory lock CTE is MATERIALIZED on Postgres v12+

---
 .github/workflows/test.yml         |  4 +++-
 lib/good_job/lockable.rb           | 19 ++++++++++++++++---
 spec/lib/good_job/lockable_spec.rb |  2 +-
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 512beac90..9bc6d1a2b 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -53,6 +53,7 @@ jobs:
     strategy:
       matrix:
         ruby: [2.5, 2.6, 2.7]
+        pg: [12.5, 10.8]
     env:
       PGHOST: localhost
       PGUSER: test_app
@@ -63,11 +64,12 @@ jobs:
       DISABLE_SPRING: 1
     services:
       postgres:
-        image: postgres:10.8
+        image: postgres:${{ matrix.pg }}
         env:
           POSTGRES_USER: test_app
           POSTGRES_DB: test_app_test
           POSTGRES_PASSWORD: ""
+          POSTGRES_HOST_AUTH_METHOD: trust
         ports: ["5432:5432"]
         options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
 
diff --git a/lib/good_job/lockable.rb b/lib/good_job/lockable.rb
index b054bdb29..20fc22691 100644
--- a/lib/good_job/lockable.rb
+++ b/lib/good_job/lockable.rb
@@ -32,11 +32,18 @@ module Lockable
         original_query = self
 
         cte_table = Arel::Table.new(:rows)
-        composed_cte = Arel::Nodes::As.new(cte_table, original_query.select(primary_key).except(:limit).arel)
+        cte_query = original_query.select(primary_key).except(:limit)
+        cte_type = if supports_cte_materialization_specifiers?
+                     'MATERIALIZED'
+                   else
+                     ''
+                   end
+
+        composed_cte = Arel::Nodes::As.new(cte_table, Arel::Nodes::SqlLiteral.new([cte_type, "(", cte_query.to_sql, ")"].join(' ')))
 
         query = cte_table.project(cte_table[:id])
-                  .with(composed_cte)
-                  .where(Arel.sql(sanitize_sql_for_conditions(["pg_try_advisory_lock(('x' || substr(md5(:table_name || #{connection.quote_table_name(cte_table.name)}.#{quoted_primary_key}::text), 1, 16))::bit(64)::bigint)", { table_name: table_name }])))
+                         .with(composed_cte)
+                         .where(Arel.sql(sanitize_sql_for_conditions(["pg_try_advisory_lock(('x' || substr(md5(:table_name || #{connection.quote_table_name(cte_table.name)}.#{quoted_primary_key}::text), 1, 16))::bit(64)::bigint)", { table_name: table_name }])))
 
         limit = original_query.arel.ast.limit
         query.limit = limit.value if limit.present?
@@ -132,6 +139,12 @@ def with_advisory_lock
           records.each(&:advisory_unlock)
         end
       end
+
+      def supports_cte_materialization_specifiers?
+        return @supports_cte_materialization_specifiers if defined?(@supports_cte_materialization_specifiers)
+
+        @supports_cte_materialization_specifiers = ActiveRecord::Base.connection.postgresql_version >= 120000
+      end
     end
 
     # Acquires an advisory lock on this record if it is not already locked by
diff --git a/spec/lib/good_job/lockable_spec.rb b/spec/lib/good_job/lockable_spec.rb
index bdf4b4967..7fb232be1 100644
--- a/spec/lib/good_job/lockable_spec.rb
+++ b/spec/lib/good_job/lockable_spec.rb
@@ -25,7 +25,7 @@
         SELECT "good_jobs".*
         FROM "good_jobs"
         WHERE "good_jobs"."id" IN (
-          WITH "rows" AS (
+          WITH "rows" AS #{'MATERIALIZED' if model_class.supports_cte_materialization_specifiers?} (
             SELECT "good_jobs"."id"
             FROM "good_jobs"
             WHERE "good_jobs"."priority" = 99