active_record_shards 2.3.1 → 2.4.1

data/lib/active_record_shards/connection_switcher.rb

@@ -28,7 +28,7 @@ module ActiveRecordShards
       if supports_sharding?
         shard_names.each do |shard|
           switch_connection(:shard => shard)
-          yield
+          yield(shard)
         end
       else
         yield

data/lib/active_record_shards/migration.rb

@@ -0,0 +1,71 @@
+module ActiveRecord
+  class Migrator
+    class << self
+      [:up, :down, :run].each do |m|
+        define_method("#{m}_with_sharding") do |*args|
+          ActiveRecord::Base.on_shard(nil) do
+            self.send("#{m}_without_sharding", *args)
+          end
+          ActiveRecord::Base.on_all_shards do
+            self.send("#{m}_without_sharding", *args)
+          end
+        end
+        alias_method_chain m.to_sym, :sharding
+      end
+
+      def bootstrap_migrations_from_nil_shard(migrations_path, this_migration=nil)
+        migrations = nil
+        ActiveRecord::Base.on_shard(nil) do
+          migrations = ActiveRecord::Migrator.new(:up, migrations_path).migrated
+        end
+
+        puts "inserting #{migrations.size} migrations on all shards..."
+        ActiveRecord::Base.on_all_shards do
+          migrator = ActiveRecord::Migrator.new(:up, migrations_path)
+          migrations.each do |m|
+            migrator.__send__(:record_version_state_after_migrating, m)
+          end
+          if this_migration
+            migrator.__send__(:record_version_state_after_migrating, this_migration)
+          end
+        end
+      end
+    end
+  end
+end
+
+module ActiveRecordShards
+  module MigrationExtension
+    attr_accessor :migration_shard
+
+    def shard(arg = nil)
+      self.migration_shard = arg
+    end
+
+    def migrate_with_forced_shard(direction)
+      if migration_shard.blank?
+        raise "Can't run migrations with a shard spec: this may be :all, :none, or a specific shard (for data-fixups) please call shard(arg) in your migration"
+      end
+
+      shard = ActiveRecord::Base.current_shard_selection.shard
+
+      if shard.nil?
+        return if migration_shard != :none
+      else
+        return if migration_shard == :none
+        return if migration_shard != :all && migration_shard.to_s != shard.to_s
+      end
+
+      migrate_without_forced_shard(direction)
+    end
+  end
+end
+
+ActiveRecord::Migration.class_eval do
+  extend ActiveRecordShards::MigrationExtension
+  class << self
+    alias_method_chain :migrate, :forced_shard
+  end
+end
+
+ActiveRecord::MigrationProxy.delegate :migration_shard, :to => :migration

data/lib/active_record_shards.rb

@@ -6,6 +6,7 @@ require 'active_record_shards/shard_selection'
 require 'active_record_shards/connection_switcher'
 require 'active_record_shards/association_collection_connection_selection'
 require 'active_record_shards/connection_pool'
+require 'active_record_shards/migration'
 
 ActiveRecord::Base.extend(ActiveRecordShards::ConfigurationParser)
 ActiveRecord::Base.extend(ActiveRecordShards::Model)

data/test/connection_switching_test.rb

@@ -48,8 +48,10 @@ class ConnectionSwitchenTest < ActiveSupport::TestCase
         assert_not_equal(@shard_0_master.select_value("SELECT DATABASE()"), @shard_1_master.select_value("SELECT DATABASE()"))
 
         @database_names = []
-        ActiveRecord::Base.on_all_shards do
+        @database_shards = []
+        ActiveRecord::Base.on_all_shards do |shard|
           @database_names << ActiveRecord::Base.connection.select_value("SELECT DATABASE()")
+          @database_shards << shard
         end
       end
 
@@ -57,6 +59,9 @@ class ConnectionSwitchenTest < ActiveSupport::TestCase
         assert_equal(2, @database_names.size)
         assert_contains(@database_names, @shard_0_master.select_value("SELECT DATABASE()"))
         assert_contains(@database_names, @shard_1_master.select_value("SELECT DATABASE()"))
+        assert_equal(2, @database_shards.size)
+        assert_contains(@database_shards, "0")
+        assert_contains(@database_shards, "1")
       end
     end
   end

data/test/database.yml

@@ -4,8 +4,7 @@ test:
   database: ars_test
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
   shard_names: ['0', '1']
   
 test_slave:
@@ -14,8 +13,7 @@ test_slave:
   database: ars_test_slave
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
 
 test_shard_0:
   adapter: mysql
@@ -23,8 +21,7 @@ test_shard_0:
   database: ars_test_shard0
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
 
 test_shard_0_slave:
   adapter: mysql
@@ -32,8 +29,7 @@ test_shard_0_slave:
   database: ars_test_shard0_slave
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
 
 test_shard_1:
   adapter: mysql
