Prevent destructive action on production database

This PR introduces a key/value type store to Active Record that can be used for storing internal values. It is an alternative implementation to #21237 cc @sgrif @matthewd.

It is possible to run your tests against your production database by accident right now. While infrequently, but as an anecdotal data point, Heroku receives a non-trivial number of requests for a database restore due to this happening. In these cases the loss can be large.

To prevent against running tests against production we can store the "environment" version that was used when migrating the database in a new internal table. Before executing tests we can see if the database is a listed in `protected_environments` and abort. There is a manual escape valve to force this check from happening with environment variable `DISABLE_DATABASE_ENVIRONMENT_CHECK=1`.

16 files changed, 232 insertions, 14 deletions

diff --git a/activerecord/lib/active_record.rb b/activerecord/lib/active_record.rb
index 264f869c68..bfd90f7fb0 100644
--- a/activerecord/lib/active_record.rb
+++ b/activerecord/lib/active_record.rb
@@ -40,6 +40,7 @@ module ActiveRecord
   autoload :CounterCache
   autoload :DynamicMatchers
   autoload :Enum
+  autoload :InternalMetadata
   autoload :Explain
   autoload :Inheritance
   autoload :Integration
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb b/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
index b50d28862c..ac77011b83 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
@@ -969,6 +969,10 @@ module ActiveRecord
         ActiveRecord::SchemaMigration.create_table
       end
 
+      def initialize_internal_metadata_table
+        ActiveRecord::InternalMetadata.create_table
+      end
+
       def assume_migrated_upto_version(version, migrations_paths)
         migrations_paths = Array(migrations_paths)
         version = version.to_i
diff --git a/activerecord/lib/active_record/internal_metadata.rb b/activerecord/lib/active_record/internal_metadata.rb
new file mode 100644
index 0000000000..e907bd0563
--- /dev/null
+++ b/activerecord/lib/active_record/internal_metadata.rb
@@ -0,0 +1,49 @@
+require 'active_record/scoping/default'
+require 'active_record/scoping/named'
+
+module ActiveRecord
+  # This class is used to create a table that keeps track of values and keys such
+  # as which environment migrations were run in.
+  class InternalMetadata < ActiveRecord::Base
+    class << self
+      def primary_key
+        "key"
+      end
+
+      def table_name
+        "#{table_name_prefix}#{ActiveRecord::Base.internal_metadata_table_name}#{table_name_suffix}"
+      end
+
+      def index_name
+        "#{table_name_prefix}unique_#{ActiveRecord::Base.internal_metadata_table_name}#{table_name_suffix}"
+      end
+
+      def store(hash)
+        hash.each do |key, value|
+          first_or_initialize(key: key).update_attributes!(value: value)
+        end
+      end
+
+      def value_for(key)
+        where(key: key).pluck(:value).first
+      end
+
+      def table_exists?
+        ActiveSupport::Deprecation.silence { connection.table_exists?(table_name) }
+      end
+
+      # Creates a internal metadata table with columns +key+ and +value+
+      def create_table
+        unless table_exists?
+          connection.create_table(table_name, id: false) do |t|
+            t.column :key,   :string
+            t.column :value, :string
+            t.timestamps
+          end
+
+          connection.add_index table_name, :key, unique: true, name: index_name
+        end
+      end
+    end
+  end
+end
diff --git a/activerecord/lib/active_record/migration.rb b/activerecord/lib/active_record/migration.rb
index 43726d795e..c61491d58c 100644
--- a/activerecord/lib/active_record/migration.rb
+++ b/activerecord/lib/active_record/migration.rb
@@ -143,6 +143,26 @@ module ActiveRecord
     end
   end
 
+  class NoEnvironmentInSchemaError < MigrationError #:nodoc:
+    def initialize
+      msg = "Environment data not found in the schema. To resolve this issue, run: \n\n\tbin/rake db:migrate"
+      if defined?(Rails.env)
+        super("#{msg} RAILS_ENV=#{::Rails.env}")
+      else
+        super(msg)
+      end
+    end
+  end
+
+  class ProtectedEnvironmentError < ActiveRecordError #:nodoc:
+    def initialize(env = "production")
+      msg = "You are attempting to run a destructive action against your '#{env}' database\n"
+      msg << "if you are sure you want to continue, run the same command with the environment variable\n"
+      msg << "DISABLE_DATABASE_ENVIRONMENT_CHECK=1"
+      super(msg)
+    end
+  end
+
   # = Active Record Migrations
   #
   # Migrations can manage the evolution of a schema used by several physical
