Add license check before enabling DB load balancing

app/models/license.rb

@@ -6,6 +6,7 @@ class License < ActiveRecord::Base
   AUDITOR_USER_FEATURE = 'GitLab_Auditor_User'.freeze
   BURNDOWN_CHARTS_FEATURE = 'GitLab_BurndownCharts'.freeze
   CONTRIBUTION_ANALYTICS_FEATURE = 'GitLab_ContributionAnalytics'.freeze
+  DB_LOAD_BALANCING_FEATURE = 'GitLab_DbLoadBalancing'.freeze
   DEPLOY_BOARD_FEATURE = 'GitLab_DeployBoard'.freeze
   ELASTIC_SEARCH_FEATURE = 'GitLab_ElasticSearch'.freeze
   EXPORT_ISSUES_FEATURE  = 'GitLab_ExportIssues'.freeze
@@ -34,6 +35,7 @@ class License < ActiveRecord::Base
   FEATURE_CODES = {
     admin_audit_log: ADMIN_AUDIT_LOG_FEATURE,
     auditor_user: AUDITOR_USER_FEATURE,
+    db_load_balancing: DB_LOAD_BALANCING_FEATURE,
     elastic_search: ELASTIC_SEARCH_FEATURE,
     geo: GEO_FEATURE,
     object_storage: OBJECT_STORAGE_FEATURE,
@@ -98,6 +100,7 @@ class License < ActiveRecord::Base
     *EES_FEATURES,
     { ADMIN_AUDIT_LOG_FEATURE => 1 },
     { AUDITOR_USER_FEATURE => 1 },
+    { DB_LOAD_BALANCING_FEATURE => 1 },
     { DEPLOY_BOARD_FEATURE => 1 },
     { FILE_LOCKS_FEATURE => 1 },
     { GEO_FEATURE => 1 },

config/initializers/load_balancing.rb

-if Gitlab::Database::LoadBalancing.enable?
-  Gitlab::Database.disable_prepared_statements
+# We need to run this initializer after migrations are done so it doesn't fail on CI
+if ActiveRecord::Base.connected? && ActiveRecord::Base.connection.table_exists?('licenses')
+  if Gitlab::Database::LoadBalancing.enable?
+    Gitlab::Database.disable_prepared_statements
 
-  Gitlab::Application.configure do |config|
-    config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
-  end
+    Gitlab::Application.configure do |config|
+      config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+    end
 
-  Gitlab::Database::LoadBalancing.configure_proxy
+    Gitlab::Database::LoadBalancing.configure_proxy
+  end
 end

lib/gitlab/database/load_balancing.rb

@@ -45,6 +45,8 @@ module Gitlab
 
       # Returns true if load balancing is to be enabled.
       def self.enable?
+        return false unless ::License.feature_available?(:db_load_balancing)
+
         program_name != 'rake' && !hosts.empty? && !Sidekiq.server? &&
           Database.postgresql?
       end

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -26,6 +26,8 @@ describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.enable?' do
+    let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
+
     it 'returns false when no hosts are specified' do
       allow(described_class).to receive(:hosts).and_return([])
 
@@ -60,6 +62,36 @@ describe Gitlab::Database::LoadBalancing do
 
       expect(described_class.enable?).to eq(true)
     end
+
+    context 'without a license' do
+      before do
+        License.destroy_all
+      end
+
+      it 'is disabled' do
+        expect(described_class.enable?).to eq(false)
+      end
+    end
+
+    context 'with an EES license' do
+      let!(:license) { create(:license, plan: ::License::STARTER_PLAN) }
+
+      it 'is disabled' do
+        expect(described_class.enable?).to eq(false)
+      end
+    end
+
+    context 'with an EEP license' do
+      let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
+
+      it 'is enabled' do
+        allow(described_class).to receive(:hosts).and_return(%w(foo))
+        allow(Sidekiq).to receive(:server?).and_return(false)
+        allow(Gitlab::Database).to receive(:postgresql?).and_return(true)
+
+        expect(described_class.enable?).to eq(true)
+      end
+    end
   end
 
   describe '.program_name' do