diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb index 3cb7373a970ca41250c147db9b3aa02816402245..5f14d95ffed159a68812c1059836f3964149a533 100644 --- a/app/controllers/application_controller.rb +++ b/app/controllers/application_controller.rb @@ -556,4 +556,4 @@ def required_signup_info end end -ApplicationController.prepend_if_ee('EE::ApplicationController') +ApplicationController.prepend_ee_mod diff --git a/app/models/vulnerability.rb b/app/models/vulnerability.rb index ab29afd0d08c170b98cc111f26754e079deda049..17e3e285b40579144d0fd984e8fe1cce1462bdae 100644 --- a/app/models/vulnerability.rb +++ b/app/models/vulnerability.rb @@ -25,4 +25,4 @@ def self.reference_postfix_escaped end end -Vulnerability.prepend_if_ee('EE::Vulnerability') +Vulnerability.prepend_ee_mod diff --git a/config/initializers/0_inject_enterprise_edition_module.rb b/config/initializers/0_inject_enterprise_edition_module.rb index b3ebb44ef255f7b9009435a84baaaf0114537d28..7478727f869e80f8403208663a9f4ba539e74a25 100644 --- a/config/initializers/0_inject_enterprise_edition_module.rb +++ b/config/initializers/0_inject_enterprise_edition_module.rb @@ -6,12 +6,7 @@ module InjectEnterpriseEditionModule def prepend_if_ee(constant, with_descendants: false) return unless Gitlab.ee? - ee_module = constant.constantize - prepend(ee_module) - - if with_descendants - descendants.each { |descendant| descendant.prepend(ee_module) } - end + prepend_module(constant.constantize, with_descendants) end def extend_if_ee(constant) @@ -21,6 +16,34 @@ def extend_if_ee(constant) def include_if_ee(constant) include(constant.constantize) if Gitlab.ee? end + + def prepend_ee_mod(with_descendants: false) + return unless Gitlab.ee? + + prepend_module(ee_module, with_descendants) + end + + def extend_ee_mod + extend(ee_module) if Gitlab.ee? + end + + def include_ee_mod + include(ee_module) if Gitlab.ee? + end + + private + + def prepend_module(mod, with_descendants) + prepend(mod) + + if with_descendants + descendants.each { |descendant| descendant.prepend(mod) } + end + end + + def ee_module + ::EE.const_get(name, false) + end end Module.prepend(InjectEnterpriseEditionModule) diff --git a/doc/development/ee_features.md b/doc/development/ee_features.md index 5be601187ca7bf8e5c754da212eac20582f573e7..c1d77abf08ea847dc43c24d865d45a47b5bded30 100644 --- a/doc/development/ee_features.md +++ b/doc/development/ee_features.md @@ -92,12 +92,36 @@ class User < ActiveRecord::Base # ... lots of code here ... end -User.prepend_if_ee('EE::User') +User.prepend_ee_mod ``` 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. +`prepend_ee_mod`, `extend_ee_mod`, or `include_ee_mod`. These methods will try to +find the relevant EE module by the name of the receiver module, for example; + +```ruby +module Vulnerabilities + class Finding + #... + end +end + +Vulnerabilities::Finding.prepend_ee_mod +``` + +will prepend the module named `::EE::Vulnerabilities::Finding`. + +If the extending module does not follow this naming convention, you can also provide the module name +by using `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, like so; + +```ruby +class User + #... +end + +User.prepend_if_ee('::EE::UserExtension') +``` 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 diff --git a/lib/api/api.rb b/lib/api/api.rb index 95e37f2c418167102cea62adbd623e957030e6b6..9ce580d0b9873f5cd5797a60a92fbb81afbb5d6f 100644 --- a/lib/api/api.rb +++ b/lib/api/api.rb @@ -313,4 +313,4 @@ class API < ::API::Base end end -API::API.prepend_if_ee('::EE::API::API') +API::API.prepend_ee_mod