@@ -1078,6 +1098,7 @@ module ActiveRecord
       validate(@migrations)
 
       Base.connection.initialize_schema_migrations_table
+      Base.connection.initialize_internal_metadata_table
     end
 
     def current_version
@@ -1202,10 +1223,24 @@ module ActiveRecord
         ActiveRecord::SchemaMigration.where(:version => version.to_s).delete_all
       else
         migrated << version
-        ActiveRecord::SchemaMigration.create!(:version => version.to_s)
+        ActiveRecord::SchemaMigration.create!(version: version.to_s)
+        environment = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+        ActiveRecord::InternalMetadata.store(environment: environment)
       end
     end
 
+    def self.last_stored_environment
+      ActiveRecord::InternalMetadata.value_for(:environment)
+    end
+
+    def self.protected_environment?
+      return false if current_version == 0
+      raise NoEnvironmentInSchemaError unless ActiveRecord::InternalMetadata.table_exists?
+
+      raise NoEnvironmentInSchemaError unless last_stored_environment
+      ActiveRecord::Base.protected_environments.include?(last_stored_environment)
+    end
+
     def up?
       @direction == :up
     end
diff --git a/activerecord/lib/active_record/model_schema.rb b/activerecord/lib/active_record/model_schema.rb
index 5df67cdbe7..b82923494f 100644
--- a/activerecord/lib/active_record/model_schema.rb
+++ b/activerecord/lib/active_record/model_schema.rb
@@ -44,6 +44,19 @@ module ActiveRecord
 
       ##
       # :singleton-method:
+      # Accessor for the name of the internal metadata table. By default, the value is "active_record_internal_metadatas"
+      class_attribute :internal_metadata_table_name, instance_accessor: false
+      self.internal_metadata_table_name = "active_record_internal_metadatas"
+
+      ##
+      # :singleton-method:
+      # Accessor for an array of names of environments where destructive actions should be prohibited. By default,
+      # the value is ["production"]
+      class_attribute :protected_environments, instance_accessor: false
+      self.protected_environments = ["production"]
+
+      ##
+      # :singleton-method:
       # Indicates whether table names should be the pluralized versions of the corresponding class names.
       # If true, the default table name for a Product class will be +products+. If false, it would just be +product+.
       # See table_name for the full rules on table/class naming. This is true, by default.
diff --git a/activerecord/lib/active_record/railties/databases.rake b/activerecord/lib/active_record/railties/databases.rake
index 5b1ea16f0b..c1203fb745 100644
--- a/activerecord/lib/active_record/railties/databases.rake
+++ b/activerecord/lib/active_record/railties/databases.rake
@@ -1,6 +1,10 @@
 require 'active_record'
 
 db_namespace = namespace :db do
+  task :check_protected_environments => [:environment, :load_config] do
+    ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+  end
+
   task :load_config do
     ActiveRecord::Base.configurations       = ActiveRecord::Tasks::DatabaseTasks.database_configuration || {}
     ActiveRecord::Migrator.migrations_paths = ActiveRecord::Tasks::DatabaseTasks.migrations_paths
@@ -18,24 +22,28 @@ db_namespace = namespace :db do
   end
 
   namespace :drop do
-    task :all => :load_config do
+    task :all => [:load_config, :check_protected_environments] do
       ActiveRecord::Tasks::DatabaseTasks.drop_all
     end
   end
 
   desc 'Drops the database from DATABASE_URL or config/database.yml for the current RAILS_ENV (use db:drop:all to drop all databases in the config). Without RAILS_ENV, it defaults to dropping the development and test databases.'
-  task :drop => [:load_config] do
+  task :drop => [:load_config, :check_protected_environments] do
+    db_namespace["drop:_unsafe"].invoke
+  end
+
+  task "drop:_unsafe" => [:load_config] do
     ActiveRecord::Tasks::DatabaseTasks.drop_current
   end
 
   namespace :purge do
-    task :all => :load_config do
+    task :all => [:load_config, :check_protected_environments] do
       ActiveRecord::Tasks::DatabaseTasks.purge_all
     end
   end
 
   # desc "Empty the database from DATABASE_URL or config/database.yml for the current RAILS_ENV (use db:purge:all to purge all databases in the config). Without RAILS_ENV it defaults to purging the development and test databases."
-  task :purge => [:load_config] do
+  task :purge => [:load_config, :check_protected_environments] do
     ActiveRecord::Tasks::DatabaseTasks.purge_current
   end
 
@@ -351,7 +359,7 @@ db_namespace = namespace :db do
     task :clone_structure => %w(db:test:deprecated db:structure:dump db:test:load_structure)
 
     # desc "Empty the test database"
