From eafa43e870cb595b8efd95323082c9fd3e50caab Mon Sep 17 00:00:00 2001
From: schneems <richard.schneeman@gmail.com>
Date: Mon, 11 Jan 2016 10:22:26 -0600
Subject: [PATCH] Allow manually setting environment value

If for some reason some one is not able to set the environment from a migration this gives us an escape valve to manually set the environment for the database see https://github.com/rails/rails/pull/22967#issuecomment-170251635.

We will also fix the migration case, but this will ensure there is always a way to set the environment.

cc/ @sgrif
---
 activerecord/lib/active_record/migration.rb     | 17 +++++++++++------
 .../lib/active_record/railties/databases.rake   |  6 ++++++
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/activerecord/lib/active_record/migration.rb b/activerecord/lib/active_record/migration.rb
index f699e83ab3..4a6e9c12fe 100644
--- a/activerecord/lib/active_record/migration.rb
+++ b/activerecord/lib/active_record/migration.rb
@@ -145,7 +145,7 @@ def initialize(message = DEFAULT_MESSAGE)
 
   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"
+      msg = "Environment data not found in the schema. To resolve this issue, run: \n\n\tbin/rake db:environment:set"
       if defined?(Rails.env)
         super("#{msg} RAILS_ENV=#{::Rails.env}")
       else
@@ -157,7 +157,7 @@ def initialize
   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 << "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
@@ -165,10 +165,15 @@ def initialize(env = "production")
 
   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"
+      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, first set the environment using:\n\n"
+      msg << "\tbin/rake db:environment:set"
+      if defined?(Rails.env)
+        super("#{msg} RAILS_ENV=#{::Rails.env}")
+      else
+        super(msg)
+      end
     end
   end
 
diff --git a/activerecord/lib/active_record/railties/databases.rake b/activerecord/lib/active_record/railties/databases.rake
index ead9407391..d81d6b54b3 100644
--- a/activerecord/lib/active_record/railties/databases.rake
+++ b/activerecord/lib/active_record/railties/databases.rake
@@ -1,6 +1,12 @@
 require 'active_record'
 
 db_namespace = namespace :db do
+  desc "Set the environment value for the database"
+  task "environment:set" => [:environment, :load_config] do
+    ActiveRecord::InternalMetadata.create_table
+    ActiveRecord::InternalMetadata[:environment] = ActiveRecord::Migrator.current_environment
+  end
+
   task :check_protected_environments => [:environment, :load_config] do
     ActiveRecord::Tasks::DatabaseTasks.check_protected_environments!
   end
-- 
GitLab