Read feature_category from request for GraphQL controller

- We also update `web_transaction` to read from `ApplicationContext`.
- We tried using `around_action` and `ApplicationContext.with`
  in the GraphqlController, but the `web_transaction` methods
  were called **after** the `ApplicationController` set the
  context, but **before** the `GraphqlController` could kick
  in.

app/controllers/graphql_controller.rb

 # frozen_string_literal: true
 
 class GraphqlController < ApplicationController
+  extend ::Gitlab::Utils::Override
+
   # Unauthenticated users have access to the API for public data
   skip_before_action :authenticate_user!
 
@@ -35,6 +37,7 @@ class GraphqlController < ApplicationController
   # callback execution order here
   around_action :sessionless_bypass_admin_mode!, if: :sessionless_user?
 
+  # The default feature category is overridden to read from request
   feature_category :not_owned
 
   def execute
@@ -64,6 +67,11 @@ class GraphqlController < ApplicationController
     render_error(exception.message, status: :unprocessable_entity)
   end
 
+  override :feature_category
+  def feature_category
+    ::Gitlab::FeatureCategories.default.from_request(request) || super
+  end
+
   private
 
   def disallow_mutations_for_get

lib/gitlab/feature_categories.rb

+# frozen_string_literal: true
+
+module Gitlab
+  class FeatureCategories
+    FEATURE_CATEGORY_DEFAULT = 'unknown'
+
+    attr_reader :categories
+
+    def self.default
+      @default ||= self.load_from_yaml
+    end
+
+    def self.load_from_yaml
+      categories = YAML.load_file(Rails.root.join('config', 'feature_categories.yml'))
+
+      new(categories)
+    end
+
+    def initialize(categories)
+      @categories = categories.to_set
+    end
+
+    # If valid, returns a feature category from the given request.
+    def from_request(request)
+      category = request.headers["HTTP_X_GITLAB_FEATURE_CATEGORY"].presence
+
+      return unless category && valid?(category)
+
+      return unless ::Gitlab::RequestForgeryProtection.verified?(request.env)
+
+      category
+    end
+
+    def valid?(category)
+      categories.include?(category.to_s)
+    end
+  end
+end

lib/gitlab/metrics/requests_rack_middleware.rb

@@ -15,7 +15,7 @@ module Gitlab
 
       HEALTH_ENDPOINT = %r{^/-/(liveness|readiness|health|metrics)/?$}.freeze
 
-      FEATURE_CATEGORY_DEFAULT = 'unknown'
+      FEATURE_CATEGORY_DEFAULT = ::Gitlab::FeatureCategories::FEATURE_CATEGORY_DEFAULT
       ENDPOINT_MISSING = 'unknown'
 
       # These were the top 5 categories at a point in time, chosen as a

lib/gitlab/metrics/web_transaction.rb

@@ -57,10 +57,6 @@ module Gitlab
 
         action = "#{controller.action_name}"
 
-        # Try to get the feature category, but don't fail when the controller is
-        # not an ApplicationController.
-        feature_category = controller.class.try(:feature_category_for_action, action).to_s
-
         # Devise exposes a method called "request_format" that does the below.
         # However, this method is not available to all controllers (e.g. certain
         # Doorkeeper controllers). As such we use the underlying code directly.
@@ -91,9 +87,6 @@ module Gitlab
         if route
           path = endpoint_paths_cache[route.request_method][route.path]
 
-          grape_class = endpoint.options[:for]
-          feature_category = grape_class.try(:feature_category_for_app, endpoint).to_s
-
           { controller: 'Grape', action: "#{route.request_method} #{path}", feature_category: feature_category }
         end
       end
@@ -109,6 +102,10 @@ module Gitlab
       def endpoint_instrumentable_path(raw_path)
         raw_path.sub('(.:format)', '').sub('/:version', '')
       end
+
+      def feature_category
+        ::Gitlab::ApplicationContext.current_context_attribute(:feature_category) || ::Gitlab::FeatureCategories::FEATURE_CATEGORY_DEFAULT
+      end
     end
   end
 end

spec/controllers/graphql_controller_spec.rb

@@ -38,6 +38,14 @@ RSpec.describe GraphqlController do
         sign_in(user)
       end
 
+      it 'sets feature category in ApplicationContext from request' do
+        request.headers["HTTP_X_GITLAB_FEATURE_CATEGORY"] = "web_ide"
+
+        post :execute
+
+        expect(::Gitlab::ApplicationContext.current_context_attribute(:feature_category)).to eq('web_ide')
+      end
+
       it 'returns 200 when user can access API' do
         post :execute
 

