From fcc33642c8271c9adb9a7d9b17c5e619094c58de Mon Sep 17 00:00:00 2001
From: Mark Chao <mchao@gitlab.com>
Date: Tue, 2 Jul 2019 13:26:26 +0800
Subject: [PATCH] ES: Forward calls from ActiveRecord to proxies

Callbacks are copied from elasticsearch_rails
---
 .../elastic/application_versioned_search.rb   | 73 +++++++++++++++++++
 ee/lib/elastic/multi_version_class_proxy.rb   | 26 +++++++
 .../elastic/multi_version_instance_proxy.rb   | 19 +++++
 ee/lib/elastic/multi_version_util.rb          | 73 +++++++++++++++++++
 .../elastic/multi_version_class_proxy_spec.rb | 47 ++++++++++++
 .../multi_version_instance_proxy_spec.rb      | 49 +++++++++++++
 6 files changed, 287 insertions(+)
 create mode 100644 ee/app/models/concerns/elastic/application_versioned_search.rb
 create mode 100644 ee/lib/elastic/multi_version_class_proxy.rb
 create mode 100644 ee/lib/elastic/multi_version_instance_proxy.rb
 create mode 100644 ee/lib/elastic/multi_version_util.rb
 create mode 100644 ee/spec/lib/elastic/multi_version_class_proxy_spec.rb
 create mode 100644 ee/spec/lib/elastic/multi_version_instance_proxy_spec.rb

