lib/elasticity/index_config.rb in es-elasticity-0.5.2 vs lib/elasticity/index_config.rb in es-elasticity-0.6.0

@@ -1,11 +1,14 @@
 module Elasticity
   class IndexConfig
-    ATTRS = [:index_base_name, :document_type, :mapping, :strategy].freeze
+    ATTRS = [:index_base_name, :document_type, :mapping, :strategy, :subclasses].freeze
+    VALIDATABLE_ATTRS = [:index_base_name, :document_type, :strategy].freeze
+
     attr_accessor *ATTRS
 
-    def initialize(elasticity_config)
+    def initialize(elasticity_config, default_document_type)
+      @document_type = default_document_type
       @elasticity_config = elasticity_config
       yield(self)
       validate!
     end
 
@@ -18,13 +21,22 @@
     def client
       @elasticity_config.client
     end
 
     def definition
-      { settings: @elasticity_config.settings, mappings: { @document_type => @mapping } }
+      return @definition if defined?(@definition)
+      @definition = { settings: @elasticity_config.settings, mappings: { @document_type => @mapping || {} } }
+      subclasses.each do |doc_type, subclass|
+        @definition[:mappings][doc_type] = subclass.constantize.mapping
+      end if subclasses.present?
+      @definition
     end
 
+    def document_types
+      @document_types ||= definition[:mappings].collect { |doc_type, mapping| doc_type }
+    end
+
     def fq_index_base_name
       return @fq_index_base_name if defined?(@fq_index_base_name)
 
       if namespace = @elasticity_config.namespace
         @fq_index_base_name = "#{namespace}_#{@index_base_name}"
@@ -40,10 +52,10 @@
     end
 
     private
 
     def validate!
-      ATTRS.each do |attr|
+      VALIDATABLE_ATTRS.each do |attr|
         raise "#{attr} is not set" if public_send(attr).nil?
       end
     end
   end
 end