fixed support for DATABASE_URL for rake db tasks

Backport for #7521

- added tests to confirm establish_connection uses DATABASE_URL and
  Rails.env correctly even when no arguments are passed in.
- updated rake db tasks to support DATABASE_URL, and added tests to
  confirm correct behavior for these rake tasks.  (Removed
  establish_connection call from some tasks since in those cases
  the :environment task already made sure the function would be called)
- updated Resolver so that when it resolves the database url, it
  removes hash values with empty strings from the config spec (e.g.
  to support connection to postgresql when no username is specified).
- updated ResolverTest to use current_adapter? to check the type of
  the current adapter.

6 files changed, 287 insertions, 37 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index aa68934099..c762210941 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,5 +1,7 @@
 ## Rails 3.2.9 (unreleased)
 
+*   Fixed support for DATABASE_URL environment variable for rake db tasks. *Grace Liu*
+
 *   Fix bug where `update_columns` and `update_column` would not let you update the primary key column.
 
     *Henrik Nyh*
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_specification.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_specification.rb
index f366392c50..04ab7e0a43 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_specification.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_specification.rb
@@ -67,7 +67,7 @@ module ActiveRecord
                    :port     => config.port,
                    :database => config.path.sub(%r{^/},""),
                    :host     => config.host }
-          spec.reject!{ |_,value| !value }
+          spec.reject!{ |_,value| value.blank? }
           spec.map { |key,value| spec[key] = URI.unescape(value) if value.is_a?(String) }
           if config.query
             options = Hash[config.query.split("&").map{ |pair| pair.split("=") }].symbolize_keys
diff --git a/activerecord/lib/active_record/railties/databases.rake b/activerecord/lib/active_record/railties/databases.rake
index 987bb558e5..ae6d9c8653 100644
--- a/activerecord/lib/active_record/railties/databases.rake
+++ b/activerecord/lib/active_record/railties/databases.rake
@@ -2,6 +2,25 @@ require 'active_support/core_ext/object/inclusion'
 require 'active_record'
 
 db_namespace = namespace :db do
+  def database_url_config
+    @database_url_config ||=
+        ActiveRecord::Base::ConnectionSpecification::Resolver.new(ENV["DATABASE_URL"], {}).spec.config.stringify_keys
+  end
+
+  def current_config(options = {})
+    options = { :env => Rails.env }.merge! options
+
+    if options[:config]
+      @current_config = options[:config]
+    else
+      @current_config ||= if ENV['DATABASE_URL']
+                            database_url_config
+                          else
+                            ActiveRecord::Base.configurations[options[:env]]
+                          end
+    end
+  end
+
   task :load_config do
     ActiveRecord::Base.configurations = Rails.application.config.database_configuration
     ActiveRecord::Migrator.migrations_paths = Rails.application.paths['db/migrate'].to_a
@@ -35,10 +54,14 @@ db_namespace = namespace :db do
     end
   end
 
-  desc 'Create the database from config/database.yml for the current Rails.env (use db:create:all to create all dbs in the config)'
+  desc 'Create the database from DATABASE_URL or config/database.yml for the current Rails.env (use db:create:all to create all dbs in the config)'
   task :create => [:load_config, :rails_env] do
-    configs_for_environment.each { |config| create_database(config) }
-    ActiveRecord::Base.establish_connection(configs_for_environment.first)
+    if ENV['DATABASE_URL']
+      create_database(database_url_config)
+    else
+      configs_for_environment.each { |config| create_database(config) }
+      ActiveRecord::Base.establish_connection(configs_for_environment.first)
+    end
   end
 
   def mysql_creation_options(config)
@@ -133,9 +156,13 @@ db_namespace = namespace :db do
     end
   end
 
-  desc 'Drops the database for the current Rails.env (use db:drop:all to drop all databases)'
+  desc 'Drops the database using DATABASE_URL or the current Rails.env (use db:drop:all to drop all databases)'
   task :drop => [:load_config, :rails_env] do
-    configs_for_environment.each { |config| drop_database_and_rescue(config) }
+    if ENV['DATABASE_URL']
+      drop_database_and_rescue(database_url_config)
+    else
+      configs_for_environment.each { |config| drop_database_and_rescue(config) }
+    end
   end
 
   def local_database?(config, &block)
