Merge pull request #22967 from schneems/schneems/generic-metadata

Prevent destructive action on production database

10 files changed, 150 insertions, 7 deletions

diff --git a/activerecord/lib/active_record.rb b/activerecord/lib/active_record.rb
index f35e1d889e..ab3846ae65 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 a918a8b035..70868ebd03 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/schema_statements.rb
@@ -968,6 +968,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/connection_adapters/abstract_mysql_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
index f63074ac21..828e46637d 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
@@ -502,6 +502,7 @@ module ActiveRecord
       end
 
       def table_exists?(table_name)
+        # Update lib/active_record/internal_metadata.rb when this gets removed
         ActiveSupport::Deprecation.warn(<<-MSG.squish)
           #table_exists? currently checks both tables and views.
           This behavior is deprecated and will be changed with Rails 5.1 to only check tables.
diff --git a/activerecord/lib/active_record/internal_metadata.rb b/activerecord/lib/active_record/internal_metadata.rb
new file mode 100644
index 0000000000..7bd66028b6
--- /dev/null
+++ b/activerecord/lib/active_record/internal_metadata.rb
@@ -0,0 +1,47 @@
+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 # :nodoc:
+    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 []=(key, value)
+        first_or_initialize(key: key).update_attributes!(value: value)
+      end
+
+      def [](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, primary_key: :key, 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 8c9eab436d..f699e83ab3 100644
--- a/activerecord/lib/active_record/migration.rb
+++ b/activerecord/lib/active_record/migration.rb
@@ -143,6 +143,35 @@ 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
+
+  class EnvironmentMismatchError < ActiveRecordError
+    def initialize(current: nil, stored: nil)
+      msg =  "You are attempting to modify a database that was last run in #{ stored } environment.\n"
+      msg << "You are running in #{ current } environment."
+      msg << "if you are sure you want to continue, run the same command with the environment variable\n"
+      msg << "DISABLE_DATABASE_ENVIRONMENT_CHECK=1"
+    end
+  end
+
   # = Active Record Migrations
   #
   # Migrations can manage the evolution of a schema used by several physical
@@ -1078,6 +1107,7 @@ module ActiveRecord
       validate(@migrations)
 
       Base.connection.initialize_schema_migrations_table
+      Base.connection.initialize_internal_metadata_table
     end
 
     def current_version
@@ -1202,10 +1232,28 @@ 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)
+        ActiveRecord::InternalMetadata[:environment] = ActiveRecord::Migrator.current_environment
       end
     end
 
+    def self.last_stored_environment
+      return nil if current_version == 0
+      raise NoEnvironmentInSchemaError unless ActiveRecord::InternalMetadata.table_exists?
+
+      environment = ActiveRecord::InternalMetadata[:environment]
+      raise NoEnvironmentInSchemaError unless environment
+      environment
+    end
+
+    def self.current_environment
+      ActiveRecord::ConnectionHandling::DEFAULT_ENV.call
+    end
+
+    def self.protected_environment?
+      ActiveRecord::Base.protected_environments.include?(last_stored_environment) if 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 a6a68f3d4b..f26c8471bc 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 9b59ee995a..ead9407391 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..784a02d2c3 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[:environment] = ActiveRecord::Migrator.current_environment
     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 b6fba0cf79..6a9af5d1b4 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -42,6 +42,22 @@ module ActiveRecord
 
       LOCAL_HOSTS    = ['127.0.0.1', 'localhost']
 
+      def check_protected_environments!
+        unless ENV['DISABLE_DATABASE_internal_metadata']
+          current = ActiveRecord::Migrator.current_environment
+          stored  = ActiveRecord::Migrator.last_stored_environment
+
+          if ActiveRecord::Migrator.protected_environment?
+            raise ActiveRecord::ProtectedEnvironmentError.new(stored)
+          end
+
+          if stored && stored != current
+            raise ActiveRecord::EnvironmentMismatchError.new(current: current, stored: stored)
+          end
+        end
+      rescue ActiveRecord::NoDatabaseError
+      end
+
       def register_task(pattern, task)
         @tasks ||= {}
         @tasks[pattern] = task
@@ -204,6 +220,8 @@ module ActiveRecord
         else
           raise ArgumentError, "unknown format #{format.inspect}"
         end
+        ActiveRecord::InternalMetadata.create_table
+        ActiveRecord::InternalMetadata[:environment] = ActiveRecord::Migrator.current_environment
       end
 
       def load_schema_for(*args)