-    task :purge => %w(environment load_config) do
+    task :purge => %w(environment load_config check_protected_environments) do
       ActiveRecord::Tasks::DatabaseTasks.purge ActiveRecord::Base.configurations['test']
     end
 
diff --git a/activerecord/lib/active_record/schema.rb b/activerecord/lib/active_record/schema.rb
index fdf9965a82..0a0ea33196 100644
--- a/activerecord/lib/active_record/schema.rb
+++ b/activerecord/lib/active_record/schema.rb
@@ -51,6 +51,9 @@ module ActiveRecord
         initialize_schema_migrations_table
         connection.assume_migrated_upto_version(info[:version], migrations_paths)
       end
+
+      ActiveRecord::InternalMetadata.create_table
+      ActiveRecord::InternalMetadata.store("environment" => ActiveRecord::ConnectionHandling::DEFAULT_ENV.call)
     end
 
     private
diff --git a/activerecord/lib/active_record/schema_dumper.rb b/activerecord/lib/active_record/schema_dumper.rb
index 2362dae9fc..65005bd44b 100644
--- a/activerecord/lib/active_record/schema_dumper.rb
+++ b/activerecord/lib/active_record/schema_dumper.rb
@@ -254,7 +254,7 @@ HEADER
       end
 
       def ignored?(table_name)
-        [ActiveRecord::Base.schema_migrations_table_name, ignore_tables].flatten.any? do |ignored|
+        [ActiveRecord::Base.schema_migrations_table_name, ActiveRecord::Base.internal_metadata_table_name, ignore_tables].flatten.any? do |ignored|
           ignored === remove_prefix_and_suffix(table_name)
         end
       end
diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index c0c29a618c..6318a0725d 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -42,6 +42,13 @@ module ActiveRecord
 
       LOCAL_HOSTS    = ['127.0.0.1', 'localhost']
 
+      def check_protected_environments!
+        if !ENV['DISABLE_DATABASE_internal_metadata'] && ActiveRecord::Migrator.protected_environment?
+          raise ActiveRecord::ProtectedEnvironmentError.new(ActiveRecord::Migrator.last_stored_environment)
+        end
+      rescue ActiveRecord::NoDatabaseError
+      end
+
       def register_task(pattern, task)
         @tasks ||= {}
         @tasks[pattern] = task
@@ -204,6 +211,8 @@ module ActiveRecord
         else
           raise ArgumentError, "unknown format #{format.inspect}"
         end
+        ActiveRecord::InternalMetadata.create_table
+        ActiveRecord::InternalMetata.store("environment" => ActiveRecord::ConnectionHandling::DEFAULT_ENV.call)
       end
 
       def load_schema_for(*args)
diff --git a/activerecord/test/cases/migration_test.rb b/activerecord/test/cases/migration_test.rb
index 19be357b6e..2f7859a7be 100644
--- a/activerecord/test/cases/migration_test.rb
+++ b/activerecord/test/cases/migration_test.rb
@@ -354,6 +354,50 @@ class MigrationTest < ActiveRecord::TestCase
     Reminder.reset_table_name
   end
 
+  def test_internal_metadata_table_name
+    original_internal_metadata_table_name = ActiveRecord::Base.internal_metadata_table_name
+
+    assert_equal "active_record_internal_metadatas", ActiveRecord::InternalMetadata.table_name
+    ActiveRecord::Base.table_name_prefix = "prefix_"
+    ActiveRecord::Base.table_name_suffix = "_suffix"
+    Reminder.reset_table_name
+    assert_equal "prefix_active_record_internal_metadatas_suffix", ActiveRecord::InternalMetadata.table_name
+    ActiveRecord::Base.internal_metadata_table_name = "changed"
+    Reminder.reset_table_name
+    assert_equal "prefix_changed_suffix", ActiveRecord::InternalMetadata.table_name
+    ActiveRecord::Base.table_name_prefix = ""
+    ActiveRecord::Base.table_name_suffix = ""
+    Reminder.reset_table_name
+    assert_equal "changed", ActiveRecord::InternalMetadata.table_name
+  ensure
+    ActiveRecord::Base.internal_metadata_table_name = original_internal_metadata_table_name
+    Reminder.reset_table_name
+  end
+
+  def test_internal_metadata_stores_environment
+    current_env     = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+    migrations_path = MIGRATIONS_ROOT + "/valid"
+    old_path        = ActiveRecord::Migrator.migrations_paths
+    ActiveRecord::Migrator.migrations_paths = migrations_path
+
+    assert_equal current_env, ActiveRecord::InternalMetadata.value_for("environment")
+
+    original_rails_env  = ENV["RAILS_ENV"]
+    original_rack_env   = ENV["RACK_ENV"]
+    ENV["RAILS_ENV"]    = ENV["RACK_ENV"] = "foofoo"
+    new_env     = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+
+    refute_equal current_env, new_env
+
+    sleep 1 # mysql by default does not store fractional seconds in the database
+    ActiveRecord::Migrator.up(migrations_path)
+    assert_equal new_env, ActiveRecord::InternalMetadata.value_for("environment")
+  ensure
+    ActiveRecord::Migrator.migrations_paths = old_path
+    ENV["RAILS_ENV"] = original_rails_env
+    ENV["RACK_ENV"]  = original_rack_env
+  end
+
   def test_proper_table_name_on_migration
     reminder_class = new_isolated_reminder_class
     migration = ActiveRecord::Migration.new
