Merge pull request #8009 from graceliu/3-2-fix_database_url_support

fixed support for DATABASE_URL for rake db tasks

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