Merge branch 'smriti-441122/preload_pats_in_expiring_worker' into 'master'

Preloading personal access tokens

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/144566



Merged-by: Imre Farkas <ifarkas@gitlab.com>
Approved-by: Imre Farkas <ifarkas@gitlab.com>
Co-authored-by: smriti <sgarg@gitlab.com>

app/models/personal_access_token.rb

@@ -42,6 +42,7 @@ class PersonalAccessToken < ApplicationRecord
   scope :owner_is_human, -> { includes(:user).references(:user).merge(User.human) }
   scope :last_used_before, -> (date) { where("last_used_at <= ?", date) }
   scope :last_used_after, -> (date) { where("last_used_at >= ?", date) }
+  scope :expiring_and_not_notified_without_impersonation, -> { where(["(revoked = false AND expire_notification_delivered = false AND expires_at >= CURRENT_DATE AND expires_at <= :date) and impersonation = false", { date: DAYS_TO_EXPIRE.days.from_now.to_date }]) }
 
   validates :scopes, presence: true
   validates :expires_at, presence: true, on: :create, unless: :allow_expires_at_to_be_empty?

app/models/user.rb

@@ -168,6 +168,8 @@ def update_tracked_fields!(request)
 
   has_many :emails
   has_many :personal_access_tokens, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent
+  has_many :expiring_soon_and_unnotified_personal_access_tokens, -> { expiring_and_not_notified_without_impersonation }, class_name: 'PersonalAccessToken'
+
   has_many :identities, dependent: :destroy, autosave: true # rubocop:disable Cop/ActiveRecordDependent
   has_many :webauthn_registrations
   has_many :chat_names, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent
@@ -615,6 +617,12 @@ def blocked?
         .where('keys.user_id = users.id')
         .expiring_soon_and_not_notified)
   end
+
+  scope :with_personal_access_tokens_expiring_soon_and_ids, ->(ids) do
+    where(id: ids)
+    .includes(:expiring_soon_and_unnotified_personal_access_tokens)
+  end
+
   scope :order_recent_sign_in, -> { reorder(arel_table[:current_sign_in_at].desc.nulls_last) }
   scope :order_oldest_sign_in, -> { reorder(arel_table[:current_sign_in_at].asc.nulls_last) }
   scope :order_recent_last_activity, -> { reorder(arel_table[:last_activity_on].desc.nulls_last, arel_table[:id].asc) }

app/workers/personal_access_tokens/expiring_worker.rb

@@ -18,21 +18,18 @@ class ExpiringWorker # rubocop:disable Scalability/IdempotentWorker
     BATCH_SIZE = 100
 
     def perform(*args)
-      limit_date = PersonalAccessToken::DAYS_TO_EXPIRE.days.from_now.to_date
-
       # rubocop: disable CodeReuse/ActiveRecord -- We need to specify batch size to avoid timing out of worker
       loop do
-        tokens = PersonalAccessToken.without_impersonation.expiring_and_not_notified(limit_date)
+        tokens = PersonalAccessToken.expiring_and_not_notified_without_impersonation
           .select(:user_id).limit(BATCH_SIZE).to_a
 
         break if tokens.empty?
 
-        users = User.where(id: tokens.pluck(:user_id).uniq)
+        users = User.with_personal_access_tokens_expiring_soon_and_ids(tokens.pluck(:user_id).uniq)
 
         users.each do |user|
           with_context(user: user) do
-            expiring_user_tokens = user.personal_access_tokens
-            .without_impersonation.expiring_and_not_notified(limit_date)
+            expiring_user_tokens = user.expiring_soon_and_unnotified_personal_access_tokens
 
             next if expiring_user_tokens.empty?
 

spec/models/personal_access_token_spec.rb

@@ -327,6 +327,28 @@
       end
     end
 
+    describe '.expiring_and_not_notified_without_impersonation' do
+      let_it_be(:expired_token) { create(:personal_access_token, expires_at: 2.days.ago) }
+      let_it_be(:revoked_token) { create(:personal_access_token, revoked: true) }
+      let_it_be(:valid_token_and_notified) { create(:personal_access_token, expires_at: 2.days.from_now, expire_notification_delivered: true) }
+      let_it_be(:valid_token) { create(:personal_access_token, expires_at: 2.days.from_now, impersonation: false) }
+      let_it_be(:long_expiry_token) { create(:personal_access_token, expires_at: described_class::MAX_PERSONAL_ACCESS_TOKEN_LIFETIME_IN_DAYS.days.from_now) }
+
+      context 'when token is there to be notified' do
+        it "has only unnotified tokens" do
+          expect(described_class.expiring_and_not_notified_without_impersonation).to contain_exactly(valid_token)
+        end
+      end
+
+      context 'when no token is there to be notified' do
+        it "return empty array" do
+          valid_token.update!(impersonation: true)
+
+          expect(described_class.expiring_and_not_notified_without_impersonation).to be_empty
+        end
+      end
+    end
+
     describe '.expired_today_and_not_notified' do
       let_it_be(:active) { create(:personal_access_token) }
       let_it_be(:expired_yesterday) { create(:personal_access_token, expires_at: Date.yesterday) }

spec/models/user_spec.rb

@@ -162,6 +162,7 @@
     it { is_expected.to have_many(:groups) }
     it { is_expected.to have_many(:keys).dependent(:destroy) }
     it { is_expected.to have_many(:expired_today_and_unnotified_keys) }
+    it { is_expected.to have_many(:expiring_soon_and_unnotified_personal_access_tokens) }
     it { is_expected.to have_many(:deploy_keys).dependent(:nullify) }
     it { is_expected.to have_many(:group_deploy_keys) }
     it { is_expected.to have_many(:events).dependent(:delete_all) }
@@ -1413,6 +1414,24 @@
       end
     end
 
+    describe '.with_personal_access_tokens_expiring_soon_and_ids' do
+      let_it_be(:user1) { create(:user) }
+      let_it_be(:user2) { create(:user) }
+      let_it_be(:pat1) { create(:personal_access_token, user: user1, expires_at: 2.days.from_now) }
+      let_it_be(:pat2) { create(:personal_access_token, user: user2, expires_at: 7.days.from_now) }
+      let_it_be(:ids) { [user1.id] }
+
+      subject(:users) { described_class.with_personal_access_tokens_expiring_soon_and_ids(ids) }
+
+      it 'filters users only by id' do
+        expect(users).to contain_exactly(user1)
+      end
+
+      it 'includes expiring personal access tokens' do
+        expect(users.first.expiring_soon_and_unnotified_personal_access_tokens).to be_loaded
+      end
+    end
+
     describe '.active_without_ghosts' do
       let_it_be(:user1) { create(:user, :external) }
       let_it_be(:user2) { create(:user, state: 'blocked') }

spec/workers/personal_access_tokens/expiring_worker_spec.rb

@@ -36,7 +36,8 @@
         create(:personal_access_token, user: user2, expires_at: 5.days.from_now)
 
         # Query count increased for the user look up
-        expect { worker.perform }.not_to exceed_all_query_limit(control).with_threshold(4)
+        # there are still 2 N+1 queries one for token name look up and another for token update.
+        expect { worker.perform }.not_to exceed_all_query_limit(control).with_threshold(2)
       end
     end