@@ -41,8 +37,7 @@ test_shard_1:
   database: ars_test_shard1
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
 
 test_shard_1_slave:
   adapter: mysql
@@ -50,8 +45,7 @@ test_shard_1_slave:
   database: ars_test_shard1_slave
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
 
 test2:
   adapter: mysql
@@ -59,8 +53,7 @@ test2:
   database: ars_test2
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
 
 test2_slave:
   adapter: mysql
@@ -68,8 +61,7 @@ test2_slave:
   database: ars_test2_slave
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
 
 alternative:
   adapter: mysql
@@ -77,6 +69,5 @@ alternative:
   database: ars_test_alternative
   username: root
   password:
-  socket: /tmp/mysql.sock
-  host: localhost
+  host: 127.0.0.1
   shard_names: ['0', '1']

data/test/failure_migration/20110824010215_failure_migration.rb

@@ -0,0 +1,14 @@
+class FailureMigration < ActiveRecord::Migration
+  shard :all
+
+  def self.up
+    $fail_at_two ||= 0
+    $fail_at_two += 1
+    raise "FAIL FAIL FAIL" if $fail_at_two == 2
+    add_column :tickets, :sharded_column, :integer
+  end
+
+  def self.down
+    remove_column :tickets, :sharded_column
+  end
+end

data/test/helper.rb

@@ -18,17 +18,16 @@ RAILS_ENV = "test"
 ActiveRecord::Base.logger = Logger.new(File.dirname(__FILE__) + "/test.log")
 ActiveRecord::Base.configurations = YAML::load(IO.read(File.dirname(__FILE__) + '/database.yml'))
 
-ActiveRecord::Base.configurations.each do |name, conf|
-  puts "Setting up the #{conf['database']} db"
-  begin
-    ActiveRecord::Base.establish_connection(name)
-    ActiveRecord::Base.connection
-  rescue Mysql::Error => e
+def init_schema
+  ActiveRecord::Base.configurations.each do |name, conf|
+    puts "Setting up the #{conf['database']} db"
+    `echo "drop DATABASE if exists #{conf['database']}" | mysql --user=#{conf['username']}`
     `echo "create DATABASE #{conf['database']}" | mysql --user=#{conf['username']}`
     ActiveRecord::Base.establish_connection(name)
+    load(File.dirname(__FILE__) + "/schema.rb")
   end
-  load(File.dirname(__FILE__) + "/schema.rb")
 end
+init_schema
 
 require 'models'
 

data/test/migrations/20110824010216_shard_migration.rb

@@ -0,0 +1,12 @@
+#This is to recreate the deleted onetime payment for DealFish
+class ShardMigration < ActiveRecord::Migration
+  shard :all
+
+  def self.up
+    add_column :emails, :sharded_column, :integer
+  end
+
+  def self.down
+    remove_column :emails, :sharded_column
+  end
+end

data/test/migrations/20110829215912_account_migration.rb

@@ -0,0 +1,13 @@
+# See ticket 130214
+# Setting up an odd initial three year term for box.net to be paid by credit card by the normal means
+class AccountMigration < ActiveRecord::Migration
+  shard :none
+
+  def self.up
+    add_column :accounts, :non_sharded_column, :integer
+  end
+
+  def self.down
+    remove_column :accounts, :non_sharded_column
+  end
+end

data/test/migrator_test.rb

@@ -0,0 +1,88 @@
+require File.expand_path('helper', File.dirname(__FILE__))
+
+class CowardlyMigration < ActiveRecord::Migration
+  def self.up
+    "not gonna happen"
+  end
+
+  def self.down
+    "uh uh"
+  end
+end
+
+class MigratorTest < ActiveSupport::TestCase
+  def setup
+    init_schema
+  end
+
+  def test_migrator
+    migration_path = File.join(File.dirname(__FILE__), "/migrations")
+    ActiveRecord::Migrator.migrate(migration_path)
+    ActiveRecord::Base.on_all_shards do
+      assert ActiveRecord::Base.connection.table_exists?(:schema_migrations), "Schema Migrations doesn't exist"
+      assert ActiveRecord::Base.connection.table_exists?(:accounts)
+      assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20110824010216'")
+      assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20110829215912'")
+    end
+
+    ActiveRecord::Base.on_all_shards do
+      assert table_has_column?("emails", "sharded_column")
+      assert !table_has_column?("accounts", "non_sharded_column")
+    end
+
+    ActiveRecord::Base.on_shard(nil) do
+      assert !table_has_column?("emails", "sharded_column")
+      assert table_has_column?("accounts", "non_sharded_column")
+    end
+
+    # now test down/ up
+    ActiveRecord::Migrator.run(:down, migration_path, 20110824010216)
+    ActiveRecord::Base.on_all_shards do
+      assert !table_has_column?("emails", "sharded_column")
+    end
+
+    ActiveRecord::Migrator.run(:down, migration_path, 20110829215912)
+    ActiveRecord::Base.on_shard(nil) do
+      assert !table_has_column?("accounts", "non_sharded_column")
+    end
+
+    ActiveRecord::Migrator.run(:up, migration_path, 20110824010216)
+    ActiveRecord::Base.on_all_shards do
+      assert table_has_column?("emails", "sharded_column")
+    end
+
+    ActiveRecord::Migrator.run(:up, migration_path, 20110829215912)
+    ActiveRecord::Base.on_shard(nil) do
+      assert table_has_column?("accounts", "non_sharded_column")
+    end
+  end
+
+  def test_bad_migration
+    assert_raises(RuntimeError) do
+      CowardlyMigration.up
+    end
+
+    assert_raises(RuntimeError) do
+      CowardlyMigration.down
+    end
+  end
+
+  def test_failing_migration
+    # like, if you have to break a migration in the middle somewhere.
+    begin
+      migration_path = File.join(File.dirname(__FILE__), "/failure_migration")
+      ActiveRecord::Migrator.migrate(migration_path)
+    rescue
+      retry
+    end
+
+    ActiveRecord::Base.on_all_shards do
+      assert table_has_column?("tickets", "sharded_column")
+    end
+  end
+
+  private
+  def table_has_column?(table, column)
+    ActiveRecord::Base.connection.select_value("show create table #{table}") =~ /#{column}/
+  end
+end

