Merge branch 'prepend-ee-helper' into 'master'

Add methods for injecting EE modules

See merge request gitlab-org/gitlab-ce!31055

5 files changed, 186 insertions, 70 deletions

diff --git a/config/initializers/0_inject_enterprise_edition_module.rb b/config/initializers/0_inject_enterprise_edition_module.rb
new file mode 100644
index 00000000000..4b21732e179
--- /dev/null
+++ b/config/initializers/0_inject_enterprise_edition_module.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module InjectEnterpriseEditionModule
+  def prepend_if_ee(constant)
+    prepend(constant.constantize) if Gitlab.ee?
+  end
+
+  def extend_if_ee(constant)
+    extend(constant.constantize) if Gitlab.ee?
+  end
+
+  def include_if_ee(constant)
+    include(constant.constantize) if Gitlab.ee?
+  end
+end
+
+Module.prepend(InjectEnterpriseEditionModule)
diff --git a/doc/development/ee_features.md b/doc/development/ee_features.md
index 7131b717353..2217dedccd3 100644
--- a/doc/development/ee_features.md
+++ b/doc/development/ee_features.md
@@ -125,20 +125,24 @@ This also applies to views.
 ### EE features based on CE features
 
 For features that build on existing CE features, write a module in the `EE`
-namespace and `prepend` it in the CE class, on the last line of the file that
-the class resides in. This makes conflicts less likely to happen during CE to EE
-merges because only one line is added to the CE class - the `prepend` line. For
-example, to prepend a module into the `User` class you would use the following
-approach:
+namespace and inject it in the CE class, on the last line of the file that the
+class resides in. This makes conflicts less likely to happen during CE to EE
+merges because only one line is added to the CE class - the line that injects
+the module. For example, to prepend a module into the `User` class you would use
+the following approach:
 
 ```ruby
 class User < ActiveRecord::Base
   # ... lots of code here ...
 end
 
-User.prepend(EE::User)
+User.prepend_if_ee('EE::User')
 ```
 
+Do not use methods such as `prepend`, `extend`, and `include`. Instead, use
+`prepend_if_ee`, `extend_if_ee`, or `include_if_ee`. These methods take a
+_String_ containing the full module name as the argument, not the module itself.
+
 Since the module would require an `EE` namespace, the file should also be
 put in an `ee/` sub-directory. For example, we want to extend the user model
 in EE, so we have a module called `::EE::User` put inside
@@ -255,7 +259,7 @@ class ApplicationController < ActionController::Base
   # ...
 end
 
-ApplicationController.prepend(EE::ApplicationController)
+ApplicationController.prepend_if_ee('EE::ApplicationController')
 ```
 
 And create a new file in the `ee/` sub-directory with the altered
@@ -504,9 +508,9 @@ EE-specific LDAP classes in `ee/lib/ee/gitlab/ldap`.
 
 ### Code in `lib/api/`
 
-It can be very tricky to extend EE features by a single line of `prepend`,
-and for each different [Grape](https://github.com/ruby-grape/grape) feature,
-we might need different strategies to extend it. To apply different strategies
+It can be very tricky to extend EE features by a single line of `prepend_if_ee`,
+and for each different [Grape](https://github.com/ruby-grape/grape) feature, we
+might need different strategies to extend it. To apply different strategies
 easily, we would use `extend ActiveSupport::Concern` in the EE module.
 
 Put the EE module files following
@@ -543,12 +547,12 @@ constants.
 We can define `params` and utilize `use` in another `params` definition to
 include params defined in EE. However, we need to define the "interface" first
 in CE in order for EE to override it. We don't have to do this in other places
-due to `prepend`, but Grape is complex internally and we couldn't easily do
-that, so we'll follow regular object-oriented practices that we define the
+due to `prepend_if_ee`, but Grape is complex internally and we couldn't easily
+do that, so we'll follow regular object-oriented practices that we define the
 interface first here.
 
 For example, suppose we have a few more optional params for EE. We can move the
-params out of the `Grape::API` class to a helper module, so we can `prepend` it
+params out of the `Grape::API` class to a helper module, so we can inject it
 before it would be used in the class.
 
 ```ruby
@@ -583,7 +587,7 @@ module API
   end
 end
 
