lib/ridgepole/dsl_parser.rb in ridgepole-0.6.4.beta4 vs lib/ridgepole/dsl_parser.rb in ridgepole-0.6.4.beta5

@@ -7,12 +7,14 @@
     end
 
     class TableDefinition
       attr_reader :__definition
 
-      def initialize
+      def initialize(table_name, base)
         @__definition = {}
+        @table_name = table_name
+        @base = base
       end
 
       def column(name, type, options = {})
         name = name.to_s
 
@@ -70,24 +72,33 @@
           column_names = args
           column_names.each {|name| column(name, column_type, options) }
         end
       end
 
+      def index(name, options = {})
+        @base.add_index(@table_name, name, options)
+      end
+
       def timestamps(*args)
         options = {:null => false}.merge(args.extract_options!)
         column(:created_at, :datetime, options)
         column(:updated_at, :datetime, options)
       end
 
       def references(*args)
         options = args.extract_options!
         polymorphic = options.delete(:polymorphic)
+        index_options = options.delete(:index)
         type = options.delete(:type) || :integer
 
         args.each do |col|
           column("#{col}_id", type, options)
           column("#{col}_type", :string, polymorphic.is_a?(Hash) ? polymorphic : options) if polymorphic
+          if index_options
+            index("#{col}_id", index_options.is_a?(Hash) ? index_options : {})
+            index("#{col}_type", index_options.is_a?(Hash) ? index_options : {}) if polymorphic
+          end
         end
       end
       alias :belongs_to :references
     end
 
@@ -112,10 +123,10 @@
       [ctx.__definition, ctx.__execute]
     end
 
     def create_table(table_name, options = {})
       table_name = table_name.to_s
-      table_definition = TableDefinition.new
+      table_definition = TableDefinition.new(table_name, self)
 
       [:primary_key].each do |key|
         options[key] = options[key].to_s if options[key]
       end