spec/lib/gitlab/feature_categories_spec.rb

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+
+RSpec.describe Gitlab::FeatureCategories do
+  let(:fake_categories) { %w(foo bar) }
+
+  subject { described_class.new(fake_categories) }
+
+  describe "#valid?" do
+    it "returns true if category is known", :aggregate_failures do
+      expect(subject.valid?('foo')).to be(true)
+      expect(subject.valid?('zzz')).to be(false)
+    end
+  end
+
+  describe "#from_request" do
+    let(:request_env) { {} }
+    let(:verified) { true }
+
+    def fake_request(request_feature_category)
+      double('request', env: request_env, headers: { "HTTP_X_GITLAB_FEATURE_CATEGORY" => request_feature_category })
+    end
+
+    before do
+      allow(::Gitlab::RequestForgeryProtection).to receive(:verified?).with(request_env).and_return(verified)
+    end
+
+    it "returns category from request when valid, otherwise returns nil", :aggregate_failures do
+      expect(subject.from_request(fake_request("foo"))).to be("foo")
+      expect(subject.from_request(fake_request("zzz"))).to be_nil
+    end
+
+    context "when request is not verified" do
+      let(:verified) { false }
+
+      it "returns nil" do
+        expect(subject.from_request(fake_request("foo"))).to be_nil
+      end
+    end
+  end
+
+  describe "#categories" do
+    it "returns a set of the given categories" do
+      expect(subject.categories).to be_a(Set)
+      expect(subject.categories).to contain_exactly(*fake_categories)
+    end
+  end
+
+  describe ".load_from_yaml" do
+    subject { described_class.load_from_yaml }
+
+    it "creates FeatureCategories from feature_categories.yml file" do
+      contents = YAML.load_file(Rails.root.join('config', 'feature_categories.yml'))
+
+      expect(subject.categories).to contain_exactly(*contents)
+    end
+  end
+
+  describe ".default" do
+    it "returns a memoization of load_from_yaml", :aggregate_failures do
+      # FeatureCategories.default could have been referenced in another spec, so we need to clean it up here
+      described_class.instance_variable_set(:@default, nil)
+
+      expect(described_class).to receive(:load_from_yaml).once.and_call_original
+
+      2.times { described_class.default }
+
+      # Uses reference equality to verify memoization
+      expect(described_class.default).to equal(described_class.default)
+      expect(described_class.default).to be_a(described_class)
+    end
+  end
+end

spec/lib/gitlab/metrics/requests_rack_middleware_spec.rb

@@ -348,7 +348,7 @@ RSpec.describe Gitlab::Metrics::RequestsRackMiddleware, :aggregate_failures do
       end
 
       it 'has every label in config/feature_categories.yml' do
-        defaults = [described_class::FEATURE_CATEGORY_DEFAULT, 'not_owned']
+        defaults = [::Gitlab::FeatureCategories::FEATURE_CATEGORY_DEFAULT, 'not_owned']
         feature_categories = YAML.load_file(Rails.root.join('config', 'feature_categories.yml')).map(&:strip) + defaults
 
         expect(described_class::FEATURE_CATEGORIES_TO_INITIALIZE).to all(be_in(feature_categories))

spec/lib/gitlab/metrics/web_transaction_spec.rb

@@ -32,7 +32,7 @@ RSpec.describe Gitlab::Metrics::WebTransaction do
 
     it 'measures with correct labels and value' do
       value = 1
-      expect(prometheus_metric).to receive(metric_method).with({ controller: 'TestController', action: 'show', feature_category: '' }, value)
+      expect(prometheus_metric).to receive(metric_method).with({ controller: 'TestController', action: 'show', feature_category: ::Gitlab::FeatureCategories::FEATURE_CATEGORY_DEFAULT }, value)
 
       transaction.send(metric_method, :bau, value)
     end
@@ -105,6 +105,9 @@ RSpec.describe Gitlab::Metrics::WebTransaction do
                           namespace: "/projects")
 
         env['api.endpoint'] = endpoint
+
+        # This is needed since we're not actually making a request, which would trigger the controller pushing to the context
+        ::Gitlab::ApplicationContext.push(feature_category: 'projects')
       end
 
       it 'provides labels with the method and path of the route in the grape endpoint' do
@@ -129,7 +132,7 @@ RSpec.describe Gitlab::Metrics::WebTransaction do
       include_context 'ActionController request'
 
       it 'tags a transaction with the name and action of a controller' do
-        expect(transaction.labels).to eq({ controller: 'TestController', action: 'show', feature_category: '' })
+        expect(transaction.labels).to eq({ controller: 'TestController', action: 'show', feature_category: ::Gitlab::FeatureCategories::FEATURE_CATEGORY_DEFAULT })
       end
 
       it 'contains only the labels defined for transactions' do
@@ -140,7 +143,7 @@ RSpec.describe Gitlab::Metrics::WebTransaction do
         let(:request) { double(:request, format: double(:format, ref: :json)) }
 
         it 'appends the mime type to the transaction action' do
-          expect(transaction.labels).to eq({ controller: 'TestController', action: 'show.json', feature_category: '' })
+          expect(transaction.labels).to eq({ controller: 'TestController', action: 'show.json', feature_category: ::Gitlab::FeatureCategories::FEATURE_CATEGORY_DEFAULT })
         end
       end
 
@@ -148,13 +151,15 @@ RSpec.describe Gitlab::Metrics::WebTransaction do
         let(:request) { double(:request, format: double(:format, ref: 'http://example.com')) }
 
         it 'does not append the MIME type to the transaction action' do
-          expect(transaction.labels).to eq({ controller: 'TestController', action: 'show', feature_category: '' })
+          expect(transaction.labels).to eq({ controller: 'TestController', action: 'show', feature_category: ::Gitlab::FeatureCategories::FEATURE_CATEGORY_DEFAULT })
         end
       end
 
       context 'when the feature category is known' do
         it 'includes it in the feature category label' do
-          expect(controller_class).to receive(:feature_category_for_action).with('show').and_return(:source_code_management)
+          # This is needed since we're not actually making a request, which would trigger the controller pushing to the context
+          ::Gitlab::ApplicationContext.push(feature_category: 'source_code_management')
+
           expect(transaction.labels).to eq({ controller: 'TestController', action: 'show', feature_category: "source_code_management" })
         end
       end