Add a Rake task to edit/remove expiration dates

Running `rake gitlab:tokens:edit` will provide an interface
to extend/remove tokens.

Gemfile

@@ -370,6 +370,8 @@ gem 'gettext', '~> 3.3', require: false, group: :development, feature_category:
 gem 'batch-loader', '~> 2.0.1' # rubocop:todo Gemfile/MissingFeatureCategory
+gem 'tty-prompt', '~> 0.23', require: false, feature_category: :shared
+
 # Perf bar
 gem 'peek', '~> 1.1' # rubocop:todo Gemfile/MissingFeatureCategory

Gemfile.lock

@@ -2071,6 +2071,7 @@ DEPENDENCIES
   timfel-krb5-auth (~> 0.8)
   toml-rb (~> 2.2.0)
   truncato (~> 0.7.12)
+  tty-prompt (~> 0.23)
   typhoeus (~> 1.4.0)
   undercover (~> 0.4.4)
   unleash (~> 3.2.2)

lib/tasks/gitlab/tokens.rake

 # frozen_string_literal: true
-module Tasks
-  module Gitlab
-    module Tokens
-      TOTAL_WIDTH = 70
-
-      class << self
-        def analyze
-          show_pat_expires_at_migration_status
-          show_most_common_pat_expiration_dates
-        end
-
-        def show_pat_expires_at_migration_status
-          sql = <<~SQL
-            SELECT * FROM batched_background_migrations
-            WHERE job_class_name = 'CleanupPersonalAccessTokensWithNilExpiresAt'
-          SQL
-
-          print_header("Personal Access Token Migration Status")
-
-          record = ApplicationRecord.connection.select_one(sql)
-
-          if record
-            puts "Started at: #{record['started_at']}"
-            puts "Finished  : #{record['finished_at']}"
-          else
-            puts "Status: Not run"
-          end
-
-          print_footer
-        end
-
-        def show_most_common_pat_expiration_dates
-          print_header "Top 10 Personal Access Token Expiration Dates"
-
-          puts "| Expiration Date | Count |"
-          puts "|-----------------|-------|"
-
-          PersonalAccessToken
-            .select(:expires_at, Arel.sql('count(*)'))
-            .where('expires_at >= NOW()')
-            .group(:expires_at)
-            .order(Arel.sql('count(*) DESC'))
-            .order(expires_at: :desc)
-            .limit(10)
-            .each do |row|
-            puts "| #{row[:expires_at].to_s.ljust(15)} | #{row[:count].to_s.ljust(5)} |"
-          end
-
-          print_footer
-        end
-
-        def print_header(title, total_width = TOTAL_WIDTH)
-          title_length = title.length
-          side_length = (total_width - title_length) / 2
-
-          left_side = "=" * side_length
-          right_side = "=" * (side_length + (total_width % 2))
-
-          header = "#{left_side} #{title} #{right_side}"
-          puts header
-        end
-
-        def print_footer
-          # Account for the spaces between the "=" in the header
-          puts "=" * (TOTAL_WIDTH + 2)
-        end
-      end
-    end
-  end
-end
+require_relative 'tokens/manage_expiry_task'
 namespace :gitlab do
   namespace :tokens do
     desc 'GitLab | Tokens | Show information about tokens'
     task analyze: :environment do |_t, _args|
-      Tasks::Gitlab::Tokens.analyze
+      Tasks::Gitlab::Tokens::ManageExpiryTask.new.analyze
+    end
+
+    desc 'GitLab | Tokens | Edit expiration dates for tokens'
+    task edit: :environment do |_t, _args|
+      Tasks::Gitlab::Tokens::ManageExpiryTask.new.edit
     end
   end
 end

lib/tasks/gitlab/tokens/manage_expiry_task.rb

