Add add_column_safely and backfill_column_safely migration helpers

CHANGELOG.md

@@ -14,6 +14,8 @@
 - Added `auto_analyze` for MySQL and MariaDB
 - Added `target_mysql_version` and `target_mariadb_version`
 - Switched to `up` for backfilling
+- Added `add_column_safely` method
+- Added `backfill_column_safely` method
 
 ## 0.5.1 (2019-12-17)
 

lib/strong_migrations.rb

@@ -2,6 +2,7 @@
 require "active_support"
 
 # modules
+require "strong_migrations/util"
 require "strong_migrations/checker"
 require "strong_migrations/database_tasks"
 require "strong_migrations/migration"
@@ -49,6 +50,18 @@ def up
   end
 end%{append}",
 
+    add_column_default_helper:
+"Adding a column with a non-null default causes the entire table to be rewritten.
+Instead, add the column without a default value, change the default and then backfill.
+
+class %{migration_name} < ActiveRecord::Migration%{migration_suffix}
+  disable_ddl_transaction!
+
+  def change
+    %{command}
+  end
+end",
+
     add_column_json:
 "There's no equality operator for the json column type, which can
 cause errors for existing SELECT DISTINCT queries. Use jsonb instead.",

lib/strong_migrations/checker.rb

@@ -1,5 +1,7 @@
 module StrongMigrations
   class Checker
+    include Util
+
     attr_accessor :direction
 
     def initialize(migration)
@@ -81,20 +83,24 @@ def perform(method, *args)
           default = options[:default]
 
           if !default.nil? && !((postgresql? && postgresql_version >= Gem::Version.new("11")) || (mysql? && mysql_version >= Gem::Version.new("8.0.12")) || (mariadb? && mariadb_version >= Gem::Version.new("10.3.2")))
-
-            if options[:null] == false
-              options = options.except(:null)
-              append = "
+            if helpers?
+              raise_error :add_column_default_helper,
+                command: command_str("add_column_safely", [table, column, type, options])
+            else
+              if options[:null] == false
+                options = options.except(:null)
+                append = "
 
 Then add the NOT NULL constraint."
-            end
+              end
 
-            raise_error :add_column_default,
-              add_command: command_str("add_column", [table, column, type, options.except(:default)]),
-              change_command: command_str("change_column_default", [table, column, default]),
-              remove_command: command_str("remove_column", [table, column]),
-              code: backfill_code(table, column, default),
-              append: append
+              raise_error :add_column_default,
+                add_command: command_str("add_column", [table, column, type, options.except(:default)]),
+                change_command: command_str("change_column_default", [table, column, default]),
+                remove_command: command_str("remove_column", [table, column]),
+                code: backfill_code(table, column, default),
+                append: append
+            end
           end
 
           if type.to_s == "json" && postgresql?
@@ -259,43 +265,6 @@ def version_safe?
       version && version <= StrongMigrations.start_after
     end
 
-    def postgresql?
-      connection.adapter_name =~ /postg/i # PostgreSQL, PostGIS
-    end
-
-    def postgresql_version
-      @postgresql_version ||= begin
-        target_version(StrongMigrations.target_postgresql_version) do
-          # only works with major versions
-          connection.select_all("SHOW server_version_num").first["server_version_num"].to_i / 10000
-        end
-      end
-    end
-
-    def mysql?
-      connection.adapter_name =~ /mysql/i && !connection.try(:mariadb?)
-    end
-
-    def mysql_version
-      @mysql_version ||= begin
-        target_version(StrongMigrations.target_mysql_version) do
-          connection.select_all("SELECT VERSION()").first["VERSION()"].split("-").first
-        end
-      end
-    end
-
-    def mariadb?
-      connection.adapter_name =~ /mysql/i && connection.try(:mariadb?)
-    end
-
-    def mariadb_version
-      @mariadb_version ||= begin
-        target_version(StrongMigrations.target_mariadb_version) do
-          connection.select_all("SELECT VERSION()").first["VERSION()"].split("-").first
-        end
-      end
-    end
-
     def target_version(target_version)
       version =
         if target_version && defined?(Rails) && (Rails.env.development? || Rails.env.test?)
