Use the generic type map for PostgreSQL OID registrations

3 files changed, 31 insertions, 45 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/oid.rb b/activerecord/lib/active_record/connection_adapters/postgresql/oid.rb
index 94afc3aec5..741f4c9b4d 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/oid.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/oid.rb
@@ -276,43 +276,6 @@ This is not reliable and will be removed in the future.
           end
         end
 
-        class TypeMap
-          def initialize
-            @mapping = {}
-          end
-
-          def []=(oid, type)
-            @mapping[oid] = type
-          end
-
-          def [](oid)
-            @mapping[oid]
-          end
-
-          def clear
-            @mapping.clear
-          end
-
-          def key?(oid)
-            @mapping.key? oid
-          end
-
-          def fetch(ftype, fmod)
-            # The type for the numeric depends on the width of the field,
-            # so we'll do something special here.
-            #
-            # When dealing with decimal columns:
-            #
-            # places after decimal  = fmod - 4 & 0xffff
-            # places before decimal = (fmod - 4) >> 16 & 0xffff
-            if ftype == 1700 && (fmod - 4 & 0xffff).zero?
-              ftype = 23
-            end
-
-            @mapping.fetch(ftype) { |oid| yield oid, fmod }
-          end
-        end
-
         # This class uses the data from PostgreSQL pg_type table to build
         # the OID -> Type mapping.
         #   - OID is and integer representing the type.
@@ -349,19 +312,19 @@ This is not reliable and will be removed in the future.
           end
 
           def register_array_type(row)
-            if subtype = @store[row['typelem'].to_i]
+            if subtype = @store.lookup(row['typelem'].to_i)
               register row['oid'], OID::Array.new(subtype)
             end
           end
 
           def register_range_type(row)
-            if subtype = @store[row['rngsubtype'].to_i]
+            if subtype = @store.lookup(row['rngsubtype'].to_i)
               register row['oid'], OID::Range.new(subtype, row['typname'].to_sym)
             end
           end
 
           def register_domain_type(row)
-            if base_type = @store[row["typbasetype"].to_i]
+            if base_type = @store.lookup(row["typbasetype"].to_i)
               register row['oid'], base_type
             else
               warn "unknown base type (OID: #{row["typbasetype"]}) for domain #{row["typname"]}."
@@ -369,7 +332,7 @@ This is not reliable and will be removed in the future.
           end
 
           def register_composite_type(row)
-            if subtype = @store[row['typelem'].to_i]
+            if subtype = @store.lookup(row['typelem'].to_i)
               register row['oid'], OID::Vector.new(row['typdelim'], subtype)
             end
           end
@@ -380,7 +343,7 @@ This is not reliable and will be removed in the future.
             raise ArgumentError, "can't register nil type for OID #{oid}" if oid_type.nil?
             return if @store.key?(oid)
 
-            @store[oid] = oid_type
+            @store.register_type(oid, oid_type)
           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 f46af35f55..ed3e884455 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -369,7 +369,7 @@ module ActiveRecord
           raise "Your version of PostgreSQL (#{postgresql_version}) is too old, please upgrade!"
         end
 
-        @type_map = OID::TypeMap.new
+        @type_map = Type::HashLookupTypeMap.new
         initialize_type_map(type_map)
         @local_tz = execute('SHOW TIME ZONE', 'SCHEMA').first["TimeZone"]
         @use_insert_returning = @config.key?(:insert_returning) ? self.class.type_cast_config_to_boolean(@config[:insert_returning]) : true
@@ -543,12 +543,29 @@ module ActiveRecord
             initialize_type_map(type_map, [oid])
           end
 
-          type_map.fetch(oid, fmod) {
+          type_map.fetch(normalize_oid_type(oid, fmod)) {
             warn "unknown OID #{oid}: failed to recognize type of '#{column_name}'. It will be treated as String."
-            type_map[oid] = Type::Value.new
+            Type::Value.new.tap do |cast_type|
+              type_map.register_type(oid, cast_type)
+            end
           }
         end
 
+        def normalize_oid_type(ftype, fmod)
+          # The type for the numeric depends on the width of the field,
+          # so we'll do something special here.
+          #
+          # When dealing with decimal columns:
+          #
+          # places after decimal  = fmod - 4 & 0xffff
+          # places before decimal = (fmod - 4) >> 16 & 0xffff
+          if ftype == 1700 && (fmod - 4 & 0xffff).zero?
+            23
+          else
+            ftype
+          end
+        end
+
         def initialize_type_map(type_map, oids = nil)
           if supports_ranges?
             query = <<-SQL
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 828ada00a7..8503d3ea1b 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
@@ -2,10 +2,16 @@ module ActiveRecord
   module ConnectionAdapters
     module Type
       class HashLookupTypeMap < TypeMap # :nodoc:
+        delegate :key?, to: :@mapping
+
         def lookup(type)
           @mapping.fetch(type, proc { default_value }).call(type)
         end
 
+        def fetch(type, &block)
+          @mapping.fetch(type, block).call(type)
+        end
+
         def alias_type(type, alias_type)
           register_type(type) { lookup(alias_type) }
         end