+# frozen_string_literal: true
+
+require 'tty-prompt'
+
+module Tasks
+  module Gitlab
+    module Tokens
+      class ManageExpiryTask
+        TOTAL_WIDTH = 70
+
+        def analyze
+          show_pat_expires_at_migration_status
+          show_most_common_pat_expiration_dates
+        end
+
+        def edit
+          loop do
+            analyze
+
+            break unless prompt_action
+          end
+        end
+
+        private
+
+        def show_pat_expires_at_migration_status
+          sql = <<~SQL
+            SELECT * FROM batched_background_migrations
+            WHERE job_class_name = 'CleanupPersonalAccessTokensWithNilExpiresAt'
+            AND table_name = 'personal_access_tokens'
+            AND column_name = 'id'
+          SQL
+
+          print_header("Personal/Project/Group Access Token Expiration Migration")
+
+          base_model = ::Gitlab::Database.database_base_models[::Gitlab::Database::MAIN_DATABASE_NAME]
+          record = base_model.connection.select_one(sql)
+
+          if record
+            puts "Started at: #{record['started_at']}"
+            puts "Finished  : #{record['finished_at']}"
+          else
+            puts "Status: Not run"
+          end
+        end
+
+        def show_most_common_pat_expiration_dates
+          print_header "Top 10 Personal/Project/Group Access Token Expiration Dates"
+
+          puts "| Expiration Date | Count |"
+          puts "|-----------------|-------|"
+
+          with_most_common_pat_expiration_dates do |row|
+            expiry_date = row[:expires_at] || "(none)"
+            puts "| #{expiry_date.to_s.ljust(15)} | #{row[:count].to_s.ljust(5)} |"
+          end
+
+          print_footer
+        end
+
+        def with_most_common_pat_expiration_dates
+          # rubocop:disable CodeReuse/ActiveRecord -- Rake task specifically for fixing an issue
+          ApplicationRecord.with_fast_read_statement_timeout(0) do # rubocop: disable Performance/ActiveRecordSubtransactionMethods -- no subtransaction here
+            PersonalAccessToken
+              .select(:expires_at, Arel.sql('count(*)'))
+              .group(:expires_at)
+              .order(Arel.sql('count(*) DESC'))
+              .order(expires_at: :desc)
+              .limit(10)
+              .each do |row|
+              yield row
+            end
+          end
+          # rubocop:enable CodeReuse/ActiveRecord
+        end
+
+        def prompt_action
+          prompt = TTY::Prompt.new
+
+          puts ""
+          user_choice = prompt.select("What do you want to do?") do |menu|
+            menu.enum "."
+
+            menu.choice "Extend expiration date", 1
+            menu.choice "Quit", 2
+          end
+
+          case user_choice
+          when 1
+            extend_expiration_date
+            true
+          when 2
+            false
+          end
+        end
+
+        def extend_expiration_date
+          old_date = prompt_expiration_date_selection
+
+          return unless old_date.is_a?(Date)
+
+          prompt = TTY::Prompt.new
+
+          num_days = ::Gitlab::CurrentSettings.max_personal_access_token_lifetime || 365
+          new_date = old_date + num_days.days
+          new_date = prompt.ask("What would you like the new expiration date to be?", default: new_date)
+
+          new_date = Date.parse(new_date) unless new_date.is_a?(Date)
+
+          puts ""
+          puts "Old expiration date: #{old_date}"
+          puts "New expiration date: #{new_date}"
+          confirmed = prompt.yes?(
+            "WARNING: This will now update #{token_count(old_date)} token(s). Are you sure?",
+            default: false
+          )
+
+          if confirmed
+            puts "Updating tokens..."
+            update_tokens_with_expiration(old_date, new_date)
+          else
+            puts "Aborting!"
+          end
+        rescue Date::Error
+          puts "Invalid date, aborting..."
+        end
+
+        def update_tokens_with_expiration(old_date, new_date)
+          total = 0
+          # rubocop:disable CodeReuse/ActiveRecord -- Rake task specifically for fixing an issue
+          PersonalAccessToken.where(expires_at: old_date).each_batch do |batch|
+            puts "Updating personal access tokens from ID #{batch.minimum(:id)} to #{batch.maximum(:id)}..."
+            total += batch.update_all(expires_at: new_date)
+          end
+          # rubocop:enable CodeReuse/ActiveRecord -- Rake task specifically for fixing an issue
+
+          puts "Updated #{total} tokens!"
+        end
+
+        def prompt_expiration_date_selection
+          prompt = TTY::Prompt.new
+
+          choices = []
+          with_most_common_pat_expiration_dates do |row|
+            choices << row[:expires_at] if row[:expires_at]
+          end
+
+          abort_choice = "--> Abort"
+          choices << abort_choice
+
+          selection = prompt.select("Select an expiration date", choices)
+          selection == abort_choice ? nil : selection
+        end
+
+        def token_count(expiration_date)
+          # rubocop:disable CodeReuse/ActiveRecord -- Rake task specifically for fixing an issue
+          ApplicationRecord.with_fast_read_statement_timeout(0) do # rubocop: disable Performance/ActiveRecordSubtransactionMethods -- no subtransaction here
+            PersonalAccessToken.where(expires_at: expiration_date).count
+          end
+          # rubocop:enable CodeReuse/ActiveRecord
+        end
+
+        def print_header(title, total_width = TOTAL_WIDTH)
+          title_length = title.length
+          side_length = (total_width - title_length) / 2
+
+          left_side = "=" * side_length
+          right_side = "=" * (side_length + (total_width % 2))
+
+          header = "#{left_side} #{title} #{right_side}"
+          puts header
+        end
+
+        def print_footer
+          # Account for the spaces between the "=" in the header
+          puts "=" * (TOTAL_WIDTH + 2)
+        end
+      end
+    end
+  end
+end