@@ -330,7 +299,7 @@ def raise_error(message_key, header: nil, **vars)
 
     def constraint_str(statement, identifiers)
       # not all identifiers are tables, but this method of quoting should be fine
-      code = statement % identifiers.map { |v| connection.quote_table_name(v) }
+      code = quote_identifiers(statement, identifiers)
       "safety_assured do\n      execute '#{code}' \n    end"
     end
 

lib/strong_migrations/migration_helpers.rb

@@ -1,5 +1,7 @@
 module StrongMigrations
   module MigrationHelpers
+    include Util
+
     def add_foreign_key_safely(from_table, to_table, **options)
       ensure_postgresql(__method__)
       ensure_not_in_transaction(__method__)
@@ -73,16 +75,85 @@ def remove_null_constraint_safely(table_name, column_name, name: nil)
       end
     end
 
+    def add_column_safely(table_name, column_name, type, **options)
+      default = options.delete(:default)
+
+      if column_can_be_added_safely?(default)
+        safety_assured { add_column(table_name, column_name, type, default: default, **options) }
+      else
+        ensure_not_in_transaction(__method__)
+
+        reversible do |dir|
+          dir.up do
+            transaction do
+              add_column(table_name, column_name, type, default: nil, **options)
+              change_column_default(table_name, column_name, default)
+            end
+
+            backfill_column_safely(table_name, column_name, default)
+
+            allow_null = options[:null]
+            add_null_constraint_safely(table_name, column_name) unless allow_null
+          end
+
+          dir.down do
+            remove_column(table_name, column_name)
+          end
+        end
+      end
+    end
+
+    def backfill_column_safely(table_name, column_name, value)
+      ensure_not_in_transaction(__method__)
+
+      table = Arel::Table.new(table_name)
+      primary_key = connection.primary_key(table_name)
+
+      start_arel = table
+        .project(table[primary_key])
+        .where(table[column_name].eq(nil))
+        .order(table[primary_key].asc)
+        .take(1)
+
+      result = connection.exec_query(start_arel.to_sql)
+      return if result.empty?
+
+      start_pk = result.first[primary_key]
+      batch_size = 10_000
+
+      loop do
+        finish_arel = table
+          .project(table[primary_key])
+          .where(table[primary_key].gteq(start_pk))
+          .order(table[primary_key].asc)
+          .skip(batch_size)
+          .take(1)
+
+        finish_result = connection.exec_query(finish_arel.to_sql).first
+
+        update_arel = Arel::UpdateManager.new
+          .table(table)
+          .set([[table[column_name], value]])
+          .where(table[primary_key].gteq(start_pk))
+
+        if finish_result
+          finish_pk = finish_result[primary_key]
+          update_arel = update_arel.where(table[primary_key].lt(finish_pk))
+          start_pk = finish_pk
+        end
+
+        safety_assured { execute(update_arel.to_sql) }
+
+        break unless finish_pk
+      end
+    end
+
     private
 
     def ensure_postgresql(method_name)
       raise StrongMigrations::Error, "`#{method_name}` is intended for Postgres only" unless postgresql?
     end
 
-    def postgresql?
-      %w(PostgreSQL PostGIS).include?(connection.adapter_name)
-    end
-
     def ensure_not_in_transaction(method_name)
       if connection.transaction_open?
         raise StrongMigrations::Error, "Cannot run `#{method_name}` inside a transaction. Use `disable_ddl_transaction` to disable the transaction."
@@ -110,8 +181,8 @@ def on_delete_update_statement(delete_or_update, action)
       end
     end
 
-    def quote_identifiers(statement, identifiers)
-      statement % identifiers.map { |v| connection.quote_table_name(v) }
+    def column_can_be_added_safely?(default)
+      default.nil? || (postgresql? && postgresql_version >= Gem::Version.new("11")) || (mysql? && mysql_version >= Gem::Version.new("8.0.12")) || (mariadb? && mariadb_version >= Gem::Version.new("10.3.2"))
     end
   end
 end

