lib/ridgepole/migration_ext.rb in ridgepole-0.2.0 vs lib/ridgepole/migration_ext.rb in ridgepole-0.2.1

@@ -1,9 +1,57 @@
 require 'active_record/migration'
 
 class ActiveRecord::Migration
+  cattr_accessor :time_recorder
+
   def write_with_logging(text = '')
     logger = Ridgepole::Logger.instance
     logger.info(text)
+    parse_text(text)
   end
   alias_method_chain :write, :logging
+
+  def parse_text(text)
+    return unless self.time_recorder
+
+    case text
+    when /\A--\s+(.+)\Z/
+      self.time_recorder.add_key($1)
+    when /\A\s+->\s+(\d+\.\d+)s\Z/
+      self.time_recorder.add_value($1.to_f)
+    end
+  end
+
+  def self.record_time
+    result = nil
+
+    begin
+      self.time_recorder = TimeRecorder.new
+      yield
+      result = self.time_recorder.result
+    ensure
+      self.time_recorder = nil
+    end
+
+    return result
+  end
+
+  class TimeRecorder
+    attr_reader :result
+
+    def initialize
+      @result = {}
+    end
+
+    def add_key(key)
+      @key = key
+    end
+
+    def add_value(value)
+      if @key
+        @result[@key] = value
+      end
+
+      @key = nil
+    end
+  end
 end