Fix PostgreSQL schema dump for `timestamptz`

1 files changed, 36 insertions, 5 deletions

diff --git a/config/initializers/active_record_data_types.rb b/config/initializers/active_record_data_types.rb
index beb97c6fce0..454d5b11902 100644
--- a/config/initializers/active_record_data_types.rb
+++ b/config/initializers/active_record_data_types.rb
@@ -4,20 +4,51 @@
 if Gitlab::Database.postgresql?
   require 'active_record/connection_adapters/postgresql_adapter'
 
-  module ActiveRecord
-    module ConnectionAdapters
-      class PostgreSQLAdapter
-        NATIVE_DATABASE_TYPES.merge!(datetime_with_timezone: { name: 'timestamptz' })
+  module ActiveRecord::ConnectionAdapters::PostgreSQL::OID
+    # Add the class `DateTimeWithTimeZone` so we can map `timestamptz` to it.
+    class DateTimeWithTimeZone < DateTime
+      def type
+        :datetime_with_timezone
+      end
+    end
+  end
+
+  module RegisterDateTimeWithTimeZone
+    # Run original `initialize_type_map` and then register `timestamptz` as a
+    # `DateTimeWithTimeZone`.
+    #
+    # Apparently it does not matter that the original `initialize_type_map`
+    # aliases `timestamptz` to `timestamp`.
+    #
+    # When schema dumping, `timestamptz` columns will be output as
+    # `t.datetime_with_timezone`.
+    def initialize_type_map(mapping)
+      super mapping
+
+      mapping.register_type 'timestamptz' do |_, _, sql_type|
+        precision = extract_precision(sql_type)
+        ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::OID::DateTimeWithTimeZone.new(precision: precision)
       end
     end
   end
+
+  class ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+    prepend RegisterDateTimeWithTimeZone
+
+    # Add column type `datetime_with_timezone` so we can do this in
+    # migrations:
+    #
+    #   add_column(:users, :datetime_with_timezone)
+    #
+    NATIVE_DATABASE_TYPES[:datetime_with_timezone] = { name: 'timestamptz' }
+  end
 elsif Gitlab::Database.mysql?
   require 'active_record/connection_adapters/mysql2_adapter'
 
   module ActiveRecord
     module ConnectionAdapters
       class AbstractMysqlAdapter
-        NATIVE_DATABASE_TYPES.merge!(datetime_with_timezone: { name: 'timestamp' })
+        NATIVE_DATABASE_TYPES[:datetime_with_timezone] = { name: 'timestamp' }
       end
     end
   end