data/test/schema.rb

@@ -1,24 +1,24 @@
 ActiveRecord::Schema.define(:version => 1) do
+  suppress_messages do
+    create_table "accounts", :force => true do |t|
+      t.string   "name"
+      t.datetime "created_at"
+      t.datetime "updated_at"
+    end
 
-  create_table "accounts", :force => true do |t|
-    t.string   "name"
-    t.datetime "created_at"
-    t.datetime "updated_at"
-  end
-
-  create_table "emails", :force => true do |t|
-    t.string   "from"
-    t.string   "to"
-    t.text     "mail"
-    t.datetime "created_at"
-    t.datetime "updated_at"
-  end
+    create_table "emails", :force => true do |t|
+      t.string   "from"
+      t.string   "to"
+      t.text     "mail"
+      t.datetime "created_at"
+      t.datetime "updated_at"
+    end
 
-  create_table "tickets", :force => true do |t|
-    t.string   "title"
-    t.integer  "account_id"
-    t.datetime "created_at"
-    t.datetime "updated_at"
+    create_table "tickets", :force => true do |t|
+      t.string   "title"
+      t.integer  "account_id"
+      t.datetime "created_at"
+      t.datetime "updated_at"
+    end
   end
-
-end
+end

metadata

@@ -1,22 +1,23 @@
 --- !ruby/object:Gem::Specification 
 name: active_record_shards
 version: !ruby/object:Gem::Version 
-  hash: 1
+  hash: 29
   prerelease: 
   segments: 
   - 2
-  - 3
+  - 4
   - 1
-  version: 2.3.1
+  version: 2.4.1
 platform: ruby
 authors: 
 - Mick Staugaard
 - Eric Chapweske
+- Ben Osheroff
 autorequire: 
 bindir: bin
 cert_chain: []
 
-date: 2011-07-20 00:00:00 -07:00
+date: 2011-08-31 00:00:00 -07:00
 default_executable: 
 dependencies: 
 - !ruby/object:Gem::Dependency 
@@ -109,6 +110,7 @@ description: Easily run queries on shard and slave databases.
 email: 
 - mick@staugaard.com
 - eac@zendesk.com
+- ben@gimbo.net
 executables: []
 
 extensions: []
@@ -120,6 +122,7 @@ files:
 - lib/active_record_shards/configuration_parser.rb
 - lib/active_record_shards/connection_pool.rb
 - lib/active_record_shards/connection_switcher.rb
+- lib/active_record_shards/migration.rb
 - lib/active_record_shards/model.rb
 - lib/active_record_shards/shard_selection.rb
 - lib/active_record_shards/tasks.rb
@@ -129,10 +132,13 @@ files:
 - test/connection_switching_test.rb
 - test/database.yml
 - test/database_parse_test.yml
+- test/failure_migration/20110824010215_failure_migration.rb
 - test/helper.rb
+- test/migrations/20110824010216_shard_migration.rb
+- test/migrations/20110829215912_account_migration.rb
+- test/migrator_test.rb
 - test/models.rb
 - test/schema.rb
-- test/test.log
 has_rdoc: true
 homepage: http://github.com/staugaard/active_record_shards
 licenses: []
@@ -172,7 +178,10 @@ test_files:
 - test/connection_switching_test.rb
 - test/database.yml
 - test/database_parse_test.yml
+- test/failure_migration/20110824010215_failure_migration.rb
 - test/helper.rb
+- test/migrations/20110824010216_shard_migration.rb
+- test/migrations/20110829215912_account_migration.rb
+- test/migrator_test.rb
 - test/models.rb
 - test/schema.rb
-- test/test.log