diff --git a/Rakefile b/Rakefile
index cb6b440..4fcea63 100644
--- a/Rakefile
+++ b/Rakefile
@@ -21,14 +21,18 @@ namespace :db do
   desc "Create the database"
   task :create do
     ActiveRecord::Base.establish_connection(admin_database_config)
-    ActiveRecord::Base.connection.create_database(database_config.fetch(:database))
     puts "Database created."
   end
 
   desc "Migrate the database"
   task :migrate do
     ActiveRecord::Base.establish_connection(database_config)
-    ActiveRecord::Migrator.migrate(migration_path)
+    if defined?(ActiveRecord::MigrationContext)
+      migration_context = ActiveRecord::MigrationContext.new(migration_path, ActiveRecord::SchemaMigration)
+      migration_context.migrate
+    else
+      ActiveRecord::Migrator.migrate(migration_path)
+    end
     Rake::Task["db:schema"].invoke
     puts "Database migrated."
   end
@@ -36,7 +40,8 @@ namespace :db do
   desc "Drop the database"
   task :drop do
     ActiveRecord::Base.establish_connection(admin_database_config)
-    ActiveRecord::Base.connection.drop_database(database_config.fetch(:database))
+    db_file = database_config.fetch(:database)
+    File.delete(db_file) if File.exist?(db_file)
     puts "Database deleted."
   end
 
diff --git a/lib/polymorphic_integer_type.rb b/lib/polymorphic_integer_type.rb
index b5bf33d..5dbcb40 100644
--- a/lib/polymorphic_integer_type.rb
+++ b/lib/polymorphic_integer_type.rb
@@ -4,6 +4,7 @@
 require "polymorphic_integer_type/extensions"
 require "polymorphic_integer_type/mapping"
 require "polymorphic_integer_type/module_generator"
+require "polymorphic_integer_type/belongs_to_valid_options_extension"
 require "polymorphic_integer_type/belongs_to_polymorphic_association_extension"
 require "polymorphic_integer_type/activerecord_5_0_0/polymorphic_array_value_extension"
 require "polymorphic_integer_type/polymorphic_foreign_association_extension"
diff --git a/lib/polymorphic_integer_type/belongs_to_polymorphic_association_extension.rb b/lib/polymorphic_integer_type/belongs_to_polymorphic_association_extension.rb
index 8840fc2..8d4b7c8 100644
--- a/lib/polymorphic_integer_type/belongs_to_polymorphic_association_extension.rb
+++ b/lib/polymorphic_integer_type/belongs_to_polymorphic_association_extension.rb
@@ -1,19 +1,25 @@
-module ActiveRecord
-  module Associations
-    class BelongsToPolymorphicAssociation < BelongsToAssociation
-      private
+module PolymorphicIntegerType
+  module BelongsToPolymorphicAssociationExtension
+    private
 
-      if Gem::Version.new(ActiveRecord::VERSION::STRING) < Gem::Version.new("6.1")
-        def replace_keys(record)
-          super
+    if Gem::Version.new(ActiveRecord::VERSION::STRING) < Gem::Version.new("6.1")
+      def replace_keys(record)
+        super
+
+        if reflection.options[:integer_type] || reflection.options[:polymorphic].is_a?(Hash)
           owner[reflection.foreign_type] = record.class.base_class unless record.nil?
         end
-      elsif
-        def replace_keys(record, force: false)
-          super
+      end
+    else
+      def replace_keys(record, force: false)
+        super
+
+        if reflection.options[:integer_type] || reflection.options[:polymorphic].is_a?(Hash)
           owner[reflection.foreign_type] = record.class.base_class unless record.nil?
         end
       end
     end
   end
 end
+
+ActiveRecord::Associations::BelongsToPolymorphicAssociation.prepend(PolymorphicIntegerType::BelongsToPolymorphicAssociationExtension)
diff --git a/lib/polymorphic_integer_type/belongs_to_valid_options_extension.rb b/lib/polymorphic_integer_type/belongs_to_valid_options_extension.rb
new file mode 100644
index 0000000..67dd805
--- /dev/null
+++ b/lib/polymorphic_integer_type/belongs_to_valid_options_extension.rb
@@ -0,0 +1,11 @@
+module PolymorphicIntegerType
+  module BelongsToValidOptionsExtension
+    def valid_options(options)
+      valid = super
+      valid << :integer_type if options[:polymorphic]
+      valid
+    end
+  end
+end
+
+ActiveRecord::Associations::Builder::BelongsTo.singleton_class.prepend(PolymorphicIntegerType::BelongsToValidOptionsExtension)
diff --git a/lib/polymorphic_integer_type/extensions.rb b/lib/polymorphic_integer_type/extensions.rb
index c6c1389..47f3611 100644
--- a/lib/polymorphic_integer_type/extensions.rb
+++ b/lib/polymorphic_integer_type/extensions.rb
@@ -9,15 +9,14 @@ module ClassMethods
 
       def belongs_to(name, scope = nil, **options)
         options = scope if scope.kind_of? Hash