lib/strong_migrations/util.rb

@@ -0,0 +1,45 @@
+module StrongMigrations
+  module Util
+    def postgresql?
+      connection.adapter_name =~ /postg/i # PostgreSQL, PostGIS
+    end
+
+    def postgresql_version
+      @postgresql_version ||= begin
+        target_version(StrongMigrations.target_postgresql_version) do
+          # only works with major versions
+          connection.select_all("SHOW server_version_num").first["server_version_num"].to_i / 10000
+        end
+      end
+    end
+
+    def mysql?
+      connection.adapter_name =~ /mysql/i && !connection.try(:mariadb?)
+    end
+
+    def mysql_version
+      @mysql_version ||= begin
+        target_version(StrongMigrations.target_mysql_version) do
+          connection.select_all("SELECT VERSION()").first["VERSION()"].split("-").first
+        end
+      end
+    end
+
+    def mariadb?
+      connection.adapter_name =~ /mysql/i && connection.try(:mariadb?)
+    end
+
+    def mariadb_version
+      @mariadb_version ||= begin
+        target_version(StrongMigrations.target_mariadb_version) do
+          connection.select_all("SELECT VERSION()").first["VERSION()"].split("-").first
+        end
+      end
+    end
+
+    def quote_identifiers(statement, identifiers)
+      # not all identifiers are tables, but this method of quoting should be fine
+      statement % identifiers.map { |v| connection.quote_table_name(v) }
+    end
+  end
+end

test/migration_helpers_test.rb

@@ -16,6 +16,22 @@ def change
   end
 end
 
+class AddColumnSafely < TestMigration
+  disable_ddl_transaction!
+
+  def change
+    add_column_safely :users, :balance, :integer, default: 10, null: false
+  end
+end
+
+class BackfillColumnSafely < TestMigration
+  disable_ddl_transaction!
+
+  def up
+    backfill_column_safely :users, :city, "San Francisco"
+  end
+end
+
 class MigrationHelpersTest < Minitest::Test
   def setup
     skip unless ENV["HELPERS"]
@@ -58,6 +74,42 @@ def test_add_null_constraint_safely
     migrate(AddNullConstraintSafely, direction: :down)
   end
 
+  def test_add_column_safely_raises_inside_transaction
+    skip unless postgresql?
+    error = assert_raises(StrongMigrations::Error) { migrate(AddColumnSafely, transaction: true) }
+    assert_match "Cannot run `add_column_safely` inside a transaction", error.message
+  end
+
+  def test_add_column_safely
+    User.reset_column_information
+    migrate(AddColumnSafely)
+
+    column = User.columns.find { |c| c.name == "balance" }
+    assert_equal :integer, column.type
+    assert_equal "10", column.default
+    assert_equal false, column.null
+  ensure
+    migrate(AddColumnSafely, direction: :down)
+    User.reset_column_information
+  end
+
+  def test_backfill_column_safely_raises_inside_transaction
+    error = assert_raises(StrongMigrations::Error) { migrate(BackfillColumnSafely, transaction: true) }
+    assert_match "Cannot run `backfill_column_safely` inside a transaction", error.message
+  end
+
+  def test_backfill_column_safely
+    user1 = User.create(name: "John", city: "Los Angeles")
+    user2 = User.create(name: "Jane", city: nil)
+    migrate(BackfillColumnSafely)
+
+    assert_equal "Los Angeles", user1.reload.city
+    assert_equal "San Francisco", user2.reload.city
+  ensure
+    User.delete_all
+    User.reset_column_information
+  end
+
   private
 
   def connection

test/test_helper.rb

@@ -53,6 +53,9 @@ def change
 end
 migrate CreateUsers
 
+class User < ActiveRecord::Base
+end
+
 module Helpers
   def postgresql?
     ENV["ADAPTER"].nil?