diff --git a/activerecord/test/cases/schema_dumper_test.rb b/activerecord/test/cases/schema_dumper_test.rb
index 9a5e4313d8..97c3d664ba 100644
--- a/activerecord/test/cases/schema_dumper_test.rb
+++ b/activerecord/test/cases/schema_dumper_test.rb
@@ -38,6 +38,7 @@ class SchemaDumperTest < ActiveRecord::TestCase
     assert_match %r{create_table "accounts"}, output
     assert_match %r{create_table "authors"}, output
     assert_no_match %r{create_table "schema_migrations"}, output
+    assert_no_match %r{create_table "internal_metadatas"}, output
   end
 
   def test_schema_dump_uses_force_cascade_on_create_table
@@ -158,6 +159,7 @@ class SchemaDumperTest < ActiveRecord::TestCase
     assert_no_match %r{create_table "accounts"}, output
     assert_match %r{create_table "authors"}, output
     assert_no_match %r{create_table "schema_migrations"}, output
+    assert_no_match %r{create_table "internal_metadatas"}, output
   end
 
   def test_schema_dump_with_regexp_ignored_table
@@ -165,6 +167,7 @@ class SchemaDumperTest < ActiveRecord::TestCase
     assert_no_match %r{create_table "accounts"}, output
     assert_match %r{create_table "authors"}, output
     assert_no_match %r{create_table "schema_migrations"}, output
+    assert_no_match %r{create_table "internal_metadatas"}, output
   end
 
   def test_schema_dumps_index_columns_in_right_order
@@ -342,6 +345,7 @@ class SchemaDumperTest < ActiveRecord::TestCase
     assert_no_match %r{create_table "foo_.+_bar"}, output
     assert_no_match %r{add_index "foo_.+_bar"}, output
     assert_no_match %r{create_table "schema_migrations"}, output
+    assert_no_match %r{create_table "internal_metadatas"}, output
 
     if ActiveRecord::Base.connection.supports_foreign_keys?
       assert_no_match %r{add_foreign_key "foo_.+_bar"}, output
diff --git a/activerecord/test/cases/tasks/database_tasks_test.rb b/activerecord/test/cases/tasks/database_tasks_test.rb
index c8f4179313..e91a222534 100644
--- a/activerecord/test/cases/tasks/database_tasks_test.rb
+++ b/activerecord/test/cases/tasks/database_tasks_test.rb
@@ -18,6 +18,34 @@ module ActiveRecord
     sqlite3:    :sqlite_tasks
   }
 
+  class DatabaseTasksUtilsTask< ActiveRecord::TestCase
+    def test_raises_an_error_when_called_with_protected_environment
+      ActiveRecord::Migrator.stubs(:current_version).returns(1)
+
+      protected_environments = ActiveRecord::Base.protected_environments.dup
+      current_env            = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+      assert !protected_environments.include?(current_env)
+      # Assert no error
+      ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+
+      ActiveRecord::Base.protected_environments << current_env
+      assert_raise(ActiveRecord::ProtectedEnvironmentError) do
+        ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+      end
+    ensure
+      ActiveRecord::Base.protected_environments = protected_environments
+    end
+
+    def test_raises_an_error_if_no_migrations_have_been_made
+      ActiveRecord::InternalMetadata.stubs(:table_exists?).returns(false)
+      ActiveRecord::Migrator.stubs(:current_version).returns(1)
+
+      assert_raise(ActiveRecord::NoEnvironmentInSchemaError) do
+        ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
+      end
+    end
+  end
+
   class DatabaseTasksRegisterTask < ActiveRecord::TestCase
     def test_register_task
       klazz = Class.new do
