Add plan limits for Premium and Ultimate

db/migrate/20210706213537_add_premium_and_ultimate_plan_limits.rb

+# frozen_string_literal: true
+
+class AddPremiumAndUltimatePlanLimits < ActiveRecord::Migration[6.1]
+  include Gitlab::Database::MigrationHelpers
+
+  class Plan < ActiveRecord::Base
+    self.inheritance_column = :_type_disabled
+
+    has_one :limits, class_name: 'PlanLimits'
+  end
+
+  class PlanLimits < ActiveRecord::Base
+    self.inheritance_column = :_type_disabled
+
+    belongs_to :plan
+  end
+
+  def copy_plan_limits(from_plan_name:, to_plan_name:)
+    source_plan = Plan.find_by(name: from_plan_name)
+    target_plan = Plan.find_by(name: to_plan_name)
+    return unless source_plan && target_plan
+    return unless source_plan.limits.present?
+    return if target_plan.limits.present?
+
+    limits = source_plan.limits.dup
+    limits.plan = target_plan
+    limits.save!
+  end
+
+  def up
+    return unless Gitlab.com?
+
+    copy_plan_limits(from_plan_name: 'gold', to_plan_name: 'ultimate')
+    copy_plan_limits(from_plan_name: 'silver', to_plan_name: 'premium')
+  end
+
+  def down
+    # no-op
+  end
+end

db/schema_migrations/20210706213537

+150463cef309e6bf69240c258dc8aede53b846a08a7e2d668ee0429709022554
\ No newline at end of file

doc/development/application_limits.md

@@ -48,7 +48,11 @@ It's recommended to create two separate migration script files.
        create_or_update_plan_limit('project_hooks', 'free', 10)
        create_or_update_plan_limit('project_hooks', 'bronze', 20)
        create_or_update_plan_limit('project_hooks', 'silver', 30)
+       create_or_update_plan_limit('project_hooks', 'premium', 30)
+       create_or_update_plan_limit('project_hooks', 'premium_trial', 30)
        create_or_update_plan_limit('project_hooks', 'gold', 100)
+       create_or_update_plan_limit('project_hooks', 'ultimate', 100)
+       create_or_update_plan_limit('project_hooks', 'ultimate_trial', 100)
      end
 
      def down
@@ -56,7 +60,11 @@ It's recommended to create two separate migration script files.
        create_or_update_plan_limit('project_hooks', 'free', 0)
        create_or_update_plan_limit('project_hooks', 'bronze', 0)
        create_or_update_plan_limit('project_hooks', 'silver', 0)
+       create_or_update_plan_limit('project_hooks', 'premium', 0)
+       create_or_update_plan_limit('project_hooks', 'premium_trial', 0)
        create_or_update_plan_limit('project_hooks', 'gold', 0)
+       create_or_update_plan_limit('project_hooks', 'ultimate', 0)
+       create_or_update_plan_limit('project_hooks', 'ultimate_trial', 0)
      end
    end
    ```
@@ -145,6 +153,10 @@ GitLab.com:
 - `free`: Namespaces and projects with a Free subscription.
 - `bronze`: Namespaces and projects with a Bronze subscription. This tier is no longer available for purchase.
 - `silver`: Namespaces and projects with a Premium subscription.
+- `premium`: Namespaces and projects with a Premium subscription.
+- `premium_trial`: Namespaces and projects with a Premium Trial subscription.
 - `gold`: Namespaces and projects with an Ultimate subscription.
+- `ultimate`: Namespaces and projects with an Ultimate subscription.
+- `ultimate_trial`: Namespaces and projects with an Ultimate Trial subscription.
 
 The `test` environment doesn't have any plans.

spec/migrations/add_premium_and_ultimate_plan_limits_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+require_migration!
+
+RSpec.describe AddPremiumAndUltimatePlanLimits, :migration do
+  shared_examples_for 'a migration that does not alter plans or plan limits' do
+    it do
+      expect { migrate! }.not_to change {
+        [
+          AddPremiumAndUltimatePlanLimits::Plan.count,
+          AddPremiumAndUltimatePlanLimits::PlanLimits.count
+        ]
+      }
+    end
+  end
+
+  describe '#up' do
+    context 'when not .com?' do
+      before do
+        allow(Gitlab).to receive(:com?).and_return false
+      end
+
+      it_behaves_like 'a migration that does not alter plans or plan limits'
+    end
+
+    context 'when .com?' do
+      before do
+        allow(Gitlab).to receive(:com?).and_return true
+      end
+
+      context 'when source plan does not exist' do
+        it_behaves_like 'a migration that does not alter plans or plan limits'
+      end
+
+      context 'when target plan does not exist' do
+        before do
+          table(:plans).create!(name: 'silver', title: 'Silver')
+          table(:plans).create!(name: 'gold', title: 'Gold')
+        end
+
+        it_behaves_like 'a migration that does not alter plans or plan limits'
+      end
+
+      context 'when source and target plans exist' do
+        let!(:silver) { table(:plans).create!(name: 'silver', title: 'Silver') }
+        let!(:gold) { table(:plans).create!(name: 'gold', title: 'Gold') }
+        let!(:premium) { table(:plans).create!(name: 'premium', title: 'Premium') }
+        let!(:ultimate) { table(:plans).create!(name: 'ultimate', title: 'Ultimate') }
+
+        let!(:silver_limits) { table(:plan_limits).create!(plan_id: silver.id, storage_size_limit: 111) }
+        let!(:gold_limits) { table(:plan_limits).create!(plan_id: gold.id, storage_size_limit: 222) }
+
+        context 'when target has plan limits' do
+          before do
+            table(:plan_limits).create!(plan_id: premium.id, storage_size_limit: 999)
+            table(:plan_limits).create!(plan_id: ultimate.id, storage_size_limit: 999)
+          end
+
+          it 'does not overwrite the limits' do
+            expect { migrate! }.not_to change {
+              [
+                AddPremiumAndUltimatePlanLimits::Plan.count,
+                AddPremiumAndUltimatePlanLimits::PlanLimits.pluck(:id, :storage_size_limit).sort
+              ]
+            }
+          end
+        end
+
+        context 'when target has no plan limits' do
+          it 'creates plan limits from the source plan' do
+            migrate!
+
+            expect(AddPremiumAndUltimatePlanLimits::PlanLimits.pluck(:plan_id, :storage_size_limit)).to match_array([
+              [silver.id, silver_limits.storage_size_limit],
+              [gold.id, gold_limits.storage_size_limit],
+              [premium.id, silver_limits.storage_size_limit],
+              [ultimate.id, gold_limits.storage_size_limit]
+            ])
+          end
+        end
+      end
+    end
+  end
+end