Merge pull request #15248 from sgrif/sg-additional-type-map-args

Allow additional arguments to be used during type map lookups

5 files changed, 29 insertions, 12 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb b/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
index 539ba38c4a..c04a1d7178 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
@@ -178,7 +178,7 @@ module ActiveRecord
         def columns(table_name)
           # Limit, precision, and scale are all handled by the superclass.
           column_definitions(table_name).map do |column_name, type, default, notnull, oid, fmod|
-            oid = get_oid_type(oid.to_i, fmod.to_i, column_name)
+            oid = get_oid_type(oid.to_i, fmod.to_i, column_name, type)
             PostgreSQLColumn.new(column_name, default, oid, type, notnull == 'f')
           end
         end
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
index ed3e884455..811c5b5c7a 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -538,12 +538,12 @@ module ActiveRecord
 
       private
 
-        def get_oid_type(oid, fmod, column_name)
+        def get_oid_type(oid, fmod, column_name, sql_type = '')
           if !type_map.key?(oid)
             initialize_type_map(type_map, [oid])
           end
 
-          type_map.fetch(normalize_oid_type(oid, fmod)) {
+          type_map.fetch(normalize_oid_type(oid, fmod), sql_type) {
             warn "unknown OID #{oid}: failed to recognize type of '#{column_name}'. It will be treated as String."
             Type::Value.new.tap do |cast_type|
               type_map.register_type(oid, cast_type)
diff --git a/activerecord/lib/active_record/connection_adapters/type/hash_lookup_type_map.rb b/activerecord/lib/active_record/connection_adapters/type/hash_lookup_type_map.rb
index 8503d3ea1b..bb1abc77ff 100644
--- a/activerecord/lib/active_record/connection_adapters/type/hash_lookup_type_map.rb
+++ b/activerecord/lib/active_record/connection_adapters/type/hash_lookup_type_map.rb
@@ -4,16 +4,16 @@ module ActiveRecord
       class HashLookupTypeMap < TypeMap # :nodoc:
         delegate :key?, to: :@mapping
 
-        def lookup(type)
-          @mapping.fetch(type, proc { default_value }).call(type)
+        def lookup(type, *args)
+          @mapping.fetch(type, proc { default_value }).call(type, *args)
         end
 
-        def fetch(type, &block)
-          @mapping.fetch(type, block).call(type)
+        def fetch(type, *args, &block)
+          @mapping.fetch(type, block).call(type, *args)
         end
 
         def alias_type(type, alias_type)
-          register_type(type) { lookup(alias_type) }
+          register_type(type) { |_, *args| lookup(alias_type, *args) }
         end
       end
     end
diff --git a/activerecord/lib/active_record/connection_adapters/type/type_map.rb b/activerecord/lib/active_record/connection_adapters/type/type_map.rb
index d89171a820..48b8b51417 100644
--- a/activerecord/lib/active_record/connection_adapters/type/type_map.rb
+++ b/activerecord/lib/active_record/connection_adapters/type/type_map.rb
@@ -6,13 +6,13 @@ module ActiveRecord
           @mapping = {}
         end
 
-        def lookup(lookup_key)
+        def lookup(lookup_key, *args)
           matching_pair = @mapping.reverse_each.detect do |key, _|
             key === lookup_key
           end
 
           if matching_pair
-            matching_pair.last.call(lookup_key)
+            matching_pair.last.call(lookup_key, *args)
           else
             default_value
           end
@@ -29,9 +29,9 @@ module ActiveRecord
         end
 
         def alias_type(key, target_key)
-          register_type(key) do |sql_type|
+          register_type(key) do |sql_type, *args|
             metadata = sql_type[/\(.*\)/, 0]
-            lookup("#{target_key}#{metadata}")
+            lookup("#{target_key}#{metadata}", *args)
           end
         end
 
diff --git a/activerecord/test/cases/connection_adapters/type/type_map_test.rb b/activerecord/test/cases/connection_adapters/type/type_map_test.rb
index 4b4d9f6b0f..3abd7a276e 100644
--- a/activerecord/test/cases/connection_adapters/type/type_map_test.rb
+++ b/activerecord/test/cases/connection_adapters/type/type_map_test.rb
@@ -88,6 +88,23 @@ module ActiveRecord
           assert_equal mapping.lookup('varchar'), binary
         end
 
+        def test_additional_lookup_args
+          mapping = TypeMap.new
+
+          mapping.register_type(/varchar/i) do |type, limit|
+            if limit > 255
+              'text'
+            else
+              'string'
+            end
+          end
+          mapping.alias_type(/string/i, 'varchar')
+
+          assert_equal mapping.lookup('varchar', 200), 'string'
+          assert_equal mapping.lookup('varchar', 400), 'text'
+          assert_equal mapping.lookup('string', 400), 'text'
+        end
+
         def test_requires_value_or_block
           mapping = TypeMap.new