lib/ridgepole/migration_ext.rb in ridgepole-0.6.4.beta5 vs lib/ridgepole/migration_ext.rb in ridgepole-0.6.4.beta6

@@ -1,58 +1,69 @@
 require 'active_record/migration'
 
-class ActiveRecord::Migration
-  cattr_accessor :time_recorder
-  cattr_accessor :disable_logging
+module Ridgepole
+  module MigrationExt
+    def self.included(klass)
+      klass.class_eval do
+        cattr_accessor :time_recorder
+        cattr_accessor :disable_logging
+      end
+    end
 
-  def write_with_logging(text = '')
-    logger = Ridgepole::Logger.instance
-    logger.info(text) unless self.disable_logging
-    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)
+    def write(text = '')
+      logger = Ridgepole::Logger.instance
+      logger.info(text) unless self.disable_logging
+      parse_text(text)
     end
-  end
 
-  def self.record_time
-    result = nil
+    def parse_text(text)
+      return unless self.time_recorder
 
-    begin
-      self.time_recorder = TimeRecorder.new
-      yield
-      result = self.time_recorder.result
-    ensure
-      self.time_recorder = nil
+      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
 
-    return result
-  end
+    def self.record_time
+      result = nil
 
-  class TimeRecorder
-    attr_reader :result
+      begin
+        self.time_recorder = TimeRecorder.new
+        yield
+        result = self.time_recorder.result
+      ensure
+        self.time_recorder = nil
+      end
 
-    def initialize
-      @result = {}
+      result
     end
 
-    def add_key(key)
-      @key = key
-    end
+    class TimeRecorder
+      attr_reader :result
 
-    def add_value(value)
-      if @key
-        @result[@key] = value
+      def initialize
+        @result = {}
       end
 
-      @key = nil
+      def add_key(key)
+        @key = key
+      end
+
+      def add_value(value)
+        if @key
+          @result[@key] = value
+        end
+
+        @key = nil
+      end
     end
+  end
+end
+
+module ActiveRecord
+  class Migration
+    include Ridgepole::MigrationExt
   end
 end