@@ -146,7 +173,6 @@ db_namespace = namespace :db do
     end
   end
 
-
   desc "Migrate the database (options: VERSION=x, VERBOSE=false)."
   task :migrate => [:environment, :load_config] do
     ActiveRecord::Migration.verbose = ENV["VERBOSE"] ? ENV["VERBOSE"] == "true" : true
@@ -201,8 +227,6 @@ db_namespace = namespace :db do
 
     desc 'Display status of migrations'
     task :status => [:environment, :load_config] do
-      config = ActiveRecord::Base.configurations[Rails.env]
-      ActiveRecord::Base.establish_connection(config)
       unless ActiveRecord::Base.connection.table_exists?(ActiveRecord::Migrator.schema_migrations_table_name)
         puts 'Schema migrations table does not exist yet.'
         next  # means "return" for rake task
@@ -222,7 +246,7 @@ db_namespace = namespace :db do
         ['up', version, '********** NO FILE **********']
       end
       # output
-      puts "\ndatabase: #{config['database']}\n\n"
+      puts "\ndatabase: #{ActiveRecord::Base.connection_config[:database]}\n\n"
       puts "#{'Status'.center(8)}  #{'Migration ID'.ljust(14)}  Migration Name"
       puts "-" * 50
       (db_list + file_list).sort_by {|migration| migration[1]}.each do |migration|
@@ -314,7 +338,6 @@ db_namespace = namespace :db do
     task :load => [:environment, :load_config] do
       require 'active_record/fixtures'
 
-      ActiveRecord::Base.establish_connection(Rails.env)
       base_dir     = File.join [Rails.root, ENV['FIXTURES_PATH'] || %w{test fixtures}].flatten
       fixtures_dir = File.join [base_dir, ENV['FIXTURES_DIR']].compact
 
@@ -353,7 +376,6 @@ db_namespace = namespace :db do
       require 'active_record/schema_dumper'
       filename = ENV['SCHEMA'] || "#{Rails.root}/db/schema.rb"
       File.open(filename, "w:utf-8") do |file|
-        ActiveRecord::Base.establish_connection(Rails.env)
         ActiveRecord::SchemaDumper.dump(ActiveRecord::Base.connection, file)
       end
       db_namespace['schema:dump'].reenable
@@ -377,25 +399,25 @@ db_namespace = namespace :db do
   namespace :structure do
     desc 'Dump the database structure to db/structure.sql. Specify another file with DB_STRUCTURE=db/my_structure.sql'
     task :dump => [:environment, :load_config] do
-      abcs = ActiveRecord::Base.configurations
+      config = current_config
       filename = ENV['DB_STRUCTURE'] || File.join(Rails.root, "db", "structure.sql")
-      case abcs[Rails.env]['adapter']
+      case config['adapter']
       when /mysql/, 'oci', 'oracle'
-        ActiveRecord::Base.establish_connection(abcs[Rails.env])
+        ActiveRecord::Base.establish_connection(config)
         File.open(filename, "w:utf-8") { |f| f << ActiveRecord::Base.connection.structure_dump }
       when /postgresql/
-        set_psql_env(abcs[Rails.env])
-        search_path = abcs[Rails.env]['schema_search_path']
+        set_psql_env(config)
+        search_path = config['schema_search_path']
         unless search_path.blank?
           search_path = search_path.split(",").map{|search_path_part| "--schema=#{Shellwords.escape(search_path_part.strip)}" }.join(" ")
         end
-        `pg_dump -i -s -x -O -f #{Shellwords.escape(filename)} #{search_path} #{Shellwords.escape(abcs[Rails.env]['database'])}`
+        `pg_dump -i -s -x -O -f #{Shellwords.escape(filename)} #{search_path} #{Shellwords.escape(config['database'])}`
         raise 'Error dumping database' if $?.exitstatus == 1
       when /sqlite/
-        dbfile = abcs[Rails.env]['database']
+        dbfile = config['database']
         `sqlite3 #{dbfile} .schema > #{filename}`
       when 'sqlserver'