spec/tasks/gitlab/tokens/manage_expiry_task_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require_relative '../../../../lib/tasks/gitlab/tokens/manage_expiry_task'
+
+RSpec.describe 'Tasks::Gitlab::Tokens::ManageExpiryTask', feature_category: :system_access do
+  # rubocop:disable RSpec/AvoidTestProf -- this is not a migration spec
+  let_it_be(:expires_at) { Date.today + 364.days }
+  let_it_be(:user) { create(:user) }
+  let_it_be(:migration_status) do
+    create(:batched_background_migration, :finished,
+      job_class_name: 'CleanupPersonalAccessTokensWithNilExpiresAt',
+      table_name: 'personal_access_tokens',
+      column_name: 'id')
+  end
+  # rubocop:enable RSpec/AvoidTestProf
+
+  let!(:personal_access_token1) { create(:personal_access_token, user: user, expires_at: expires_at) }
+  let!(:personal_access_token2) { create(:personal_access_token, user: user, expires_at: expires_at) }
+  let!(:personal_access_token3) { create(:personal_access_token, user: user, expires_at: expires_at + 1.day) }
+
+  subject(:task) { Tasks::Gitlab::Tokens::ManageExpiryTask.new }
+
+  describe '.analyze' do
+    it 'calls the expected methods' do
+      expect(task).to receive(:show_pat_expires_at_migration_status)
+      expect(task).to receive(:show_most_common_pat_expiration_dates)
+
+      task.analyze
+    end
+  end
+
+  describe '.edit' do
+    it 'calls analyze and prompts for action' do
+      expect(task).to receive(:analyze).at_least(:once)
+      expect(task).to receive(:prompt_action).at_least(:once).and_return(false)
+
+      task.edit
+    end
+  end
+
+  describe '.show_pat_expires_at_migration_status' do
+    it 'prints the migration status' do
+      expect { task.send(:show_pat_expires_at_migration_status) }.to output(
+        /Started at: #{migration_status[:started_at]}\nFinished  : #{migration_status[:finished_at]}/).to_stdout
+    end
+  end
+
+  describe '.show_most_common_pat_expiration_dates' do
+    let(:second) { personal_access_token3.expires_at }
+
+    it 'shows the two groups of expiration dates' do
+      expect { task.send(:show_most_common_pat_expiration_dates) }.to output(
+        /#{expires_at}.*\|\s+2\s+\|\n\|\s+#{second}\s+\|\s+1\s+/).to_stdout
+    end
+  end
+
+  describe '.extend_expiration_date' do
+    context 'with no max personal access token lifetime set' do
+      it 'extends the expiration date for selected tokens' do
+        new_date = expires_at + 1.day
+        default_date = expires_at + 365.days
+        prompt = instance_double(TTY::Prompt)
+
+        expect(task).to receive(:prompt_expiration_date_selection).and_return(expires_at)
+        expect(TTY::Prompt).to receive(:new).and_return(prompt)
+        expect(prompt).to receive(:ask).with(anything, default: default_date).and_return(new_date.to_s)
+        expect(prompt).to receive(:yes?).and_return(true)
+        expect(task).to receive(:update_tokens_with_expiration).with(expires_at, new_date).and_call_original
+
+        expect { task.send(:extend_expiration_date) }.to output(/Updated 2 tokens!/).to_stdout
+
+        expect(personal_access_token1.reload.expires_at).to eq(new_date)
+        expect(personal_access_token2.reload.expires_at).to eq(new_date)
+      end
+    end
+
+    context 'with max personal access token token lifetime set' do
+      before do
+        stub_application_setting(max_personal_access_token_lifetime: 30)
+      end
+
+      it 'asks with the max_personal_access_token_lifetime default' do
+        new_date = expires_at + 29.days
+        default_date = expires_at + 30.days
+        prompt = instance_double(TTY::Prompt)
+
+        expect(task).to receive(:prompt_expiration_date_selection).and_return(expires_at)
+        expect(TTY::Prompt).to receive(:new).and_return(prompt)
+        expect(prompt).to receive(:ask).with(anything, default: default_date).and_return(new_date.to_s)
+
+        expect(prompt).to receive(:yes?).and_return(true)
+        expect(task).to receive(:update_tokens_with_expiration).with(expires_at, new_date).and_call_original
+
+        expect { task.send(:extend_expiration_date) }.to output(/Updated 2 tokens!/).to_stdout
+
+        expect(personal_access_token1.reload.expires_at).to eq(new_date)
+        expect(personal_access_token2.reload.expires_at).to eq(new_date)
+      end
+    end
+  end
+end