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`.

9 files changed, 129 insertions, 7 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)