1 files changed, 39 insertions, 0 deletions
diff --git a/lib/gitlab/database/shared_model.rb b/lib/gitlab/database/shared_model.rb
new file mode 100644
index 00000000000..8f256758961
--- /dev/null
+++ b/lib/gitlab/database/shared_model.rb
@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class SharedModel < ActiveRecord::Base
+      self.abstract_class = true
+
+      class << self
+        def using_connection(connection)
+          raise 'cannot nest connection overrides for shared models' unless overriding_connection.nil?
+
+          self.overriding_connection = connection
+
+          yield
+        ensure
+          self.overriding_connection = nil
+        end
+
+        def connection
+          if connection = self.overriding_connection
+            connection
+          else
+            super
+          end
+        end
+
+        private
+
+        def overriding_connection
+          Thread.current[:overriding_connection]
+        end
+
+        def overriding_connection=(connection)
+          Thread.current[:overriding_connection] = connection
+        end
+      end
+    end
+  end
+end