diff --git a/ee/app/models/concerns/elastic/application_versioned_search.rb b/ee/app/models/concerns/elastic/application_versioned_search.rb
new file mode 100644
index 00000000000..d5076ed7ce5
--- /dev/null
+++ b/ee/app/models/concerns/elastic/application_versioned_search.rb
@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+module Elastic
+  module ApplicationVersionedSearch
+    extend ActiveSupport::Concern
+
+    FORWARDABLE_INSTANCE_METHODS = [:es_id, :es_parent].freeze
+    FORWARDABLE_CLASS_METHODS = [:elastic_search, :es_import, :nested?, :es_type, :index_name, :document_type].freeze
+
+    def __elasticsearch__(&block)
+      @__elasticsearch__ ||= ::Elastic::MultiVersionInstanceProxy.new(self)
+    end
+
+    # Should be overridden in the models where some records should be skipped
+    def searchable?
+      self.use_elasticsearch?
+    end
+
+    def use_elasticsearch?
+      self.project&.use_elasticsearch?
+    end
+
+    def es_type
+      self.class.es_type
+    end
+
+    included do
+      delegate(*FORWARDABLE_INSTANCE_METHODS, to: :__elasticsearch__)
+
+      class << self
+        delegate(*FORWARDABLE_CLASS_METHODS, to: :__elasticsearch__)
+      end
+
+      # Add to the registry if it's a class (and not in intermediate module)
+      Elasticsearch::Model::Registry.add(self) if self.is_a?(Class)
+
+      after_commit on: :create do
+        if Gitlab::CurrentSettings.elasticsearch_indexing? && self.searchable?
+          ElasticIndexerWorker.perform_async(:index, self.class.to_s, self.id, self.es_id)
+        end
+      end
+
+      after_commit on: :update do
+        if Gitlab::CurrentSettings.elasticsearch_indexing? && self.searchable?
+          ElasticIndexerWorker.perform_async(
+            :update,
+            self.class.to_s,
+            self.id,
+            self.es_id,
+            changed_fields: self.previous_changes.keys
+          )
+        end
+      end
+
+      after_commit on: :destroy do
+        if Gitlab::CurrentSettings.elasticsearch_indexing? && self.searchable?
+          ElasticIndexerWorker.perform_async(
+            :delete,
+            self.class.to_s,
+            self.id,
+            self.es_id,
+            es_parent: self.es_parent
+          )
+        end
+      end
+    end
+
+    class_methods do
+      def __elasticsearch__
+        @__elasticsearch__ ||= ::Elastic::MultiVersionClassProxy.new(self)
+      end
+    end
+  end
+end
diff --git a/ee/lib/elastic/multi_version_class_proxy.rb b/ee/lib/elastic/multi_version_class_proxy.rb
new file mode 100644
index 00000000000..a6b2ae2ac7d
--- /dev/null
+++ b/ee/lib/elastic/multi_version_class_proxy.rb
@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+# Agnostic proxy to decide which version of elastic_target to use based on method being reads or writes
+module Elastic
+  class MultiVersionClassProxy
+    include MultiVersionUtil
+
+    def initialize(data_target)
+      @data_target = data_target
+      @data_class = get_data_class(data_target)
+
+      generate_forwarding
+    end
+
+    def version(version)
+      super.tap do |elastic_target|
+        elastic_target.extend Elasticsearch::Model::Importing::ClassMethods
+        elastic_target.extend Elasticsearch::Model::Adapter.from_class(@data_class).importing_mixin
+      end
+    end
+
+    def proxy_class_name
+      "#{@data_class.name}ClassProxy"
+    end
+  end
+end
diff --git a/ee/lib/elastic/multi_version_instance_proxy.rb b/ee/lib/elastic/multi_version_instance_proxy.rb
new file mode 100644
index 00000000000..f93604cc358
--- /dev/null
+++ b/ee/lib/elastic/multi_version_instance_proxy.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+# Agnostic proxy to decide which version of elastic_target to use based on method being reads or writes
+module Elastic
+  class MultiVersionInstanceProxy
+    include MultiVersionUtil
+
+    def initialize(data_target)
+      @data_target = data_target
+      @data_class = get_data_class(data_target.class)
+
+      generate_forwarding
+    end
+
+    def proxy_class_name
+      "#{@data_class.name}InstanceProxy"
+    end
+  end
+end
diff --git a/ee/lib/elastic/multi_version_util.rb b/ee/lib/elastic/multi_version_util.rb
new file mode 100644
index 00000000000..8664a698225
--- /dev/null
+++ b/ee/lib/elastic/multi_version_util.rb
@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module Elastic
+  module MultiVersionUtil
+    extend ActiveSupport::Concern
+    include Gitlab::Utils::StrongMemoize
+
+    attr_reader :data_class, :data_target
+
+    # @params version [String, Module] can be a string "V12p1" or module (Elastic::V12p1)
+    def version(version)
+      version = Elastic.const_get(version) if version.is_a?(String)
+      version.const_get(proxy_class_name).new(data_target)
+    end
+
+    private
+
+    # TODO: load from db table https://gitlab.com/gitlab-org/gitlab-ee/issues/12555
+    def elastic_reading_target
+      strong_memoize(:elastic_reading_target) do
+        version('V12p1')
+      end
+    end
+
+    # TODO: load from db table https://gitlab.com/gitlab-org/gitlab-ee/issues/12555
+    def elastic_writing_targets
+      strong_memoize(:elastic_writing_targets) do
+        [elastic_reading_target]
+      end
+    end
+
+    def get_data_class(klass)
+      klass < ActiveRecord::Base ? klass.base_class : klass
+    end
+
+    def generate_forwarding
+      write_methods = elastic_writing_targets.first.real_class.write_methods
+
+      write_methods.each do |method|
+        self.class.forward_write_method(method)
+      end
+
+      read_methods = elastic_reading_target.real_class.public_instance_methods
+      read_methods -= write_methods
+      read_methods -= self.class.instance_methods
+      read_methods.delete(:method_missing)
+
+      read_methods.each do |method|
+        self.class.forward_read_method(method)
+      end
+    end
+
+    class_methods do
+      def forward_read_method(method)
+        return if respond_to?(method)
+
+        delegate method, to: :elastic_reading_target
+      end
+
+      def forward_write_method(method)
+        return if respond_to?(method)
+
+        define_method(method) do |*args|
+          responses = elastic_writing_targets.map do |elastic_target|
+            elastic_target.public_send(method, *args) # rubocop:disable GitlabSecurity/PublicSend
+          end
+
+          responses.find { |response| response['_shards']['successful'] == 0 } || responses.last
+        end
+      end
+    end
+  end
+end
diff --git a/ee/spec/lib/elastic/multi_version_class_proxy_spec.rb b/ee/spec/lib/elastic/multi_version_class_proxy_spec.rb
new file mode 100644
index 00000000000..4cf1ce1a5c9
--- /dev/null
+++ b/ee/spec/lib/elastic/multi_version_class_proxy_spec.rb
@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Elastic::MultiVersionClassProxy do
+  subject { described_class.new(ProjectSnippet) }
+
+  describe '#version' do
+    it 'returns class proxy in specified version' do
+      result = subject.version('V12p1')
+
+      expect(result).to be_a(Elastic::V12p1::SnippetClassProxy)
+      expect(result.target).to eq(ProjectSnippet)
+    end
+  end
+
+  describe 'method forwarding' do
+    let(:old_target) { double(:old_target) }
+    let(:new_target) { double(:new_target) }
+    let(:response) do
+      { "_index" => "gitlab-test", "_type" => "doc", "_id" => "snippet_1", "_version" => 3, "result" => "updated", "_shards" => { "total" => 2, "successful" => 1, "failed" => 0 }, "created" => false }
+    end
+
+    before do
+      allow(subject).to receive(:elastic_reading_target).and_return(old_target)
+      allow(subject).to receive(:elastic_writing_targets).and_return([old_target, new_target])
+    end
+
+    it 'forwards write methods to all targets' do
+      Elastic::V12p1::SnippetClassProxy.write_methods.each do |method|
+        expect(new_target).to receive(method).and_return(response)
+        expect(old_target).to receive(method).and_return(response)
+
+        subject.public_send(method)
+      end
+    end
+
+    it 'forwards read methods to only reading target' do
+      expect(old_target).to receive(:search)
+      expect(new_target).not_to receive(:search)
+
+      subject.search
+
+      expect(subject).not_to respond_to(:method_missing)
+    end
+  end
+end
diff --git a/ee/spec/lib/elastic/multi_version_instance_proxy_spec.rb b/ee/spec/lib/elastic/multi_version_instance_proxy_spec.rb
new file mode 100644
index 00000000000..47fe35d9f9c
--- /dev/null
+++ b/ee/spec/lib/elastic/multi_version_instance_proxy_spec.rb
@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Elastic::MultiVersionInstanceProxy do
+  let(:snippet) { create(:project_snippet) }
+
+  subject { described_class.new(snippet) }
+
+  describe '#version' do
+    it 'returns instance proxy in specified version' do
+      result = subject.version('V12p1')
+
+      expect(result).to be_a(Elastic::V12p1::SnippetInstanceProxy)
+      expect(result.target).to eq(snippet)
+    end
+  end
+
+  describe 'method forwarding' do
+    let(:old_target) { double(:old_target) }
+    let(:new_target) { double(:new_target) }
+    let(:response) do
+      { "_index" => "gitlab-test", "_type" => "doc", "_id" => "snippet_1", "_version" => 3, "result" => "updated", "_shards" => { "total" => 2, "successful" => 1, "failed" => 0 }, "created" => false }
+    end
+
+    before do
+      allow(subject).to receive(:elastic_reading_target).and_return(old_target)
+      allow(subject).to receive(:elastic_writing_targets).and_return([old_target, new_target])
+    end
+
+    it 'forwards write methods to all targets' do
+      Elastic::V12p1::SnippetInstanceProxy.write_methods.each do |method|
+        expect(new_target).to receive(method).and_return(response)
+        expect(old_target).to receive(method).and_return(response)
+
+        subject.public_send(method)
+      end
+    end
+
+    it 'forwards read methods to only reading target' do
+      expect(old_target).to receive(:as_indexed_json)
+      expect(new_target).not_to receive(:as_indexed_json)
+
+      subject.as_indexed_json
+
+      expect(subject).not_to respond_to(:method_missing)
+    end
+  end
+end
-- 
2.30.9