-        integer_type = options.delete :integer_type
         super
-        if options[:polymorphic] && (integer_type || options[:polymorphic].is_a?(Hash))
+        if options[:polymorphic] && (options[:integer_type] || options[:polymorphic].is_a?(Hash))
           mapping =
-            case integer_type
+            case options[:integer_type]
             when true then PolymorphicIntegerType::Mapping[name]
             when nil then options[:polymorphic]
             else
-              raise ArgumentError, "Unknown integer_type value: #{integer_type.inspect}"
+              raise ArgumentError, "Unknown integer_type value: #{options[:integer_type].inspect}"
             end.dup
 
           foreign_type = reflections[name.to_s].foreign_type
diff --git a/spec/polymorphic_integer_type_spec.rb b/spec/polymorphic_integer_type_spec.rb
index 6e426cf..08a25c0 100644
--- a/spec/polymorphic_integer_type_spec.rb
+++ b/spec/polymorphic_integer_type_spec.rb
@@ -179,6 +179,7 @@
       end
     end
   end
+
   context "When a link is given polymorphic record" do
     let(:link) { Link.create(source: source) }
     let(:source) { cat }
@@ -193,6 +194,17 @@
     end
   end
 
+  context "When a link is given polymorphic record via a non-integer type association" do
+    let(:link) { Link.create(legacy_source: source) }
+    let(:source) { Place.create(name: "Main Street") }
+
+    it "appropriately sets the source_id and source_type to the polymorphic_name" do
+      expect(link.legacy_source_id).to eql source.id
+      expect(link.legacy_source_type).to eql Place.polymorphic_name
+      expect(link.legacy_source).to eql source
+    end
+  end
+
   context "When a link is given polymorphic id and type" do
     let(:link) { Link.create(source_id: source.id, source_type: source.class.to_s) }
     let(:source) { cat }
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index fdcbd3d..c802cb9 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -13,6 +13,7 @@
 require 'support/profile'
 require 'support/profile_history'
 require 'support/namespaced_activity'
+require 'support/place'
 require 'byebug'
 require 'pry'
 
diff --git a/spec/support/link.rb b/spec/support/link.rb
index f49bdd1..00deb6f 100644
--- a/spec/support/link.rb
+++ b/spec/support/link.rb
@@ -1,6 +1,12 @@
 class Link < ActiveRecord::Base
   include PolymorphicIntegerType::Extensions
 
+  def self.polymorphic_class_for(name)
+    name == 'Country' ? Place : super
+  end
+
   belongs_to :source, polymorphic: true, integer_type: true
   belongs_to :target, polymorphic: true, integer_type: true
+
+  belongs_to :legacy_source, polymorphic: true
 end
diff --git a/spec/support/migrations/11_add_legacy_source_reference_to_links_table.rb b/spec/support/migrations/11_add_legacy_source_reference_to_links_table.rb
new file mode 100644
index 0000000..5f67911
--- /dev/null
+++ b/spec/support/migrations/11_add_legacy_source_reference_to_links_table.rb
@@ -0,0 +1,7 @@
+class AddLegacySourceReferenceToLinksTable < ActiveRecord::Migration[5.0]
+  def change
+    add_reference :links, :legacy_source, polymorphic: true
+  end
+end
+
+
diff --git a/spec/support/place.rb b/spec/support/place.rb
new file mode 100644
index 0000000..fc6b981
--- /dev/null
+++ b/spec/support/place.rb
@@ -0,0 +1,7 @@
+class Place < ActiveRecord::Base
+  has_many :source_links, as: :target, inverse_of: :legacy_source, class_name: "Link"
+
+  def self.polymorphic_name
+    'Country'
+  end
+end