-        `smoscript -s #{abcs[Rails.env]['host']} -d #{abcs[Rails.env]['database']} -u #{abcs[Rails.env]['username']} -p #{abcs[Rails.env]['password']} -f #{filename} -A -U`
+        `smoscript -s #{config['host']} -d #{config['database']} -u #{config['username']} -p #{config['password']} -f #{filename} -A -U`
       when "firebird"
         set_firebird_env(abcs[Rails.env])
         db_string = firebird_db_string(abcs[Rails.env])
@@ -412,36 +434,34 @@ db_namespace = namespace :db do
 
     # desc "Recreate the databases from the structure.sql file"
     task :load => [:environment, :load_config] do
-      env = Rails.env
-
-      abcs = ActiveRecord::Base.configurations
+      config = current_config
       filename = ENV['DB_STRUCTURE'] || File.join(Rails.root, "db", "structure.sql")
-      case abcs[env]['adapter']
+      case config['adapter']
       when /mysql/
-        ActiveRecord::Base.establish_connection(abcs[env])
+        ActiveRecord::Base.establish_connection(config)
         ActiveRecord::Base.connection.execute('SET foreign_key_checks = 0')
         IO.read(filename).split("\n\n").each do |table|
           ActiveRecord::Base.connection.execute(table)
         end
       when /postgresql/
-        set_psql_env(abcs[env])
-        `psql -f "#{filename}" #{abcs[env]['database']}`
+        set_psql_env(config)
+        `psql -f "#{filename}" #{config['database']}`
       when /sqlite/
-        dbfile = abcs[env]['database']
+        dbfile = config['database']
         `sqlite3 #{dbfile} < "#{filename}"`
       when 'sqlserver'
-        `sqlcmd -S #{abcs[env]['host']} -d #{abcs[env]['database']} -U #{abcs[env]['username']} -P #{abcs[env]['password']} -i #{filename}`
+        `sqlcmd -S #{config['host']} -d #{config['database']} -U #{config['username']} -P #{config['password']} -i #{filename}`
       when 'oci', 'oracle'
-        ActiveRecord::Base.establish_connection(abcs[env])
+        ActiveRecord::Base.establish_connection(config)
         IO.read(filename).split(";\n\n").each do |ddl|
           ActiveRecord::Base.connection.execute(ddl)
         end
       when 'firebird'
-        set_firebird_env(abcs[env])
-        db_string = firebird_db_string(abcs[env])
+        set_firebird_env(config)
+        db_string = firebird_db_string(config)
         sh "isql -i #{filename} #{db_string}"
       else
-        raise "Task not supported by '#{abcs[env]['adapter']}'"
+        raise "Task not supported by '#{config['adapter']}'"
       end
     end
 
@@ -465,10 +485,10 @@ db_namespace = namespace :db do
     # desc "Recreate the test database from an existent structure.sql file"
     task :load_structure => 'db:test:purge' do
       begin
-        old_env, ENV['RAILS_ENV'] = ENV['RAILS_ENV'], 'test'
+        current_config(:config => ActiveRecord::Base.configurations['test'])
         db_namespace["structure:load"].invoke
       ensure
-        ENV['RAILS_ENV'] = old_env
+        current_config(:config => nil)
       end
     end
 
diff --git a/activerecord/test/cases/connection_specification/resolver_test.rb b/activerecord/test/cases/connection_specification/resolver_test.rb
index 46fc1a252f..451652dfcb 100644
--- a/activerecord/test/cases/connection_specification/resolver_test.rb
+++ b/activerecord/test/cases/connection_specification/resolver_test.rb
@@ -13,7 +13,6 @@ module ActiveRecord
           spec = resolve 'mysql://foo?encoding=utf8'
           assert_equal({
             :adapter  => "mysql",
-            :database => "",
             :host     => "foo",
             :encoding => "utf8" }, spec)
         end
@@ -33,7 +32,6 @@ module ActiveRecord
           spec = resolve 'mysql://foo:123?encoding=utf8'
           assert_equal({
             :adapter  => "mysql",
-            :database => "",
             :port     => 123,
             :host     => "foo",
             :encoding => "utf8" }, spec)
