lib/drafting.rb in drafting-0.3.0 vs lib/drafting.rb in drafting-0.3.1

@@ -5,5 +5,11 @@
 require "drafting/class_methods"
 require "drafting/instance_methods"
 require "drafting/draft"
 
 ActiveRecord::Base.extend Drafting::BaseClassMethods
+
+Drafting::MIGRATION_BASE_CLASS = if ActiveRecord::VERSION::MAJOR >= 5
+  ActiveRecord::Migration[5.0]
+else
+  ActiveRecord::Migration
+end