lib/tasks/databases.rake in yam-db-charmer-1.7.01 vs lib/tasks/databases.rake in yam-db-charmer-1.7.4.0

@@ -21,11 +21,11 @@
     end
   end
 
   desc 'Create the databases defined in config/database.yml for the current RAILS_ENV'
   task :create => "db:load_config" do
-    create_core_and_sub_database(ActiveRecord::Base.configurations[RAILS_ENV])
+    create_core_and_sub_database(ActiveRecord::Base.configurations[Rails.env])
   end
 
   def create_core_and_sub_database(config)
     create_database(config)
     config.each_value do | sub_config |
@@ -43,15 +43,20 @@
         next unless config['database']
         # Only connect to local databases
         local_database?(config) { drop_core_and_sub_database(config) }
       end
     end
+
+    task :schema_shards => :environment do
+      config = ::ActiveRecord::Base.configurations[Rails.env || 'development']
+      drop_schema_shard_databases(config)
+    end      
   end
 
   desc 'Drops the database for the current RAILS_ENV'
   task :drop => "db:load_config" do
-    config = ::ActiveRecord::Base.configurations[RAILS_ENV || 'development']
+    config = ::ActiveRecord::Base.configurations[Rails.env || 'development']
     begin
       drop_core_and_sub_database(config)
     rescue Exception => e
       puts "Couldn't drop #{config['database']} : #{e.inspect}"
     end
@@ -66,17 +71,38 @@
     end
   end
 end
 
 def drop_core_and_sub_database(config)
+  exit unless Rails.env=='test'
   drop_database(config)
   config.each_value do | sub_config |
     next unless sub_config.is_a?(Hash)
     next unless sub_config['database']
     begin
       drop_database(sub_config)
-    rescue
-      $stderr.puts "#{config['database']} not exists"
+    rescue => e
+      $stderr.puts "#{e.to_s}, #{config['database']} not exists"
     end
   end
 end
+
+# find all schema sharded databases and drop them
+def drop_schema_shard_databases(config)
+  exit unless Rails.env=='test'
+  config.each do |name, sub_config|
+    next unless sub_config.is_a?(Hash)
+    next unless sub_config['database']
+
+    # find the database connection for the schema admin db
+    next unless sub_config['shard_db_name_prefix']
+    connection = DbCharmer::Sharding.sharded_connection_by_connection_name(name)
+    next unless connection
+
+    # itereate through entries in the shards_info table to find the
+    # databases that will be dropped
+    dbgm = DbCharmer::Sharding::Method::DbBlockSchemaMap.new(connection.config)
+    dbgm.drop_all_shard_databases
+  end
+end
+