-API::Helpers::ProjectsHelpers.prepend(EE::API::Helpers::ProjectsHelpers)
+API::Helpers::ProjectsHelpers.prepend_if_ee('EE::API::Helpers::ProjectsHelpers')
 ```
 
 We could override it in EE module:
@@ -624,7 +628,7 @@ module API
   end
 end
 
-API::JobArtifacts.prepend(EE::API::JobArtifacts)
+API::JobArtifacts.prepend_if_ee('EE::API::JobArtifacts')
 ```
 
 And then we can follow regular object-oriented practices to override it:
@@ -677,7 +681,7 @@ module API
   end
 end
 
-API::MergeRequests.prepend(EE::API::MergeRequests)
+API::MergeRequests.prepend_if_ee('EE::API::MergeRequests')
 ```
 
 Note that `update_merge_request_ee` doesn't do anything in CE, but
@@ -717,8 +721,8 @@ Sometimes we need to use different arguments for a particular API route, and we
 can't easily extend it with an EE module because Grape has different context in
 different blocks. In order to overcome this, we need to move the data to a class
 method that resides in a separate module or class. This allows us to extend that
-module or class before its data is used, without having to place a `prepend` in
-the middle of CE code.
+module or class before its data is used, without having to place a
+`prepend_if_ee` in the middle of CE code.
 
 For example, in one place we need to pass an extra argument to
 `at_least_one_of` so that the API could consider an EE-only argument as the
@@ -739,7 +743,7 @@ module API
   end
 end
 
-API::MergeRequests::Parameters.prepend(EE::API::MergeRequests::Parameters)
+API::MergeRequests::Parameters.prepend_if_ee('EE::API::MergeRequests::Parameters')
 
 # api/merge_requests.rb
 module API
@@ -789,7 +793,7 @@ class Identity < ActiveRecord::Base
     [:provider]
   end
 
-  prepend EE::Identity
+  prepend_if_ee('EE::Identity')
 
   validates :extern_uid,
     allow_blank: true,
@@ -841,7 +845,7 @@ class Identity < ActiveRecord::Base
   end
 end
 
-Identity::UniquenessScopes.prepend(EE::Identity::UniquenessScopes)
+Identity::UniquenessScopes.prepend_if_ee('EE::Identity::UniquenessScopes')
 
 # app/models/identity.rb
 class Identity < ActiveRecord::Base
diff --git a/rubocop/cop/inject_enterprise_edition_module.rb b/rubocop/cop/inject_enterprise_edition_module.rb
index 1d37b1bd12d..e0e1b2d6c7d 100644
--- a/rubocop/cop/inject_enterprise_edition_module.rb
+++ b/rubocop/cop/inject_enterprise_edition_module.rb
@@ -6,23 +6,39 @@ module RuboCop
     # the last line of a file. Injecting a module in the middle of a file will
     # cause merge conflicts, while placing it on the last line will not.
     class InjectEnterpriseEditionModule < RuboCop::Cop::Cop
-      MSG = 'Injecting EE modules must be done on the last line of this file' \
-        ', outside of any class or module definitions'
+      INVALID_LINE = 'Injecting EE modules must be done on the last line of this file' \
+          ', outside of any class or module definitions'
 
-      METHODS = Set.new(%i[include extend prepend]).freeze
+      DISALLOWED_METHOD =
+        'EE modules must be injected using `include_if_ee`, `extend_if_ee`, or `prepend_if_ee`'
+
+      INVALID_ARGUMENT = 'EE modules to inject must be specified as a String'
+
+      CHECK_LINE_METHODS =
+        Set.new(%i[include_if_ee extend_if_ee prepend_if_ee]).freeze
+
+      DISALLOW_METHODS = Set.new(%i[include extend prepend]).freeze
 
       def ee_const?(node)
         line = node.location.expression.source_line
 
         # We use `match?` here instead of RuboCop's AST matching, as this makes
         # it far easier to handle nested constants such as `EE::Foo::Bar::Baz`.
-        line.match?(/(\s|\()(::)?EE::/)
+        line.match?(/(\s|\()('|")?(::)?EE::/)
       end
 
       def on_send(node)
-        return unless METHODS.include?(node.children[1])
-        return unless ee_const?(node.children[2])
+        return unless check_method?(node)
 
+        if DISALLOW_METHODS.include?(node.children[1])
+          add_offense(node, message: DISALLOWED_METHOD)
+        else
+          verify_line_number(node)
+          verify_argument_type(node)
+        end
+      end
+
+      def verify_line_number(node)
         line = node.location.line
         buffer = node.location.expression.source_buffer
         last_line = buffer.last_line
@@ -32,7 +48,25 @@ module RuboCop
         # the expression is the last line _of code_.
         last_line -= 1 if buffer.source.end_with?("\n")
 
-        add_offense(node) if line < last_line
+        add_offense(node, message: INVALID_LINE) if line < last_line
+      end
+
+      def verify_argument_type(node)
+        argument = node.children[2]
+
+        return if argument.str_type?
+
+        add_offense(argument, message: INVALID_ARGUMENT)
+      end
+
+      def check_method?(node)
+        name = node.children[1]
+
+        if CHECK_LINE_METHODS.include?(name) || DISALLOW_METHODS.include?(name)
+          ee_const?(node.children[2])
+        else
+          false
+        end
       end
 
       # Automatically correcting these offenses is not always possible, as
diff --git a/spec/fast_spec_helper.rb b/spec/fast_spec_helper.rb
index 91ef7653822..f5a487b4d57 100644
--- a/spec/fast_spec_helper.rb
+++ b/spec/fast_spec_helper.rb
@@ -4,6 +4,7 @@ ENV['GITLAB_ENV'] = 'test'
 ENV['IN_MEMORY_APPLICATION_SETTINGS'] = 'true'
 
 require 'active_support/dependencies'
+require_relative '../config/initializers/0_inject_enterprise_edition_module'
 require_relative '../config/settings'
 require_relative 'support/rspec'
 require 'active_support/all'
diff --git a/spec/rubocop/cop/inject_enterprise_edition_module_spec.rb b/spec/rubocop/cop/inject_enterprise_edition_module_spec.rb
index 0ff777388e5..27df42c0aee 100644
--- a/spec/rubocop/cop/inject_enterprise_edition_module_spec.rb
+++ b/spec/rubocop/cop/inject_enterprise_edition_module_spec.rb
@@ -10,91 +10,91 @@ describe RuboCop::Cop::InjectEnterpriseEditionModule do
 
   subject(:cop) { described_class.new }
 
-  it 'flags the use of `prepend EE` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend EE::Foo
-      ^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee 'EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'does not flag the use of `prepend EEFoo` in the middle of a file' do
+  it 'does not flag the use of `prepend_if_ee EEFoo` in the middle of a file' do
     expect_no_offenses(<<~SOURCE)
     class Foo
-      prepend EEFoo
+      prepend_if_ee 'EEFoo'
     end
     SOURCE
   end
 
-  it 'flags the use of `prepend EE::Foo::Bar` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee EE::Foo::Bar` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend EE::Foo::Bar
-      ^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee 'EE::Foo::Bar'
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `prepend(EE::Foo::Bar)` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee(EE::Foo::Bar)` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend(EE::Foo::Bar)
-      ^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee('EE::Foo::Bar')
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `prepend EE::Foo::Bar::Baz` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee EE::Foo::Bar::Baz` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend EE::Foo::Bar::Baz
-      ^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee 'EE::Foo::Bar::Baz'
+      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `prepend ::EE` in the middle of a file' do
+  it 'flags the use of `prepend_if_ee ::EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      prepend ::EE::Foo
-      ^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      prepend_if_ee '::EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `include EE` in the middle of a file' do
+  it 'flags the use of `include_if_ee EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      include EE::Foo
-      ^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      include_if_ee 'EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `include ::EE` in the middle of a file' do
+  it 'flags the use of `include_if_ee ::EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      include ::EE::Foo
-      ^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      include_if_ee '::EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `extend EE` in the middle of a file' do
+  it 'flags the use of `extend_if_ee EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      extend EE::Foo
-      ^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      extend_if_ee 'EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
 
