ruby 用于将多个模型并置为单个索引中的类型的Rake任务。

Posted 2021-05-16

# A Rake tasks to facilitate importing data from your models into a common Elasticsearch index.
#
# All models should declare a common index_name, and a document_type: 
#
#  class Article
#    include Elasticsearch::Model
#
#    index_name 'app_scoped_index'
#    document_type 'articles'
#
#    mappings do
#       ...
#    end
#  end
#
#

STDOUT.sync = true
STDERR.sync = true

begin; require 'ansi/progressbar'; rescue LoadError; end

namespace :elasticsearch do

  task :import => 'import:model'

  namespace :import do
   
    desc <<-DESC.gsub(/    /, '')
      Import all mappings from `app/models` (or use DIR environment variable) into a single index.

      All classes should declare a common `index_name`:

      class Article
        include Elasticsearch::Model
      
        index_name 'app_scoped_index'
      
        mappings do
          ...
        end
      end

      Usage: 
        $ rake environment elasticsearch:import:combined DIR=app/models
    DESC
    task :combined do
      dir = ENV['DIR'].to_s != '' ? ENV['DIR'] : Rails.root.join("app/models")
      
      puts "[IMPORT] Loading models from: #{dir} into a single index."

      all_mappings = {}
      index_klass = nil

      Dir.glob(File.join("#{dir}/**/*.rb")).each do |path|
        model_filename = path[/#{Regexp.escape(dir.to_s)}\/([^\.]+).rb/, 1]

        next if model_filename.match(/^concerns\//i) # Skip concerns/ folder

        begin
          klass = model_filename.camelize.constantize
        rescue NameError
          require(path) ? retry : raise(RuntimeError, "Cannot load class '#{klass}'")
        end

        # Skip if the class doesn't have Elasticsearch integration
        next unless klass.respond_to?(:__elasticsearch__)
        next unless klass.respond_to?(:mappings)

        puts "[IMPORT] Processing mappings for: #{klass}..."

        index_klass = klass
        all_mappings.merge! klass.mappings.to_hash
      end

      
      ## Create the combined index
      index_klass.__elasticsearch__.client.indices.create(
        { 
          index: index_klass.index_name, 
          body: { 
            mappings: all_mappings 
          }  
        })    

      ## Import data into the newly created index
      Rake::Task["elasticsearch:import:all"].invoke

      puts
    end

  end

end