Move LB calls for API helpers to Core

ee/lib/ee/api/helpers.rb

@@ -37,20 +37,6 @@ module EE
         unauthorized! unless ::Gitlab::Geo.allowed_ip?(request.ip)
       end
 
-      override :current_user
-      def current_user
-        strong_memoize(:current_user) do
-          user = super
-
-          if user
-            ::Gitlab::Database::LoadBalancing::RackMiddleware
-              .stick_or_unstick(env, :user, user.id)
-          end
-
-          user
-        end
-      end
-
       def authorization_header_valid?
         return unless gitlab_geo_node_token?
 

ee/lib/ee/api/helpers/runner.rb

@@ -6,30 +6,6 @@ module EE
       module Runner
         extend ::Gitlab::Utils::Override
 
-        override :current_job
-        def current_job
-          id = params[:id]
-
-          if id
-            ::Gitlab::Database::LoadBalancing::RackMiddleware
-              .stick_or_unstick(env, :build, id)
-          end
-
-          super
-        end
-
-        override :current_runner
-        def current_runner
-          token = params[:token]
-
-          if token
-            ::Gitlab::Database::LoadBalancing::RackMiddleware
-              .stick_or_unstick(env, :runner, token)
-          end
-
-          super
-        end
-
         override :track_ci_minutes_usage!
         def track_ci_minutes_usage!(build, runner)
           ::Ci::Minutes::TrackLiveConsumptionService.new(build).execute

ee/spec/lib/ee/api/helpers_spec.rb

@@ -26,44 +26,6 @@ RSpec.describe EE::API::Helpers do
     helper
   end
 
-  describe '#current_user' do
-    let(:user) { build(:user, id: 42) }
-
-    before do
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-    end
-
-    it 'handles sticking when a user could be found' do
-      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
-
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .to receive(:stick_or_unstick).with(any_args, :user, 42)
-
-      get 'user'
-
-      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'id' => user.id })
-    end
-
-    it 'does not handle sticking if no user could be found' do
-      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(nil)
-
-      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
-        .not_to receive(:stick_or_unstick)
-
-      get 'user'
-
-      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'found' => false })
-    end
-
-    it 'returns the user if one could be found' do
-      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
-
-      get 'user'
-
-      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'id' => user.id })
-    end
-  end
-
   describe '#authenticate_by_gitlab_geo_node_token!' do
     let(:invalid_geo_auth_header) { "#{::Gitlab::Geo::BaseRequest::GITLAB_GEO_AUTH_TOKEN_TYPE}...Test" }
 

lib/api/helpers.rb

@@ -74,6 +74,11 @@ module API
 
       save_current_user_in_env(@current_user) if @current_user
 
+      if @current_user
+        ::Gitlab::Database::LoadBalancing::RackMiddleware
+          .stick_or_unstick(env, :user, @current_user.id)
+      end
+
       @current_user
     end
     # rubocop:enable Gitlab/ModuleWithInstanceVariables

lib/api/helpers/runner.rb

@@ -34,8 +34,15 @@ module API
       end
 
       def current_runner
+        token = params[:token]
+
+        if token
+          ::Gitlab::Database::LoadBalancing::RackMiddleware
+            .stick_or_unstick(env, :runner, token)
+        end
+
         strong_memoize(:current_runner) do
-          ::Ci::Runner.find_by_token(params[:token].to_s)
+          ::Ci::Runner.find_by_token(token.to_s)
         end
       end
 
@@ -65,8 +72,15 @@ module API
       end
 
       def current_job
+        id = params[:id]
+
+        if id
+          ::Gitlab::Database::LoadBalancing::RackMiddleware
+            .stick_or_unstick(env, :build, id)
+        end
+
         strong_memoize(:current_job) do
-          ::Ci::Build.find_by_id(params[:id])
+          ::Ci::Build.find_by_id(id)
         end
       end
 

ee/spec/lib/ee/api/helpers/runner_spec.rb → spec/lib/api/helpers/runner_spec.rb

@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe EE::API::Helpers::Runner do
+RSpec.describe API::Helpers::Runner do
   let(:helper) { Class.new { include API::Helpers::Runner }.new }
 
   before do

spec/lib/api/helpers_spec.rb

@@ -7,6 +7,66 @@ RSpec.describe API::Helpers do
 
   subject { Class.new.include(described_class).new }
 
+  describe '#current_user' do
+    include Rack::Test::Methods
+
+    let(:user) { build(:user, id: 42) }
+
+    let(:helper) do
+      Class.new(Grape::API::Instance) do
+        helpers API::APIGuard::HelperMethods
+        helpers API::Helpers
+        format :json
+
+        get 'user' do
+          current_user ? { id: current_user.id } : { found: false }
+        end
+
+        get 'protected' do
+          authenticate_by_gitlab_geo_node_token!
+        end
+      end
+    end
+
+    def app
+      helper
+    end
+
+    before do
+      allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+    end
+
+    it 'handles sticking when a user could be found' do
+      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
+
+      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
+        .to receive(:stick_or_unstick).with(any_args, :user, 42)
+
+      get 'user'
+
+      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'id' => user.id })
+    end
+
+    it 'does not handle sticking if no user could be found' do
+      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(nil)
+
+      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
+        .not_to receive(:stick_or_unstick)
+
+      get 'user'
+
+      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'found' => false })
+    end
+
+    it 'returns the user if one could be found' do
+      allow_any_instance_of(API::Helpers).to receive(:initial_current_user).and_return(user)
+
+      get 'user'
+
+      expect(Gitlab::Json.parse(last_response.body)).to eq({ 'id' => user.id })
+    end
+  end
+
   describe '#find_project' do
     let(:project) { create(:project) }