-  it 'flags the use of `extend ::EE` in the middle of a file' do
+  it 'flags the use of `extend_if_ee ::EE` in the middle of a file' do
     expect_offense(<<~SOURCE)
     class Foo
-      extend ::EE::Foo
-      ^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
+      extend_if_ee '::EE::Foo'
+      ^^^^^^^^^^^^^^^^^^^^^^^^ Injecting EE modules must be done on the last line of this file, outside of any class or module definitions
     end
     SOURCE
   end
@@ -102,7 +102,7 @@ describe RuboCop::Cop::InjectEnterpriseEditionModule do
   it 'does not flag prepending of regular modules' do
     expect_no_offenses(<<~SOURCE)
     class Foo
-      prepend Foo
+      prepend_if_ee 'Foo'
     end
     SOURCE
   end
@@ -110,7 +110,7 @@ describe RuboCop::Cop::InjectEnterpriseEditionModule do
   it 'does not flag including of regular modules' do
     expect_no_offenses(<<~SOURCE)
     class Foo
-      include Foo
+      include_if_ee 'Foo'
     end
     SOURCE
   end
@@ -118,51 +118,111 @@ describe RuboCop::Cop::InjectEnterpriseEditionModule do
   it 'does not flag extending using regular modules' do
     expect_no_offenses(<<~SOURCE)
     class Foo
