Merge pull request #6761 from freelancing-god/db-rake

Get logic out of db rake tasks, and into classes and objects

4 files changed, 252 insertions, 0 deletions

diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
new file mode 100644
index 0000000000..10cc679a05
--- /dev/null
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -0,0 +1,82 @@
+class ActiveRecord::Tasks::DatabaseTasks
+  TASKS_PATTERNS = {
+    /mysql/      => ActiveRecord::Tasks::MySQLDatabaseTasks,
+    /postgresql/ => ActiveRecord::Tasks::PostgreSQLDatabaseTasks,
+    /sqlite/     => ActiveRecord::Tasks::SQLiteDatabaseTasks
+  }
+  LOCAL_HOSTS    = ['127.0.0.1', 'localhost']
+
+  def self.create(*arguments)
+    configuration = arguments.first
+    class_for_adapter(configuration['adapter']).new(*arguments).create
+  rescue Exception => error
+    $stderr.puts error, *(error.backtrace)
+    $stderr.puts "Couldn't create database for #{configuration.inspect}"
+  end
+
+  def self.create_all
+    each_local_configuration { |configuration| create configuration }
+  end
+
+  def self.create_current(environment = Rails.env)
+    each_current_configuration(environment) { |configuration|
+      create configuration
+    }
+    ActiveRecord::Base.establish_connection environment
+  end
+
+  def self.drop(*arguments)
+    configuration = arguments.first
+    class_for_adapter(configuration['adapter']).new(*arguments).drop
+  rescue Exception => error
+    $stderr.puts error, *(error.backtrace)
+    $stderr.puts "Couldn't drop #{configuration['database']}"
+  end
+
+  def self.drop_all
+    each_local_configuration { |configuration| drop configuration }
+  end
+
+  def self.drop_current(environment = Rails.env)
+    each_current_configuration(environment) { |configuration|
+      drop configuration
+    }
+  end
+
+  def self.purge(configuration)
+    class_for_adapter(configuration['adapter']).new(configuration).purge
+  end
+
+  private
+
+  def self.class_for_adapter(adapter)
+    key = TASKS_PATTERNS.keys.detect { |pattern| adapter[pattern] }
+    TASKS_PATTERNS[key]
+  end
+
+  def self.each_current_configuration(environment)
+    environments = [environment]
+    environments << 'test' if environment.development?
+
+    configurations = ActiveRecord::Base.configurations.values_at(*environments)
+    configurations.compact.each do |configuration|
+      yield configuration unless configuration['database'].blank?
+    end
+  end
+
+  def self.each_local_configuration
+    ActiveRecord::Base.configurations.each_value do |configuration|
+      next unless configuration['database']
+
+      if local_database?(configuration)
+        yield configuration
+      else
+        $stderr.puts "This task only modifies local databases. #{configuration['database']} is on a remote host."
+      end
+    end
+  end
+
+  def self.local_database?(configuration)
+    configuration['host'].in?(LOCAL_HOSTS) || configuration['host'].blank?
+  end
+end
diff --git a/activerecord/lib/active_record/tasks/mysql_database_tasks.rb b/activerecord/lib/active_record/tasks/mysql_database_tasks.rb
new file mode 100644
index 0000000000..eb598629fa
--- /dev/null
+++ b/activerecord/lib/active_record/tasks/mysql_database_tasks.rb
@@ -0,0 +1,88 @@
+class ActiveRecord::Tasks::MySQLDatabaseTasks
+  DEFAULT_CHARSET     = ENV['CHARSET']   || 'utf8'
+  DEFAULT_COLLATION   = ENV['COLLATION'] || 'utf8_unicode_ci'
+  ACCESS_DENIED_ERROR = 1045
+
+  delegate :connection, :establish_connection, :to => ActiveRecord::Base
+
+  def initialize(configuration)
+    @configuration = configuration
+  end
+
+  def create
+    establish_connection configuration_without_database
+    connection.create_database configuration['database'], creation_options
+    establish_connection configuration
+  rescue error_class => error
+    raise error unless error.errno == ACCESS_DENIED_ERROR
+
+    $stdout.print error.error
+    establish_connection root_configuration_without_database
+    connection.create_database configuration['database'], creation_options
+    connection.execute grant_statement.gsub(/\s+/, ' ').strip
+    establish_connection configuration
+  rescue error_class => error
+    $stderr.puts error.error
+    $stderr.puts "Couldn't create database for #{configuration.inspect}, #{creation_options.inspect}"
+    $stderr.puts "(If you set the charset manually, make sure you have a matching collation)" if configuration['charset']
+  end
+
+  def drop
+    establish_connection configuration
+    connection.drop_database configuration['database']
+  end
+
+  def purge
+    establish_connection :test
+    connection.recreate_database configuration['database'], creation_options
+  end
+
+  private
+
+  def configuration
+    @configuration
+  end
+
+  def configuration_without_database
+    configuration.merge('database' => nil)
+  end
+
+  def creation_options
+    {
+      :charset   => (configuration['charset']   || DEFAULT_CHARSET),
+      :collation => (configuration['collation'] || DEFAULT_COLLATION)
+    }
+  end
+
+  def error_class
+    case configuration['adapter']
+    when /jdbc/
+      require 'active_record/railties/jdbcmysql_error'
+      ArJdbcMySQL::Error
+    when /mysql2/
+      Mysql2::Error
+    else
+      Mysql::Error
+    end
+  end
+
+  def grant_statement
+    <<-SQL
+GRANT ALL PRIVILEGES ON #{configuration['database']}.*
+  TO '#{configuration['username']}'@'localhost'
+IDENTIFIED BY '#{configuration['password']}' WITH GRANT OPTION;
+    SQL
+  end
+
+  def root_configuration_without_database
+    configuration_without_database.merge(
+      'username' => 'root',
+      'password' => root_password
+    )
+  end
+
+  def root_password
+    $stdout.print "Please provide the root password for your mysql installation\n>"
+    $stdin.gets.strip
+  end
+end
diff --git a/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb b/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb
new file mode 100644
index 0000000000..331825d3ec
--- /dev/null
+++ b/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb
@@ -0,0 +1,45 @@
+class ActiveRecord::Tasks::PostgreSQLDatabaseTasks
+  DEFAULT_ENCODING = ENV['CHARSET'] || 'utf8'
+
+  delegate :connection, :establish_connection, :clear_active_connections!,
+    :to => ActiveRecord::Base
+
+  def initialize(configuration)
+    @configuration = configuration
+  end
+
+  def create(master_established = false)
+    establish_master_connection unless master_established
+    connection.create_database configuration['database'],
+      configuration.merge('encoding' => encoding)
+    establish_connection configuration
+  end
+
+  def drop
+    establish_master_connection
+    connection.drop_database configuration['database']
+  end
+
+  def purge
+    clear_active_connections!
+    drop
+    create true
+  end
+
+  private
+
+  def configuration
+    @configuration
+  end
+
+  def encoding
+    configuration['encoding'] || DEFAULT_ENCODING
+  end
+
+  def establish_master_connection
+    establish_connection configuration.merge(
+      'database'           => 'postgres',
+      'schema_search_path' => 'public'
+    )
+  end
+end
diff --git a/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb b/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb
new file mode 100644
index 0000000000..4d3cba8a50
--- /dev/null
+++ b/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb
@@ -0,0 +1,37 @@
+class ActiveRecord::Tasks::SQLiteDatabaseTasks
+  delegate :connection, :establish_connection, :to => ActiveRecord::Base
+
+  def initialize(configuration, root = Rails.root)
+    @configuration, @root = configuration, root
+  end
+
+  def create
+    if File.exist? configuration['database']
+      $stderr.puts "#{configuration['database']} already exists"
+      return
+    end
+
+    establish_connection configuration
+    connection
+  end
+
+  def drop
+    require 'pathname'
+    path = Pathname.new configuration['database']
+    file = path.absolute? ? path.to_s : File.join(root, path)
+
+    FileUtils.rm(file)
+  end
+
+  alias :purge :drop
+
+  private
+
+  def configuration
+    @configuration
+  end
+
+  def root
+    @root
+  end
+end