extract pg type map initialization process to `TypeMapInitializer`.

2 files changed, 74 insertions, 65 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/oid.rb b/activerecord/lib/active_record/connection_adapters/postgresql/oid.rb
index b173163e41..1e89f8cfd6 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/oid.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/oid.rb
@@ -374,6 +374,77 @@ This is not reliable and will be removed in the future.
           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.
+        #   - Type is an OID::Type object.
+        # This class has side effects on the +store+ passed during initialization.
+        class TypeMapInitializer # :nodoc:
+          def initialize(store)
+            @store = store
+          end
+
+          def run(records)
+            mapped, nodes = records.partition { |row| OID.registered_type? row['typname'] }
+            ranges, nodes = nodes.partition { |row| row['typtype'] == 'r' }
+            enums, nodes = nodes.partition { |row| row['typtype'] == 'e' }
+            domains, nodes = nodes.partition { |row| row['typtype'] == 'd' }
+            arrays, nodes = nodes.partition { |row| row['typinput'] == 'array_in' }
+            composites, nodes = nodes.partition { |row| row['typelem'] != '0' }
+
+            mapped.each     { |row| register_mapped_type(row)    }
+            enums.each      { |row| register_enum_type(row)      }
+            domains.each    { |row| register_domain_type(row)    }
+            arrays.each     { |row| register_array_type(row)     }
+            ranges.each     { |row| register_range_type(row)     }
+            composites.each { |row| register_composite_type(row) }
+          end
+
+          private
+          def register_mapped_type(row)
+            register row['oid'], OID::NAMES[row['typname']]
+          end
+
+          def register_enum_type(row)
+            register row['oid'], OID::Enum.new
+          end
+
+          def register_array_type(row)
+            if subtype = @store[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]
+              register row['oid'], OID::Range.new(subtype)
+            end
+          end
+
+          def register_domain_type(row)
+            if base_type = @store[row["typbasetype"].to_i]
+              register row['oid'], base_type
+            else
+              warn "unknown base type (OID: #{row["typbasetype"]}) for domain #{row["typname"]}."
+            end
+          end
+
+          def register_composite_type(row)
+            if subtype = @store[row['typelem'].to_i]
+              register row['oid'], OID::Vector.new(row['typdelim'], subtype)
+            end
+          end
+
+          def register(oid, oid_type)
+            oid = oid.to_i
+
+            raise ArgumentError, "can't register nil type for OID #{oid}" if oid_type.nil?
+            return if @store.key?(oid)
+
+            @store[oid] = oid_type
+          end
+        end
+
         # When the PG adapter connects, the pg_type table is queried. The
         # key of this hash maps to the `typname` column from the table.
         # type_map is then dynamically built with oids as the key and type
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
index 764cb576d9..4908c5a47f 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -571,25 +571,6 @@ module ActiveRecord
           initialize_type_map(type_map)
         end
 
-        def add_oid(row, records_by_oid, type_map)
-          return type_map if type_map.key? row['type_elem'].to_i
-
-          if OID.registered_type? row['typname']
-            # this composite type is explicitly registered
-            vector = OID::NAMES[row['typname']]
-          else
-            # use the default for composite types
-            unless type_map.key? row['typelem'].to_i
-              add_oid records_by_oid[row['typelem']], records_by_oid, type_map
-            end
-
-            vector = OID::Vector.new row['typdelim'], type_map[row['typelem'].to_i]
-          end
-
-          type_map[row['oid'].to_i] = vector
-          type_map
-        end
-
         def initialize_type_map(type_map, oids = nil)
           if supports_ranges?
             query = <<-SQL
@@ -608,52 +589,9 @@ module ActiveRecord
             query += "WHERE t.oid::integer IN (%s)" % oids.join(", ")
           end
 
-          result = execute(query, 'SCHEMA')
-          ranges, nodes = result.partition { |row| row['typtype'] == 'r' }
-          enums, nodes = nodes.partition { |row| row['typtype'] == 'e' }
-          domains, nodes = nodes.partition { |row| row['typtype'] == 'd' }
-          arrays, nodes = nodes.partition { |row| row['typinput'] == 'array_in' }
-          leaves, nodes = nodes.partition { |row| row['typelem'] == '0' }
-
-          # populate the enum types
-          enums.each do |row|
-            type_map[row['oid'].to_i] = OID::Enum.new
-          end
-
-          # populate the base types
-          leaves.find_all { |row| OID.registered_type? row['typname'] }.each do |row|
-            type_map[row['oid'].to_i] = OID::NAMES[row['typname']]
-          end
-
-          records_by_oid = result.group_by { |row| row['oid'] }
-
-          # populate composite types
-          nodes.each do |row|
-            add_oid row, records_by_oid, type_map
-          end
-
-          # populate array types
-          arrays.find_all { |row| type_map.key? row['typelem'].to_i }.each do |row|
-            array = OID::Array.new  type_map[row['typelem'].to_i]
-            type_map[row['oid'].to_i] = array
-          end
-
-          # populate range types
-          ranges.find_all { |row| type_map.key? row['rngsubtype'].to_i }.each do |row|
-            subtype = type_map[row['rngsubtype'].to_i]
-            range = OID::Range.new subtype
-            type_map[row['oid'].to_i] = range
-          end
-
-          # populate domain types
-          domains.each do |row|
-            base_type_oid = row["typbasetype"].to_i
-            if base_type = type_map[base_type_oid]
-              type_map[row['oid'].to_i] = base_type
-            else
-              warn "unknown base type (OID: #{base_type_oid}) for domain #{row["typname"]}."
-            end
-          end
+          initializer = OID::TypeMapInitializer.new(type_map)
+          records = execute(query, 'SCHEMA')
+          initializer.run(records)
         end
 
         FEATURE_NOT_SUPPORTED = "0A000" #:nodoc: