2 files changed, 81 insertions, 14 deletions

diff --git a/config/initializers/active_record_data_types.rb b/config/initializers/active_record_data_types.rb
index beb97c6fce0..fef591c397d 100644
--- a/config/initializers/active_record_data_types.rb
+++ b/config/initializers/active_record_data_types.rb
@@ -4,21 +4,78 @@
 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' })
+  module RegisterDateTimeWithTimeZone
+    # Run original `initialize_type_map` and then register `timestamp` as a
+    # `MysqlDateTimeWithTimeZone`.
+    #
+    # When schema dumping, `timestamp` columns will be output as
+    # `t.datetime_with_timezone`.
+    def initialize_type_map(mapping)
+      super mapping
+
+      mapping.register_type(%r(timestamp)i) do |sql_type|
+        precision = extract_precision(sql_type)
+        ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter::MysqlDateTimeWithTimeZone.new(precision: precision)
       end
     end
   end
+
+  class ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter
+    prepend RegisterDateTimeWithTimeZone
+
+    # Add the class `DateTimeWithTimeZone` so we can map `timestamp` to it.
+    class MysqlDateTimeWithTimeZone < MysqlDateTime
+      def type
+        :datetime_with_timezone
+      end
+    end
+
+    # 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: 'timestamp' }
+  end
 end
diff --git a/config/initializers/active_record_table_definition.rb b/config/initializers/active_record_table_definition.rb
index 4f59e35f4da..8e3a1c7a62f 100644
--- a/config/initializers/active_record_table_definition.rb
+++ b/config/initializers/active_record_table_definition.rb
@@ -3,15 +3,15 @@
 
 require 'active_record/connection_adapters/abstract/schema_definitions'
 
-# Appends columns `created_at` and `updated_at` to a table.
-#
-# It is used in table creation like:
-# create_table 'users' do |t|
-#   t.timestamps_with_timezone
-# end
 module ActiveRecord
   module ConnectionAdapters
     class TableDefinition
+      # Appends columns `created_at` and `updated_at` to a table.
+      #
+      # It is used in table creation like:
+      # create_table 'users' do |t|
+      #   t.timestamps_with_timezone
+      # end
       def timestamps_with_timezone(**options)
         options[:null] = false if options[:null].nil?
 
@@ -19,6 +19,16 @@ module ActiveRecord
           column(column_name, :datetime_with_timezone, options)
         end
       end
+
+      # Adds specified column with appropriate timestamp type
+      #
+      # It is used in table creation like:
+      # create_table 'users' do |t|
+      #   t.datetime_with_timezone :did_something_at
+      # end
+      def datetime_with_timezone(column_name, **options)
+        column(column_name, :datetime_with_timezone, options)
+      end
     end
   end
 end