-      extend Foo
+      extend_if_ee 'Foo'
     end
     SOURCE
   end
 
-  it 'does not flag the use of `prepend EE` on the last line' do
+  it 'does not flag the use of `prepend_if_ee EE` on the last line' do
     expect_no_offenses(<<~SOURCE)
     class Foo
     end
 
-    Foo.prepend(EE::Foo)
+    Foo.prepend_if_ee('EE::Foo')
     SOURCE
   end
 
-  it 'does not flag the use of `include EE` on the last line' do
+  it 'does not flag the use of `include_if_ee EE` on the last line' do
     expect_no_offenses(<<~SOURCE)
     class Foo
     end
 
-    Foo.include(EE::Foo)
+    Foo.include_if_ee('EE::Foo')
     SOURCE
   end
 
-  it 'does not flag the use of `extend EE` on the last line' do
+  it 'does not flag the use of `extend_if_ee EE` on the last line' do
     expect_no_offenses(<<~SOURCE)
     class Foo
     end
 
-    Foo.extend(EE::Foo)
+    Foo.extend_if_ee('EE::Foo')
     SOURCE
   end
 
   it 'autocorrects offenses by just disabling the Cop' do
     source = <<~SOURCE
     class Foo
-      prepend EE::Foo
-      include Bar
+      prepend_if_ee 'EE::Foo'
+      include_if_ee 'Bar'
     end
     SOURCE
 
     expect(autocorrect_source(source)).to eq(<<~SOURCE)
     class Foo
-      prepend EE::Foo # rubocop: disable Cop/InjectEnterpriseEditionModule
-      include Bar
+      prepend_if_ee 'EE::Foo' # rubocop: disable Cop/InjectEnterpriseEditionModule
+      include_if_ee 'Bar'
+    end
+    SOURCE
+  end
+
+  it 'disallows the use of prepend to inject an EE module' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.prepend(EE::Foo)
+    ^^^^^^^^^^^^^^^^^^^^ EE modules must be injected using `include_if_ee`, `extend_if_ee`, or `prepend_if_ee`
+    SOURCE
+  end
+
+  it 'disallows the use of extend to inject an EE module' do
+    expect_offense(<<~SOURCE)
+    class Foo
     end
+
+    Foo.extend(EE::Foo)
+    ^^^^^^^^^^^^^^^^^^^ EE modules must be injected using `include_if_ee`, `extend_if_ee`, or `prepend_if_ee`
+    SOURCE
+  end
+
+  it 'disallows the use of include to inject an EE module' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.include(EE::Foo)
+    ^^^^^^^^^^^^^^^^^^^^ EE modules must be injected using `include_if_ee`, `extend_if_ee`, or `prepend_if_ee`
+    SOURCE
+  end
+
+  it 'disallows the use of prepend_if_ee without a String' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.prepend_if_ee(EE::Foo)
+                      ^^^^^^^ EE modules to inject must be specified as a String
+    SOURCE
+  end
+
+  it 'disallows the use of include_if_ee without a String' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.include_if_ee(EE::Foo)
+                      ^^^^^^^ EE modules to inject must be specified as a String
+    SOURCE
+  end
+
+  it 'disallows the use of extend_if_ee without a String' do
+    expect_offense(<<~SOURCE)
+    class Foo
+    end
+
+    Foo.extend_if_ee(EE::Foo)
+                     ^^^^^^^ EE modules to inject must be specified as a String
     SOURCE
   end
 end