diff --git a/railties/test/application/initializers/frameworks_test.rb b/railties/test/application/initializers/frameworks_test.rb
index cf6c4d8fc2..b4c9741b05 100644
--- a/railties/test/application/initializers/frameworks_test.rb
+++ b/railties/test/application/initializers/frameworks_test.rb
@@ -193,5 +193,41 @@ module ApplicationTests
       require "#{app_path}/config/environment"
       assert_nil defined?(ActiveRecord::Base)
     end
+
+    test "active record establish_connection uses Rails.env if DATABASE_URL is not set" do
+      begin
+        require "#{app_path}/config/environment"
+        orig_database_url = ENV.delete("DATABASE_URL")
+        orig_rails_env, Rails.env = Rails.env, 'development'
+
+        ActiveRecord::Base.establish_connection
+
+        assert ActiveRecord::Base.connection
+        assert_match /#{ActiveRecord::Base.configurations[Rails.env]['database']}/, ActiveRecord::Base.connection_config[:database]
+      ensure
+        ActiveRecord::Base.remove_connection
+        ENV["DATABASE_URL"] = orig_database_url if orig_database_url
+        Rails.env = orig_rails_env if orig_rails_env
+      end
+    end
+
+    test "active record establish_connection uses DATABASE_URL even if Rails.env is set" do
+      begin
+        require "#{app_path}/config/environment"
+        orig_database_url = ENV.delete("DATABASE_URL")
+        orig_rails_env, Rails.env = Rails.env, 'development'
+        database_url_db_name = "db/database_url_db.sqlite3"
+        ENV["DATABASE_URL"] = "sqlite3://:@localhost/#{database_url_db_name}"
+
+        ActiveRecord::Base.establish_connection
+
+        assert ActiveRecord::Base.connection
+        assert_match /#{database_url_db_name}/, ActiveRecord::Base.connection_config[:database]
+      ensure
+        ActiveRecord::Base.remove_connection
+        ENV["DATABASE_URL"] = orig_database_url if orig_database_url
+        Rails.env = orig_rails_env if orig_rails_env
+      end
+    end
   end
 end