diff --git a/railties/test/application/bin_setup_test.rb b/railties/test/application/bin_setup_test.rb
index 1bdced02e9..c96bbe4c70 100644
--- a/railties/test/application/bin_setup_test.rb
+++ b/railties/test/application/bin_setup_test.rb
@@ -28,7 +28,7 @@ module ApplicationTests
         assert_not File.exist?("tmp/restart.txt")
         `bin/setup 2>&1`
         assert_equal 0, File.size("log/my.log")
-        assert_equal '["articles", "schema_migrations"]', list_tables.call
+        assert_equal '["articles", "schema_migrations", "internal_metadatas"]', list_tables.call
         assert File.exist?("tmp/restart.txt")
       end
     end
diff --git a/railties/test/application/loading_test.rb b/railties/test/application/loading_test.rb
index 2cc599ca6f..40abaf860d 100644
--- a/railties/test/application/loading_test.rb
+++ b/railties/test/application/loading_test.rb
@@ -116,11 +116,11 @@ class LoadingTest < ActiveSupport::TestCase
     require "#{rails_root}/config/environment"
     setup_ar!
 
-    assert_equal [ActiveRecord::SchemaMigration], ActiveRecord::Base.descendants
+    assert_equal [ActiveRecord::SchemaMigration, ActiveRecord::InternalMetadata], ActiveRecord::Base.descendants
     get "/load"
-    assert_equal [ActiveRecord::SchemaMigration, Post], ActiveRecord::Base.descendants
+    assert_equal [ActiveRecord::SchemaMigration, ActiveRecord::InternalMetadata, Post], ActiveRecord::Base.descendants
     get "/unload"
-    assert_equal [ActiveRecord::SchemaMigration], ActiveRecord::Base.descendants
+    assert_equal [ActiveRecord::SchemaMigration, ActiveRecord::InternalMetadata], ActiveRecord::Base.descendants
   end
 
   test "initialize cant be called twice" do
diff --git a/railties/test/application/rake/dbs_test.rb b/railties/test/application/rake/dbs_test.rb
index 0b0fb50fe1..47ab65e57b 100644
--- a/railties/test/application/rake/dbs_test.rb
+++ b/railties/test/application/rake/dbs_test.rb
@@ -85,7 +85,7 @@ module ApplicationTests
 
       test 'db:drop failure because database does not exist' do
         Dir.chdir(app_path) do
-          output = `bin/rake db:drop 2>&1`
+          output = `bin/rake db:drop:_unsafe --trace 2>&1`
           assert_match(/does not exist/, output)
           assert_equal 0, $?.exitstatus
         end
@@ -222,14 +222,14 @@ module ApplicationTests
 
           assert_equal '["posts"]', list_tables[]
           `bin/rake db:schema:load`
-          assert_equal '["posts", "comments", "schema_migrations"]', list_tables[]
+          assert_equal '["posts", "comments", "schema_migrations", "internal_metadatas"]', list_tables[]
 
           app_file 'db/structure.sql', <<-SQL
             CREATE TABLE "users" ("id" INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, "name" varchar(255));
           SQL
 
           `bin/rake db:structure:load`
-          assert_equal '["posts", "comments", "schema_migrations", "users"]', list_tables[]
+          assert_equal '["posts", "comments", "schema_migrations", "internal_metadatas", "users"]', list_tables[]
         end
       end
 
diff --git a/railties/test/application/rake_test.rb b/railties/test/application/rake_test.rb
index 1d3550add9..cf4e5ddbf6 100644
--- a/railties/test/application/rake_test.rb
+++ b/railties/test/application/rake_test.rb
@@ -24,6 +24,26 @@ module ApplicationTests
       assert $task_loaded
     end
 
+    def test_the_test_rake_task_is_protected_when_previous_migration_was_production
+      Dir.chdir(app_path) do
+        output = `bin/rails generate model product name:string;
+         env RAILS_ENV=production bin/rake db:create db:migrate;
+         env RAILS_ENV=production bin/rake db:test:prepare test 2>&1`
+
+        assert_match /ActiveRecord::ProtectedEnvironmentError/, output
+      end
+    end
+
+    def test_not_protected_when_previous_migration_was_not_production
+      Dir.chdir(app_path) do
+        output = `bin/rails generate model product name:string;
+         env RAILS_ENV=test bin/rake db:create db:migrate;
+         env RAILS_ENV=test bin/rake db:test:prepare test 2>&1`
+
+        refute_match /ActiveRecord::ProtectedEnvironmentError/, output
+      end
+    end
+
     def test_environment_is_required_in_rake_tasks
       app_file "config/environment.rb", <<-RUBY
         SuperMiddleware = Struct.new(:app)