summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--lib/gitlab/cache/request_store_wrap.rb60
-rw-r--r--lib/gitlab/user_access.rb10
-rw-r--r--spec/serializers/pipeline_serializer_spec.rb2
3 files changed, 69 insertions, 3 deletions
diff --git a/lib/gitlab/cache/request_store_wrap.rb b/lib/gitlab/cache/request_store_wrap.rb
new file mode 100644
index 00000000000..3e0a5f06b53
--- /dev/null
+++ b/lib/gitlab/cache/request_store_wrap.rb
@@ -0,0 +1,60 @@
+module Gitlab
+ module Cache
+ # This module provides a simple way to cache values in RequestStore,
+ # and the cache key would be based on the class name, method name,
+ # customized instance level values, and arguments.
+ #
+ # A simple example:
+ #
+ # class UserAccess
+ # extend Gitlab::Cache::RequestStoreWrap
+ #
+ # request_store_wrap_key do
+ # [user.id, project.id]
+ # end
+ #
+ # request_store_wrap def can_push_to_branch?(ref)
+ # # ...
+ # end
+ # end
+ #
+ # This way, the result of `can_push_to_branch?` would be cached in
+ # `RequestStore.store` based on the cache key.
+ module RequestStoreWrap
+ def self.extended(klass)
+ return if klass < self
+
+ extension = Module.new
+ klass.const_set(:RequestStoreWrapExtension, extension)
+ klass.prepend(extension)
+ end
+
+ def request_store_wrap_key(&block)
+ if block_given?
+ @request_store_wrap_key = block
+ else
+ @request_store_wrap_key
+ end
+ end
+
+ def request_store_wrap(method_name)
+ const_get(:RequestStoreWrapExtension)
+ .send(:define_method, method_name) do |*args|
+ return super(*args) unless RequestStore.active?
+
+ klass = self.class
+ key = [klass.name,
+ method_name,
+ *instance_exec(&klass.request_store_wrap_key),
+ *args].join(':')
+
+ if RequestStore.store.key?(key)
+ RequestStore.store[key]
+ else
+ RequestStore.store[key] = super(*args)
+ end
+ end
+ end
+ end
+ end
+end
diff --git a/lib/gitlab/user_access.rb b/lib/gitlab/user_access.rb
index bb05c474fa2..d8b043f5021 100644
--- a/lib/gitlab/user_access.rb
+++ b/lib/gitlab/user_access.rb
@@ -1,5 +1,11 @@
module Gitlab
class UserAccess
+ extend Gitlab::Cache::RequestStoreWrap
+
+ request_store_wrap_key do
+ [user&.id, project&.id]
+ end
+
attr_reader :user, :project
def initialize(user, project: nil)
@@ -52,7 +58,7 @@ module Gitlab
can_push_to_branch?(ref) || can_merge_to_branch?(ref)
end
- def can_push_to_branch?(ref)
+ request_store_wrap def can_push_to_branch?(ref)
return false unless can_access_git?
if ProtectedBranch.protected?(project, ref)
@@ -64,7 +70,7 @@ module Gitlab
end
end
- def can_merge_to_branch?(ref)
+ request_store_wrap def can_merge_to_branch?(ref)
return false unless can_access_git?
if ProtectedBranch.protected?(project, ref)
diff --git a/spec/serializers/pipeline_serializer_spec.rb b/spec/serializers/pipeline_serializer_spec.rb
index 44813656aff..8dc666586c7 100644
--- a/spec/serializers/pipeline_serializer_spec.rb
+++ b/spec/serializers/pipeline_serializer_spec.rb
@@ -110,7 +110,7 @@ describe PipelineSerializer do
it 'verifies number of queries', :request_store do
recorded = ActiveRecord::QueryRecorder.new { subject }
- expect(recorded.count).to be_within(1).of(57)
+ expect(recorded.count).to be_within(1).of(59)
expect(recorded.cached_count).to eq(0)
end