diff --git a/railties/test/application/rake/dbs_test.rb b/railties/test/application/rake/dbs_test.rb
new file mode 100644
index 0000000000..479ed2612b
--- /dev/null
+++ b/railties/test/application/rake/dbs_test.rb
@@ -0,0 +1,194 @@
+require "isolation/abstract_unit"
+
+module ApplicationTests
+  module RakeTests
+    class RakeDbsTest < Test::Unit::TestCase
+      include ActiveSupport::Testing::Isolation
+
+      def setup
+        build_app
+        boot_rails
+        FileUtils.rm_rf("#{app_path}/config/environments")
+      end
+
+      def teardown
+        teardown_app
+      end
+
+      def database_url_db_name
+        "db/database_url_db.sqlite3"
+      end
+
+      def set_database_url
+        ENV['DATABASE_URL'] = "sqlite3://:@localhost/#{database_url_db_name}"
+      end
+
+      def expected
+        @expected ||= {}
+      end
+
+      def db_create_and_drop
+        Dir.chdir(app_path) do
+          output = `bundle exec rake db:create`
+          assert_equal output, ""
+          assert File.exists?(expected[:database])
+          assert_equal expected[:database],
+                        ActiveRecord::Base.connection_config[:database]
+          output = `bundle exec rake db:drop`
+          assert_equal output, ""
+          assert !File.exists?(expected[:database])
+        end
+      end
+
+      test 'db:create and db:drop without database url' do
+        require "#{app_path}/config/environment"
+        expected[:database] = ActiveRecord::Base.configurations[Rails.env]['database']
+        db_create_and_drop
+      end
+
+      test 'db:create and db:drop with database url' do
+        require "#{app_path}/config/environment"
+        set_database_url
+        expected[:database] = database_url_db_name
+        db_create_and_drop
+      end
+
+      def db_migrate_and_status
+        Dir.chdir(app_path) do
+          `rails generate model book title:string`
+          `bundle exec rake db:migrate`
+          output = `bundle exec rake db:migrate:status`
+          assert_match(/database:\s+\S*#{expected[:database]}/, output)
+          assert_match(/up\s+\d{14}\s+Create books/, output)
+        end
+      end
+
+      test 'db:migrate and db:migrate:status without database_url' do
+        require "#{app_path}/config/environment"
+        expected[:database] = ActiveRecord::Base.configurations[Rails.env]['database']
+        db_migrate_and_status
+      end
+
+      test 'db:migrate and db:migrate:status with database_url' do
+        require "#{app_path}/config/environment"
+        set_database_url
+        expected[:database] = database_url_db_name
+        db_migrate_and_status
+      end
+
+      def db_schema_dump
+        Dir.chdir(app_path) do
+          `rails generate model book title:string`
+          `rake db:migrate`
+          `rake db:schema:dump`
+
+          assert File.exists?("db/schema.rb"), "db/schema.rb doesn't exist"
+
+          schema_dump = File.read("db/schema.rb")
+
+          assert_match(/create_table \"books\"/, schema_dump)
+        end
+      end
+
+      test 'db:schema:dump without database_url' do
+        db_schema_dump
+      end
+
+      test 'db:schema:dump with database_url' do
+        set_database_url
+        db_schema_dump
+      end
+
+      def db_fixtures_load
+        Dir.chdir(app_path) do
+          `rails generate model book title:string`
+          `bundle exec rake db:migrate`
+          `bundle exec rake db:fixtures:load`
+          assert_match /#{expected[:database]}/,
+                    ActiveRecord::Base.connection_config[:database]
+          require "#{app_path}/app/models/book"
+          assert_equal 2, Book.count
+        end
+      end
+
+      test 'db:fixtures:load without database_url' do
+        require "#{app_path}/config/environment"
+        expected[:database] =  ActiveRecord::Base.configurations[Rails.env]['database']
+        db_fixtures_load
+      end
+
+      test 'db:fixtures:load with database_url' do
+        require "#{app_path}/config/environment"
+        set_database_url
+        expected[:database] = database_url_db_name
+        db_fixtures_load
+      end
+
+      def db_structure_dump_and_load
+        Dir.chdir(app_path) do
+          `rails generate model book title:string`
+          `bundle exec rake db:create`
+          `bundle exec rake db:migrate`
+          `bundle exec rake db:structure:dump`
+
+          assert File.exists?("db/structure.sql"), "db/structure.sql doesn't exist"
+
+          structure_dump = File.read("db/structure.sql")
+
+          assert_match /CREATE TABLE \"books\"/, structure_dump
+
+          `bundle exec rake db:drop`
+          `bundle exec rake db:structure:load`
+
+          assert_match /#{expected[:database]}/,
+                       ActiveRecord::Base.connection_config[:database]
+
+          require "#{app_path}/app/models/book"
+          #if structure is not loaded correctly, exception would be raised
+          assert_equal Book.count, 0
+        end
+      end
+
+      test 'db:structure:dump and db:structure:load without database_url' do
+        require "#{app_path}/config/environment"
+        expected[:database] = ActiveRecord::Base.configurations[Rails.env]['database']
+        db_structure_dump_and_load
+      end
+
+      test 'db:structure:dump and db:structure:load with database_url' do
+        require "#{app_path}/config/environment"
+        set_database_url
+        expected[:database] = database_url_db_name
+        db_structure_dump_and_load
+      end
+
+      def db_test_load_structure
+        Dir.chdir(app_path) do
+          `rails generate model book title:string`
+          `bundle exec rake db:migrate`
+          `bundle exec rake db:structure:dump`
+          `bundle exec rake db:test:load_structure`
+          ActiveRecord::Base.configurations = Rails.application.config.database_configuration
+          ActiveRecord::Base.establish_connection 'test'
+          require "#{app_path}/app/models/book"
+
+          #if structure is not loaded correctly, exception would be raised
+          assert_equal Book.count, 0
+          assert_match /#{ActiveRecord::Base.configurations['test']['database']}/,
+                       ActiveRecord::Base.connection_config[:database]
+        end
+      end
+
+      test 'db:test:load_structure without database_url' do
+        require "#{app_path}/config/environment"
+        db_test_load_structure
+      end
+
+      test 'db:test:load_structure with database_url' do
+        require "#{app_path}/config/environment"
+        set_database_url
+        db_test_load_structure
+      end
+    end
+  end
+end