├── test ├── schemas │ ├── ars_test3.sql │ ├── ars_test3_shard0.sql │ ├── ars_test_shard0.sql │ ├── ars_test_shard1.sql │ ├── ars_test_shard0_replica.sql │ ├── ars_test_shard1_replica.sql │ ├── ars_test2.sql │ ├── ars_test2_replica.sql │ ├── ars_test.sql │ └── ars_test_replica.sql ├── support │ ├── cowardly_migration │ │ └── 20110824010215_cowardly_migration.rb │ ├── migrations │ │ ├── 20110824010216_shard_migration.rb │ │ └── 20110829215912_account_migration.rb │ ├── separate_migrations │ │ ├── 20190121112234_separate_sharded_migration.rb │ │ └── 20190121112233_separate_unsharded_migration.rb │ ├── sharded_schema.rb │ ├── failure_migration │ │ └── 20110824010215_failure_migration.rb │ ├── unsharded_schema.rb │ ├── database_tasks.yml │ ├── db_helper.rb │ └── tcp_proxy.rb ├── models.rb ├── active_record_shards │ ├── sql_comments_test.rb │ ├── schema_dumper_extension_test.rb │ └── configuration_parser_test.rb ├── active_record │ └── connection_adapters │ │ └── mysql2_fake_adapter.rb ├── database.yml ├── active_record_shards_test.rb ├── check_performance.rb ├── tasks_test.rb ├── migrator_test.rb ├── helper.rb ├── thread_safety_test.rb ├── on_replica_by_default_test.rb └── connection_switching_test.rb ├── Gemfile ├── .github ├── term-check.yaml ├── CODEOWNERS └── workflows │ ├── publish.yml │ └── ci.yml ├── .document ├── gemfiles ├── rails5.1.gemfile ├── rails5.2.gemfile ├── rails6.0.gemfile ├── rails6.1.gemfile ├── rails7.0.gemfile └── common.rb ├── .gitignore ├── Rakefile ├── lib ├── active_record_shards │ ├── default_shard.rb │ ├── sql_comments.rb │ ├── connection_switcher-5-1.rb │ ├── schema_dumper_extension.rb │ ├── shard_support.rb │ ├── connection_switcher-6-0.rb │ ├── connection_switcher-7-0.rb │ ├── connection_switcher-6-1.rb │ ├── association_collection_connection_selection.rb │ ├── configuration_parser.rb │ ├── model.rb │ ├── shard_selection.rb │ ├── migration.rb │ ├── tasks.rb │ ├── default_replica_patches.rb │ └── connection_switcher.rb └── active_record_shards.rb ├── LICENSE ├── active_record_shards.gemspec ├── .rubocop.yml ├── Changelog.md └── README.md /test/schemas/ars_test3.sql: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /test/schemas/ars_test3_shard0.sql: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'gemfiles/rails6.0.gemfile' 2 | -------------------------------------------------------------------------------- /.github/term-check.yaml: -------------------------------------------------------------------------------- 1 | ignore: 2 | - Changelog.md 3 | -------------------------------------------------------------------------------- /.github/CODEOWNERS: -------------------------------------------------------------------------------- 1 | * @zendesk/ruby-core @zendesk/database-gem-owners 2 | -------------------------------------------------------------------------------- /.document: -------------------------------------------------------------------------------- 1 | README.rdoc 2 | lib/**/*.rb 3 | bin/* 4 | features/**/*.feature 5 | LICENSE 6 | -------------------------------------------------------------------------------- /gemfiles/rails5.1.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem 'activerecord', '~> 5.1.6' 4 | -------------------------------------------------------------------------------- /gemfiles/rails5.2.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem 'activerecord', '~> 5.2.0' 4 | -------------------------------------------------------------------------------- /gemfiles/rails6.0.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem 'activerecord', '~> 6.0.0' 4 | -------------------------------------------------------------------------------- /gemfiles/rails6.1.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem 'activerecord', '~> 6.1.0' 4 | -------------------------------------------------------------------------------- /gemfiles/rails7.0.gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'common.rb' 2 | 3 | gem 'activerecord', '~> 7.0.0' 4 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | pkg/* 2 | *.gem 3 | .bundle 4 | *.log 5 | Gemfile.lock 6 | gemfiles/*.lock 7 | .tags* 8 | tmp/* 9 | -------------------------------------------------------------------------------- /test/support/cowardly_migration/20110824010215_cowardly_migration.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | class CowardlyMigration < BaseMigration 4 | def self.up 5 | # shouldn't get here. must fail. 6 | end 7 | 8 | def self.down 9 | end 10 | end 11 | -------------------------------------------------------------------------------- /gemfiles/common.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | source 'https://rubygems.org' 4 | 5 | gemspec path: Bundler.root.to_s.sub('/gemfiles', '') 6 | 7 | group :test do 8 | gem 'pry-byebug', platforms: [:mri] 9 | end 10 | 11 | gem 'benchmark-ips' 12 | gem 'debug', '>= 1.0.0' 13 | -------------------------------------------------------------------------------- /test/support/migrations/20110824010216_shard_migration.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | class ShardMigration < BaseMigration 4 | shard :all 5 | 6 | def self.up 7 | add_column :tickets, :sharded_column, :integer 8 | end 9 | 10 | def self.down 11 | remove_column :tickets, :sharded_column 12 | end 13 | end 14 | -------------------------------------------------------------------------------- /test/support/separate_migrations/20190121112234_separate_sharded_migration.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | class SeparateShardedMigration < BaseMigration 4 | shard :all 5 | 6 | def self.up 7 | create_table :sharded_table do |t| 8 | t.string :name 9 | end 10 | end 11 | 12 | def self.down 13 | end 14 | end 15 | -------------------------------------------------------------------------------- /test/support/migrations/20110829215912_account_migration.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | class AccountMigration < BaseMigration 4 | shard :none 5 | 6 | def self.up 7 | add_column :accounts, :non_sharded_column, :integer 8 | end 9 | 10 | def self.down 11 | remove_column :accounts, :non_sharded_column 12 | end 13 | end 14 | -------------------------------------------------------------------------------- /test/support/separate_migrations/20190121112233_separate_unsharded_migration.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | class SeparateUnshardedMigration < BaseMigration 4 | shard :none 5 | 6 | def self.up 7 | create_table :unsharded_table do |t| 8 | t.string :name 9 | end 10 | end 11 | 12 | def self.down 13 | end 14 | end 15 | -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish Gem 2 | 3 | on: 4 | push: 5 | tags: v* 6 | 7 | jobs: 8 | call-workflow: 9 | uses: zendesk/gw/.github/workflows/ruby-gem-publication.yml@main 10 | secrets: 11 | RUBY_GEMS_API_KEY: ${{ secrets.RUBY_GEMS_API_KEY }} 12 | RUBY_GEMS_TOTP_DEVICE: ${{ secrets.RUBY_GEMS_TOTP_DEVICE }} 13 | -------------------------------------------------------------------------------- /test/support/sharded_schema.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | ActiveRecord::Migration.verbose = false 4 | 5 | ActiveRecord::Schema.define(version: 1) do 6 | create_table "tickets", force: true do |t| 7 | t.string "title" 8 | t.integer "account_id" 9 | t.datetime "created_at" 10 | t.datetime "updated_at" 11 | end 12 | end 13 | -------------------------------------------------------------------------------- /test/support/failure_migration/20110824010215_failure_migration.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | class FailureMigration < BaseMigration 4 | shard :all 5 | 6 | def self.up 7 | @@fail_at_two ||= 0 8 | @@fail_at_two += 1 9 | raise "ERROR_IN_MIGRATION" if @@fail_at_two == 2 10 | 11 | add_column :tickets, :sharded_column, :integer 12 | end 13 | 14 | def self.down 15 | remove_column :tickets, :sharded_column 16 | end 17 | end 18 | -------------------------------------------------------------------------------- /Rakefile: -------------------------------------------------------------------------------- 1 | require 'bundler/setup' 2 | require 'bump/tasks' 3 | require 'rubocop/rake_task' 4 | 5 | Bundler::GemHelper.install_tasks 6 | 7 | require 'rake/testtask' 8 | Rake::TestTask.new(:test) do |test| 9 | test.pattern = './test/**/*_test.rb' 10 | test.verbose = false 11 | test.warning = true 12 | end 13 | 14 | task default: ["rubocop", "test"] 15 | 16 | RuboCop::RakeTask.new 17 | 18 | desc 'Run an IRB console with ActiveRecordShards loaded' 19 | task :console do 20 | require 'irb' 21 | require 'irb/completion' 22 | require 'active_record_shards' 23 | ARGV.clear 24 | IRB.start 25 | end 26 | -------------------------------------------------------------------------------- /test/support/unsharded_schema.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | ActiveRecord::Migration.verbose = false 4 | 5 | ActiveRecord::Schema.define(version: 1) do 6 | create_table "accounts", force: true do |t| 7 | t.string "name" 8 | t.datetime "created_at" 9 | t.datetime "updated_at" 10 | end 11 | 12 | create_table "account_things", force: true do |t| 13 | t.integer "account_id" 14 | t.boolean "enabled", default: true 15 | end 16 | 17 | create_table "account_people", force: true, id: false do |t| 18 | t.integer "account_id" 19 | t.integer "person_id" 20 | end 21 | 22 | create_table "people", force: true do |t| 23 | t.string "name" 24 | t.string "type" 25 | t.datetime "created_at" 26 | t.datetime "updated_at" 27 | end 28 | end 29 | -------------------------------------------------------------------------------- /lib/active_record_shards/default_shard.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module ActiveRecordShards 4 | module DefaultShard 5 | def default_shard=(new_default_shard) 6 | if ars_shard_type?(new_default_shard) 7 | ActiveRecordShards::ShardSelection.ars_default_shard = new_default_shard 8 | switch_connection(shard: new_default_shard) 9 | else 10 | super 11 | end 12 | end 13 | 14 | private 15 | 16 | def ars_shard_type?(shard) 17 | return true if ActiveRecord.version < Gem::Version.new('6.1') 18 | return true if shard.nil? 19 | return true if shard == :_no_shard 20 | return true if shard.is_a?(Integer) 21 | 22 | false 23 | end 24 | end 25 | end 26 | 27 | ActiveRecord::Base.singleton_class.prepend(ActiveRecordShards::DefaultShard) 28 | -------------------------------------------------------------------------------- /test/support/database_tasks.yml: -------------------------------------------------------------------------------- 1 | <% mysql = URI(ENV['MYSQL_URL'] || 'mysql://root@127.0.0.1:3306') %> 2 | 3 | mysql: &MYSQL 4 | encoding: utf8 5 | username: <%= mysql.user %> 6 | password: <%= mysql.password %> 7 | host: <%= mysql.host %> 8 | port: <%= mysql.port %> 9 | 10 | test: 11 | <<: *MYSQL 12 | adapter: mysql2 13 | database: ars_tasks_test 14 | replica: 15 | database: ars_tasks_test_replica 16 | shards: 17 | 0: 18 | database: ars_tasks_test_shard_a 19 | 1: 20 | database: ars_tasks_test_shard_b 21 | 22 | test_adapter: 23 | <<: *MYSQL 24 | adapter: mysql2_fake 25 | database: ars_tasks_adapter_test 26 | replica: 27 | database: ars_tasks_adapter_test_replica 28 | shards: 29 | 0: 30 | database: ars_tasks_adapter_test_shard_a 31 | 1: 32 | database: ars_tasks_adapter_test_shard_b 33 | -------------------------------------------------------------------------------- /test/models.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | class Account < ActiveRecord::Base 4 | # attributes: id, name, updated_at, created_at 5 | not_sharded 6 | 7 | has_many :tickets 8 | has_many :account_things 9 | has_and_belongs_to_many :people, join_table: 'account_people' 10 | end 11 | 12 | class AccountThing < ActiveRecord::Base 13 | not_sharded 14 | 15 | scope(:enabled, -> { where(enabled: true) }) 16 | end 17 | 18 | class AccountInherited < Account 19 | end 20 | 21 | class Ticket < ActiveRecord::Base 22 | # attributes: id, title, account_id, updated_at, created_at 23 | belongs_to :account 24 | end 25 | 26 | class Person < ActiveRecord::Base 27 | not_sharded 28 | end 29 | 30 | class User < Person 31 | # Makes `User.new` a bit more complicated. Don't change without changing the 32 | # corresponding tests. 33 | default_scope { where(type: 'User') } 34 | end 35 | -------------------------------------------------------------------------------- /lib/active_record_shards/sql_comments.rb: -------------------------------------------------------------------------------- 1 | # show which connection was picked to debug primary/replica slowness when both servers are the same 2 | module ActiveRecordShards 3 | module SqlComments 4 | module Methods 5 | def execute(query, name = nil, **kwargs) 6 | shard = ActiveRecord::Base.current_shard_selection.shard 7 | shard_text = shard ? "shard #{shard}" : 'unsharded' 8 | replica = ActiveRecord::Base.current_shard_selection.on_replica? 9 | replica_text = replica ? 'replica' : 'primary' 10 | query = "/* #{shard_text} #{replica_text} */ " + query 11 | super(query, name, **kwargs) 12 | end 13 | end 14 | 15 | def self.enable 16 | ActiveRecord::Base.on_replica do 17 | ActiveRecord::Base.on_shard(nil) do 18 | ActiveRecord::Base.connection.class.prepend(Methods) 19 | end 20 | end 21 | end 22 | end 23 | end 24 | -------------------------------------------------------------------------------- /test/schemas/ars_test_shard0.sql: -------------------------------------------------------------------------------- 1 | DROP TABLE IF EXISTS `ar_internal_metadata`; 2 | CREATE TABLE `ar_internal_metadata` ( 3 | `key` varchar(255) CHARACTER SET utf8 NOT NULL, 4 | `value` varchar(255) DEFAULT NULL, 5 | `created_at` datetime NOT NULL, 6 | `updated_at` datetime NOT NULL, 7 | PRIMARY KEY (`key`) 8 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 9 | 10 | DROP TABLE IF EXISTS `schema_migrations`; 11 | CREATE TABLE `schema_migrations` ( 12 | `version` varchar(255) CHARACTER SET utf8 NOT NULL, 13 | PRIMARY KEY (`version`) 14 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 15 | 16 | DROP TABLE IF EXISTS `tickets`; 17 | CREATE TABLE `tickets` ( 18 | `id` int(11) NOT NULL AUTO_INCREMENT, 19 | `title` varchar(255) DEFAULT NULL, 20 | `account_id` int(11) DEFAULT NULL, 21 | `created_at` datetime DEFAULT NULL, 22 | `updated_at` datetime DEFAULT NULL, 23 | PRIMARY KEY (`id`) 24 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 25 | 26 | INSERT INTO `schema_migrations` (version) VALUES ('1'); 27 | -------------------------------------------------------------------------------- /test/schemas/ars_test_shard1.sql: -------------------------------------------------------------------------------- 1 | DROP TABLE IF EXISTS `ar_internal_metadata`; 2 | CREATE TABLE `ar_internal_metadata` ( 3 | `key` varchar(255) CHARACTER SET utf8 NOT NULL, 4 | `value` varchar(255) DEFAULT NULL, 5 | `created_at` datetime NOT NULL, 6 | `updated_at` datetime NOT NULL, 7 | PRIMARY KEY (`key`) 8 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 9 | 10 | DROP TABLE IF EXISTS `schema_migrations`; 11 | CREATE TABLE `schema_migrations` ( 12 | `version` varchar(255) CHARACTER SET utf8 NOT NULL, 13 | PRIMARY KEY (`version`) 14 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 15 | 16 | DROP TABLE IF EXISTS `tickets`; 17 | CREATE TABLE `tickets` ( 18 | `id` int(11) NOT NULL AUTO_INCREMENT, 19 | `title` varchar(255) DEFAULT NULL, 20 | `account_id` int(11) DEFAULT NULL, 21 | `created_at` datetime DEFAULT NULL, 22 | `updated_at` datetime DEFAULT NULL, 23 | PRIMARY KEY (`id`) 24 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 25 | 26 | INSERT INTO `schema_migrations` (version) VALUES ('1'); 27 | -------------------------------------------------------------------------------- /test/schemas/ars_test_shard0_replica.sql: -------------------------------------------------------------------------------- 1 | DROP TABLE IF EXISTS `ar_internal_metadata`; 2 | CREATE TABLE `ar_internal_metadata` ( 3 | `key` varchar(255) CHARACTER SET utf8 NOT NULL, 4 | `value` varchar(255) DEFAULT NULL, 5 | `created_at` datetime NOT NULL, 6 | `updated_at` datetime NOT NULL, 7 | PRIMARY KEY (`key`) 8 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 9 | 10 | DROP TABLE IF EXISTS `schema_migrations`; 11 | CREATE TABLE `schema_migrations` ( 12 | `version` varchar(255) CHARACTER SET utf8 NOT NULL, 13 | PRIMARY KEY (`version`) 14 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 15 | 16 | DROP TABLE IF EXISTS `tickets`; 17 | CREATE TABLE `tickets` ( 18 | `id` int(11) NOT NULL AUTO_INCREMENT, 19 | `title` varchar(255) DEFAULT NULL, 20 | `account_id` int(11) DEFAULT NULL, 21 | `created_at` datetime DEFAULT NULL, 22 | `updated_at` datetime DEFAULT NULL, 23 | PRIMARY KEY (`id`) 24 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 25 | 26 | INSERT INTO `schema_migrations` (version) VALUES ('1'); 27 | -------------------------------------------------------------------------------- /test/schemas/ars_test_shard1_replica.sql: -------------------------------------------------------------------------------- 1 | DROP TABLE IF EXISTS `ar_internal_metadata`; 2 | CREATE TABLE `ar_internal_metadata` ( 3 | `key` varchar(255) CHARACTER SET utf8 NOT NULL, 4 | `value` varchar(255) DEFAULT NULL, 5 | `created_at` datetime NOT NULL, 6 | `updated_at` datetime NOT NULL, 7 | PRIMARY KEY (`key`) 8 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 9 | 10 | DROP TABLE IF EXISTS `schema_migrations`; 11 | CREATE TABLE `schema_migrations` ( 12 | `version` varchar(255) CHARACTER SET utf8 NOT NULL, 13 | PRIMARY KEY (`version`) 14 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 15 | 16 | DROP TABLE IF EXISTS `tickets`; 17 | CREATE TABLE `tickets` ( 18 | `id` int(11) NOT NULL AUTO_INCREMENT, 19 | `title` varchar(255) DEFAULT NULL, 20 | `account_id` int(11) DEFAULT NULL, 21 | `created_at` datetime DEFAULT NULL, 22 | `updated_at` datetime DEFAULT NULL, 23 | PRIMARY KEY (`id`) 24 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 25 | 26 | INSERT INTO `schema_migrations` (version) VALUES ('1'); 27 | -------------------------------------------------------------------------------- /lib/active_record_shards/connection_switcher-5-1.rb: -------------------------------------------------------------------------------- 1 | module ActiveRecordShards 2 | module ConnectionSwitcher 3 | def connection_specification_name 4 | name = current_shard_selection.resolve_connection_name(sharded: is_sharded?, configurations: configurations) 5 | 6 | unless configurations[name] || name == "primary" 7 | raise ActiveRecord::AdapterNotSpecified, "No database defined by #{name} in your database config. (configurations: #{configurations.to_h.keys.inspect})" 8 | end 9 | 10 | name 11 | end 12 | 13 | private 14 | 15 | def ensure_shard_connection 16 | # See if we've connected before. If not, call `#establish_connection` 17 | # so that ActiveRecord can resolve connection_specification_name to an 18 | # ARS connection. 19 | spec_name = connection_specification_name 20 | 21 | pool = connection_handler.retrieve_connection_pool(spec_name) 22 | connection_handler.establish_connection(spec_name.to_sym) if pool.nil? 23 | end 24 | end 25 | end 26 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2009 Zendesk 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining 4 | a copy of this software and associated documentation files (the 5 | "Software"), to deal in the Software without restriction, including 6 | without limitation the rights to use, copy, modify, merge, publish, 7 | distribute, sublicense, and/or sell copies of the Software, and to 8 | permit persons to whom the Software is furnished to do so, subject to 9 | the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be 12 | included in all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 15 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 16 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 17 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 18 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 19 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 20 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /test/active_record_shards/sql_comments_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative '../helper' 4 | require 'active_record_shards/sql_comments' 5 | require 'active_record/connection_adapters/mysql2_adapter' 6 | 7 | describe ActiveRecordShards::SqlComments do 8 | with_fresh_databases 9 | 10 | class CustomAdapter < ActiveRecord::ConnectionAdapters::Mysql2Adapter 11 | prepend ActiveRecordShards::SqlComments::Methods 12 | end 13 | 14 | before do 15 | ActiveRecord::Base.establish_connection(RAILS_ENV.to_sym) 16 | end 17 | 18 | it "adds sql comment" do 19 | old_logger = ActiveRecord::Base.logger 20 | new_logger = StringIO.new 21 | ActiveRecord::Base.logger = Logger.new(new_logger) 22 | config = Account.connection.instance_variable_get(:@config) 23 | custom_connection = CustomAdapter.new(Mysql2::Client.new(config), nil, nil, config) 24 | 25 | Account.stub :connection, custom_connection do 26 | Account.first 27 | end 28 | 29 | assert_includes(new_logger.string, "/* unsharded primary */") 30 | ensure 31 | ActiveRecord::Base.logger = old_logger 32 | end 33 | end 34 | -------------------------------------------------------------------------------- /lib/active_record_shards/schema_dumper_extension.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module ActiveRecordShards 4 | module SchemaDumperExtension 5 | def dump(stream) 6 | stream = super(stream) 7 | original_connection = @connection 8 | 9 | if ActiveRecord::Base.supports_sharding? 10 | ActiveRecord::Base.on_first_shard do 11 | @connection = ActiveRecord::Base.connection 12 | shard_header(stream) 13 | extensions(stream) 14 | tables(stream) 15 | shard_trailer(stream) 16 | end 17 | end 18 | 19 | stream 20 | ensure 21 | @connection = original_connection 22 | end 23 | 24 | def shard_header(stream) 25 | define_params = @version ? "version: #{@version}" : "" 26 | 27 | stream.puts <<~HEADER 28 | 29 | 30 | # This section generated by active_record_shards 31 | 32 | ActiveRecord::Base.on_all_shards do 33 | ActiveRecord::Schema.define(#{define_params}) do 34 | 35 | HEADER 36 | end 37 | 38 | def shard_trailer(stream) 39 | stream.puts "end\nend" 40 | end 41 | end 42 | end 43 | -------------------------------------------------------------------------------- /active_record_shards.gemspec: -------------------------------------------------------------------------------- 1 | Gem::Specification.new "active_record_shards", "5.5.1" do |s| 2 | s.authors = ["Benjamin Quorning", "Gabe Martin-Dempesy", "Pierre Schambacher", "Mick Staugaard", "Eric Chapweske", "Ben Osheroff"] 3 | s.email = ["bquorning@zendesk.com", "gabe@zendesk.com", "pschambacher@zendesk.com", "mick@staugaard.com"] 4 | s.homepage = "https://github.com/zendesk/active_record_shards" 5 | s.summary = "Simple database switching for ActiveRecord." 6 | s.description = "Easily run queries on shard and replica databases." 7 | s.license = "MIT" 8 | 9 | s.required_ruby_version = ">= 2.6" 10 | 11 | s.add_runtime_dependency("activerecord", ">= 5.1", "< 7.1") 12 | s.add_runtime_dependency("activesupport", ">= 5.1", "< 7.1") 13 | 14 | s.add_development_dependency("bump") 15 | s.add_development_dependency("minitest", ">= 5.10.0") 16 | s.add_development_dependency("mysql2") 17 | s.add_development_dependency("rake", '~> 12.0') 18 | s.add_development_dependency("rubocop", "~> 0.77.0") 19 | s.add_development_dependency("rubocop-minitest", "~> 0.5.0") 20 | s.add_development_dependency("rubocop-performance", "~> 1.5.1") 21 | 22 | s.files = Dir["lib/**/*"] + ["README.md"] 23 | end 24 | -------------------------------------------------------------------------------- /lib/active_record_shards/shard_support.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module ActiveRecordShards 4 | class ShardSupport 5 | class ShardEnumerator 6 | include Enumerable 7 | 8 | def each(&block) 9 | ActiveRecord::Base.on_all_shards(&block) 10 | end 11 | end 12 | 13 | def initialize(scope) 14 | @scope = scope 15 | end 16 | 17 | def enum 18 | ShardEnumerator.new 19 | end 20 | 21 | def find(*find_args) 22 | ensure_concrete! 23 | 24 | exception = nil 25 | enum.each do 26 | record = @scope.find(*find_args) 27 | return record if record 28 | rescue ActiveRecord::RecordNotFound => e 29 | exception = e 30 | end 31 | raise exception 32 | end 33 | ruby2_keywords(:find) if respond_to?(:ruby2_keywords, true) 34 | 35 | def count 36 | enum.inject(0) { |accum, _shard| @scope.clone.count + accum } 37 | end 38 | 39 | def to_a 40 | enum.flat_map { @scope.clone.to_a } 41 | end 42 | 43 | private 44 | 45 | def ensure_concrete! 46 | raise "Please call this method on a concrete model, not an abstract class!" if @scope.abstract_class? 47 | end 48 | end 49 | end 50 | -------------------------------------------------------------------------------- /lib/active_record_shards/connection_switcher-6-0.rb: -------------------------------------------------------------------------------- 1 | module ActiveRecordShards 2 | module ConnectionSwitcher 3 | def connection_specification_name 4 | name = current_shard_selection.resolve_connection_name(sharded: is_sharded?, configurations: configurations) 5 | 6 | @_ars_connection_specification_names ||= {} 7 | unless @_ars_connection_specification_names.include?(name) 8 | unless configurations[name] || name == "primary" 9 | raise ActiveRecord::AdapterNotSpecified, "No database defined by #{name} in your database config. (configurations: #{configurations.to_h.keys.inspect})" 10 | end 11 | 12 | @_ars_connection_specification_names[name] = true 13 | end 14 | 15 | name 16 | end 17 | 18 | private 19 | 20 | def ensure_shard_connection 21 | # See if we've connected before. If not, call `#establish_connection` 22 | # so that ActiveRecord can resolve connection_specification_name to an 23 | # ARS connection. 24 | spec_name = connection_specification_name 25 | 26 | pool = connection_handler.retrieve_connection_pool(spec_name) 27 | connection_handler.establish_connection(spec_name.to_sym) if pool.nil? 28 | end 29 | end 30 | end 31 | -------------------------------------------------------------------------------- /test/active_record/connection_adapters/mysql2_fake_adapter.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | # ActiveRecord expects adapter definition to be under active_record/connection_adapter in the LOAD_PATH, see 4 | # https://github.com/rails/rails/blob/fdf3f0b9306ba8145e6e3acb84a50e5d23dfe48c/activerecord/lib/active_record/connection_adapters/connection_specification.rb#L168 5 | 6 | require 'active_record' 7 | require 'active_record/connection_adapters/mysql2_adapter' 8 | require 'mysql2' 9 | 10 | module ActiveRecord 11 | class Base 12 | def self.mysql2_fake_connection(config) 13 | # Based on `mysql2_connection`: https://github.com/rails/rails/blob/fdf3f0b9306ba8145e6e3acb84a50e5d23dfe48c/activerecord/lib/active_record/connection_adapters/mysql2_adapter.rb#L12 14 | config = config.symbolize_keys 15 | config[:flags] ||= 0 16 | 17 | if config[:flags].is_a? Array 18 | config[:flags].push "FOUND_ROWS" 19 | else 20 | config[:flags] |= Mysql2::Client::FOUND_ROWS 21 | end 22 | 23 | client = Mysql2::Client.new(config) 24 | ConnectionAdapters::Mysql2FakeAdapter.new(client, logger, nil, config) 25 | end 26 | end 27 | 28 | module ConnectionAdapters 29 | class Mysql2FakeAdapter < Mysql2Adapter; end 30 | end 31 | end 32 | -------------------------------------------------------------------------------- /lib/active_record_shards/connection_switcher-7-0.rb: -------------------------------------------------------------------------------- 1 | module ActiveRecordShards 2 | module ConnectionSwitcher 3 | def connection_specification_name 4 | name = current_shard_selection.resolve_connection_name(sharded: is_sharded?, configurations: configurations) 5 | 6 | @_ars_connection_specification_names ||= {} 7 | unless @_ars_connection_specification_names.include?(name) 8 | unless configurations.configs_for(env_name: name, include_hidden: true).any? || name == "ActiveRecord::Base" 9 | raise ActiveRecord::AdapterNotSpecified, "No database defined by #{name} in your database config. (configurations: #{configurations.configurations.map(&:env_name).inspect})" 10 | end 11 | 12 | @_ars_connection_specification_names[name] = true 13 | end 14 | 15 | name 16 | end 17 | 18 | private 19 | 20 | def ensure_shard_connection 21 | # See if we've connected before. If not, call `#establish_connection` 22 | # so that ActiveRecord can resolve connection_specification_name to an 23 | # ARS connection. 24 | spec_name = connection_specification_name 25 | 26 | pool = connection_handler.retrieve_connection_pool(spec_name) 27 | 28 | connection_handler.establish_connection(spec_name.to_sym) if pool.nil? 29 | end 30 | end 31 | end 32 | -------------------------------------------------------------------------------- /lib/active_record_shards/connection_switcher-6-1.rb: -------------------------------------------------------------------------------- 1 | module ActiveRecordShards 2 | module ConnectionSwitcher 3 | def connection_specification_name 4 | name = current_shard_selection.resolve_connection_name(sharded: is_sharded?, configurations: configurations) 5 | 6 | @_ars_connection_specification_names ||= {} 7 | unless @_ars_connection_specification_names.include?(name) 8 | unless configurations.configs_for(env_name: name, include_replicas: true).any? || name == "ActiveRecord::Base" 9 | raise ActiveRecord::AdapterNotSpecified, "No database defined by #{name} in your database config. (configurations: #{configurations.configurations.map(&:env_name).inspect})" 10 | end 11 | 12 | @_ars_connection_specification_names[name] = true 13 | end 14 | 15 | name 16 | end 17 | 18 | private 19 | 20 | def ensure_shard_connection 21 | # See if we've connected before. If not, call `#establish_connection` 22 | # so that ActiveRecord can resolve connection_specification_name to an 23 | # ARS connection. 24 | spec_name = connection_specification_name 25 | 26 | pool = connection_handler.retrieve_connection_pool(spec_name) 27 | 28 | connection_handler.establish_connection(spec_name.to_sym) if pool.nil? 29 | end 30 | end 31 | end 32 | -------------------------------------------------------------------------------- /lib/active_record_shards/association_collection_connection_selection.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module ActiveRecordShards 4 | module AssociationCollectionConnectionSelection 5 | def on_replica_if(condition) 6 | condition ? on_replica : self 7 | end 8 | 9 | def on_replica_unless(condition) 10 | on_replica_if(!condition) 11 | end 12 | 13 | def on_primary_if(condition) 14 | condition ? on_primary : self 15 | end 16 | 17 | def on_primary_unless(condition) 18 | on_primary_if(!condition) 19 | end 20 | 21 | def on_replica 22 | PrimaryReplicaProxy.new(self, :replica) 23 | end 24 | 25 | def on_primary 26 | PrimaryReplicaProxy.new(self, :primary) 27 | end 28 | 29 | class PrimaryReplicaProxy 30 | def initialize(association_collection, which) 31 | @association_collection = association_collection 32 | @which = which 33 | end 34 | 35 | def method_missing(method, *args, &block) # rubocop:disable Style/MethodMissingSuper, Style/MissingRespondToMissing 36 | reflection = @association_collection.proxy_association.reflection 37 | reflection.klass.on_cx_switch_block(@which) { @association_collection.send(method, *args, &block) } 38 | end 39 | ruby2_keywords(:method_missing) if respond_to?(:ruby2_keywords, true) 40 | end 41 | end 42 | end 43 | -------------------------------------------------------------------------------- /test/active_record_shards/schema_dumper_extension_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative '../helper' 4 | 5 | describe ActiveRecordShards::SchemaDumperExtension do 6 | describe "schema dump" do 7 | let(:schema_file) { Tempfile.new('active_record_shards_schema.rb') } 8 | 9 | with_fresh_databases 10 | 11 | before do 12 | ActiveRecord::Base.establish_connection(RAILS_ENV.to_sym) 13 | 14 | # create shard-specific columns 15 | ActiveRecord::Migrator.migrations_paths = [File.join(__dir__, "../support/migrations")] 16 | migrator.migrate 17 | end 18 | 19 | after { schema_file.unlink } 20 | 21 | it "includes the sharded tables" do 22 | ActiveRecord::SchemaDumper.dump(ActiveRecord::Base.connection, schema_file) 23 | schema_file.close 24 | 25 | # Recreate the database without loading the schema 26 | DbHelper.drop_databases 27 | DbHelper.create_databases 28 | 29 | load(schema_file) 30 | 31 | ActiveRecord::Base.on_all_shards do 32 | assert table_exists?(:schema_migrations), "Schema Migrations doesn't exist" 33 | assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20110824010216'") 34 | assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20110829215912'") 35 | end 36 | 37 | ActiveRecord::Base.on_all_shards do 38 | assert table_has_column?("tickets", "sharded_column") 39 | end 40 | 41 | ActiveRecord::Base.on_shard(nil) do 42 | assert table_has_column?("accounts", "non_sharded_column") 43 | end 44 | end 45 | end 46 | end 47 | -------------------------------------------------------------------------------- /.rubocop.yml: -------------------------------------------------------------------------------- 1 | require: 2 | - rubocop-minitest 3 | - rubocop-performance 4 | 5 | AllCops: 6 | CacheRootDirectory: tmp 7 | Exclude: 8 | - .git/**/* 9 | - gemfiles/vendor/**/* 10 | - vendor/**/* 11 | TargetRubyVersion: 2.6 12 | 13 | # Configured cops 14 | 15 | Layout/FirstHashElementIndentation: 16 | EnforcedStyle: consistent 17 | 18 | Layout/HashAlignment: 19 | EnforcedHashRocketStyle: 20 | - table 21 | 22 | Naming/FileName: 23 | Exclude: 24 | - 'lib/active_record_shards/connection_switcher-5-1.rb' 25 | - 'lib/active_record_shards/connection_switcher-6-0.rb' 26 | - 'lib/active_record_shards/connection_switcher-6-1.rb' 27 | - 'lib/active_record_shards/connection_switcher-7-0.rb' 28 | 29 | Style/Alias: 30 | EnforcedStyle: prefer_alias_method 31 | 32 | Style/EmptyMethod: 33 | EnforcedStyle: expanded 34 | 35 | Style/FormatString: 36 | EnforcedStyle: percent 37 | 38 | # Disabled cops 39 | 40 | Lint/AssignmentInCondition: 41 | Enabled: false 42 | 43 | Metrics: 44 | Enabled: false 45 | 46 | Style/AsciiComments: 47 | Enabled: false 48 | 49 | Style/ClassVars: 50 | Enabled: false 51 | 52 | Style/ConditionalAssignment: 53 | Enabled: false 54 | 55 | Style/Documentation: 56 | Enabled: false 57 | 58 | Style/DoubleNegation: 59 | Enabled: false 60 | 61 | Style/FormatStringToken: 62 | Enabled: false 63 | 64 | Style/FrozenStringLiteralComment: 65 | Enabled: false 66 | 67 | Style/GuardClause: 68 | Enabled: false 69 | 70 | Style/IfUnlessModifier: 71 | Enabled: false 72 | 73 | Style/NumericLiterals: 74 | Enabled: false 75 | 76 | Style/StringLiterals: 77 | Enabled: false 78 | 79 | Style/SymbolArray: 80 | Enabled: false 81 | 82 | Style/WordArray: 83 | Enabled: false 84 | -------------------------------------------------------------------------------- /test/schemas/ars_test2.sql: -------------------------------------------------------------------------------- 1 | DROP TABLE IF EXISTS `account_people`; 2 | CREATE TABLE `account_people` ( 3 | `account_id` int(11) DEFAULT NULL, 4 | `person_id` int(11) DEFAULT NULL 5 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 6 | 7 | DROP TABLE IF EXISTS `account_things`; 8 | CREATE TABLE `account_things` ( 9 | `id` int(11) NOT NULL AUTO_INCREMENT, 10 | `account_id` int(11) DEFAULT NULL, 11 | `enabled` tinyint(1) DEFAULT '1', 12 | PRIMARY KEY (`id`) 13 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 14 | 15 | DROP TABLE IF EXISTS `accounts`; 16 | CREATE TABLE `accounts` ( 17 | `id` int(11) NOT NULL AUTO_INCREMENT, 18 | `name` varchar(255) DEFAULT NULL, 19 | `created_at` datetime DEFAULT NULL, 20 | `updated_at` datetime DEFAULT NULL, 21 | PRIMARY KEY (`id`) 22 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 23 | 24 | DROP TABLE IF EXISTS `ar_internal_metadata`; 25 | CREATE TABLE `ar_internal_metadata` ( 26 | `key` varchar(255) CHARACTER SET utf8 NOT NULL, 27 | `value` varchar(255) DEFAULT NULL, 28 | `created_at` datetime NOT NULL, 29 | `updated_at` datetime NOT NULL, 30 | PRIMARY KEY (`key`) 31 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 32 | 33 | DROP TABLE IF EXISTS `people`; 34 | CREATE TABLE `people` ( 35 | `id` int(11) NOT NULL AUTO_INCREMENT, 36 | `name` varchar(255) DEFAULT NULL, 37 | `type` varchar(255) DEFAULT NULL, 38 | `created_at` datetime DEFAULT NULL, 39 | `updated_at` datetime DEFAULT NULL, 40 | PRIMARY KEY (`id`) 41 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 42 | 43 | DROP TABLE IF EXISTS `schema_migrations`; 44 | CREATE TABLE `schema_migrations` ( 45 | `version` varchar(255) CHARACTER SET utf8 NOT NULL, 46 | PRIMARY KEY (`version`) 47 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 48 | 49 | INSERT INTO `schema_migrations` (version) VALUES ('1'); 50 | -------------------------------------------------------------------------------- /test/schemas/ars_test2_replica.sql: -------------------------------------------------------------------------------- 1 | DROP TABLE IF EXISTS `account_people`; 2 | CREATE TABLE `account_people` ( 3 | `account_id` int(11) DEFAULT NULL, 4 | `person_id` int(11) DEFAULT NULL 5 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 6 | 7 | DROP TABLE IF EXISTS `account_things`; 8 | CREATE TABLE `account_things` ( 9 | `id` int(11) NOT NULL AUTO_INCREMENT, 10 | `account_id` int(11) DEFAULT NULL, 11 | `enabled` tinyint(1) DEFAULT '1', 12 | PRIMARY KEY (`id`) 13 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 14 | 15 | DROP TABLE IF EXISTS `accounts`; 16 | CREATE TABLE `accounts` ( 17 | `id` int(11) NOT NULL AUTO_INCREMENT, 18 | `name` varchar(255) DEFAULT NULL, 19 | `created_at` datetime DEFAULT NULL, 20 | `updated_at` datetime DEFAULT NULL, 21 | PRIMARY KEY (`id`) 22 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 23 | 24 | DROP TABLE IF EXISTS `ar_internal_metadata`; 25 | CREATE TABLE `ar_internal_metadata` ( 26 | `key` varchar(255) CHARACTER SET utf8 NOT NULL, 27 | `value` varchar(255) DEFAULT NULL, 28 | `created_at` datetime NOT NULL, 29 | `updated_at` datetime NOT NULL, 30 | PRIMARY KEY (`key`) 31 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 32 | 33 | DROP TABLE IF EXISTS `people`; 34 | CREATE TABLE `people` ( 35 | `id` int(11) NOT NULL AUTO_INCREMENT, 36 | `name` varchar(255) DEFAULT NULL, 37 | `type` varchar(255) DEFAULT NULL, 38 | `created_at` datetime DEFAULT NULL, 39 | `updated_at` datetime DEFAULT NULL, 40 | PRIMARY KEY (`id`) 41 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 42 | 43 | DROP TABLE IF EXISTS `schema_migrations`; 44 | CREATE TABLE `schema_migrations` ( 45 | `version` varchar(255) CHARACTER SET utf8 NOT NULL, 46 | PRIMARY KEY (`version`) 47 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 48 | 49 | INSERT INTO `schema_migrations` (version) VALUES ('1'); 50 | -------------------------------------------------------------------------------- /test/schemas/ars_test.sql: -------------------------------------------------------------------------------- 1 | DROP TABLE IF EXISTS `account_people`; 2 | CREATE TABLE `account_people` ( 3 | `account_id` int(11) DEFAULT NULL, 4 | `person_id` int(11) DEFAULT NULL 5 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 6 | 7 | DROP TABLE IF EXISTS `account_things`; 8 | CREATE TABLE `account_things` ( 9 | `id` int(11) NOT NULL AUTO_INCREMENT, 10 | `account_id` int(11) DEFAULT NULL, 11 | `enabled` tinyint(1) DEFAULT '1', 12 | PRIMARY KEY (`id`) 13 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 14 | 15 | DROP TABLE IF EXISTS `accounts`; 16 | CREATE TABLE `accounts` ( 17 | `id` int(11) NOT NULL AUTO_INCREMENT, 18 | `name` varchar(255) DEFAULT NULL, 19 | `created_at` datetime DEFAULT NULL, 20 | `updated_at` datetime DEFAULT NULL, 21 | PRIMARY KEY (`id`) 22 | ) ENGINE=InnoDB AUTO_INCREMENT=1001 DEFAULT CHARSET=utf8mb4; 23 | 24 | DROP TABLE IF EXISTS `ar_internal_metadata`; 25 | CREATE TABLE `ar_internal_metadata` ( 26 | `key` varchar(255) CHARACTER SET utf8 NOT NULL, 27 | `value` varchar(255) DEFAULT NULL, 28 | `created_at` datetime NOT NULL, 29 | `updated_at` datetime NOT NULL, 30 | PRIMARY KEY (`key`) 31 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 32 | 33 | DROP TABLE IF EXISTS `people`; 34 | CREATE TABLE `people` ( 35 | `id` int(11) NOT NULL AUTO_INCREMENT, 36 | `name` varchar(255) DEFAULT NULL, 37 | `type` varchar(255) DEFAULT NULL, 38 | `created_at` datetime DEFAULT NULL, 39 | `updated_at` datetime DEFAULT NULL, 40 | PRIMARY KEY (`id`) 41 | ) ENGINE=InnoDB AUTO_INCREMENT=11 DEFAULT CHARSET=utf8mb4; 42 | 43 | DROP TABLE IF EXISTS `schema_migrations`; 44 | CREATE TABLE `schema_migrations` ( 45 | `version` varchar(255) CHARACTER SET utf8 NOT NULL, 46 | PRIMARY KEY (`version`) 47 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 48 | 49 | INSERT INTO `schema_migrations` (version) VALUES ('1'); 50 | -------------------------------------------------------------------------------- /test/schemas/ars_test_replica.sql: -------------------------------------------------------------------------------- 1 | DROP TABLE IF EXISTS `account_people`; 2 | CREATE TABLE `account_people` ( 3 | `account_id` int(11) DEFAULT NULL, 4 | `person_id` int(11) DEFAULT NULL 5 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 6 | 7 | DROP TABLE IF EXISTS `account_things`; 8 | CREATE TABLE `account_things` ( 9 | `id` int(11) NOT NULL AUTO_INCREMENT, 10 | `account_id` int(11) DEFAULT NULL, 11 | `enabled` tinyint(1) DEFAULT '1', 12 | PRIMARY KEY (`id`) 13 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 14 | 15 | DROP TABLE IF EXISTS `accounts`; 16 | CREATE TABLE `accounts` ( 17 | `id` int(11) NOT NULL AUTO_INCREMENT, 18 | `name` varchar(255) DEFAULT NULL, 19 | `created_at` datetime DEFAULT NULL, 20 | `updated_at` datetime DEFAULT NULL, 21 | PRIMARY KEY (`id`) 22 | ) ENGINE=InnoDB AUTO_INCREMENT=1002 DEFAULT CHARSET=utf8mb4; 23 | 24 | DROP TABLE IF EXISTS `ar_internal_metadata`; 25 | CREATE TABLE `ar_internal_metadata` ( 26 | `key` varchar(255) CHARACTER SET utf8 NOT NULL, 27 | `value` varchar(255) DEFAULT NULL, 28 | `created_at` datetime NOT NULL, 29 | `updated_at` datetime NOT NULL, 30 | PRIMARY KEY (`key`) 31 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 32 | 33 | DROP TABLE IF EXISTS `people`; 34 | CREATE TABLE `people` ( 35 | `id` int(11) NOT NULL AUTO_INCREMENT, 36 | `name` varchar(255) DEFAULT NULL, 37 | `type` varchar(255) DEFAULT NULL, 38 | `created_at` datetime DEFAULT NULL, 39 | `updated_at` datetime DEFAULT NULL, 40 | PRIMARY KEY (`id`) 41 | ) ENGINE=InnoDB AUTO_INCREMENT=21 DEFAULT CHARSET=utf8mb4; 42 | 43 | DROP TABLE IF EXISTS `schema_migrations`; 44 | CREATE TABLE `schema_migrations` ( 45 | `version` varchar(255) CHARACTER SET utf8 NOT NULL, 46 | PRIMARY KEY (`version`) 47 | ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; 48 | 49 | INSERT INTO `schema_migrations` (version) VALUES ('1'); 50 | -------------------------------------------------------------------------------- /test/database.yml: -------------------------------------------------------------------------------- 1 | <% mysql = URI(ENV['MYSQL_URL'] || 'mysql://root@127.0.0.1:3306') %> 2 | 3 | mysql: &MYSQL 4 | encoding: utf8 5 | adapter: mysql2 6 | username: <%= mysql.user %> 7 | host: <%= mysql.host %> 8 | port: <%= mysql.port %> 9 | password: <%= mysql.password %> 10 | ssl_mode: :disabled 11 | reaping_frequency: 0 # Prevents ActiveRecord from spawning reaping threads. 12 | 13 | # We connect to the unsharded primary database on a different port, via a proxy, 14 | # so we can make the connection unavailable when testing on_replica_by_default 15 | # behavior. 16 | test: 17 | <<: *MYSQL 18 | database: ars_test 19 | shard_names: [0, 1] 20 | host: 127.0.0.1 21 | port: 13306 22 | 23 | test_replica: 24 | <<: *MYSQL 25 | database: ars_test_replica 26 | 27 | # We connect to this sharded primary database on a different port, via a proxy, 28 | # so we can make the connection unavailable when testing on_replica_by_default 29 | # behavior. 30 | test_shard_0: 31 | <<: *MYSQL 32 | database: ars_test_shard0 33 | host: 127.0.0.1 34 | port: 13307 35 | 36 | test_shard_0_replica: 37 | <<: *MYSQL 38 | database: ars_test_shard0_replica 39 | 40 | # We connect to this sharded primary database on a different port, via a proxy, 41 | # so we can make the connection unavailable when testing on_replica_by_default 42 | # behavior. 43 | test_shard_1: 44 | <<: *MYSQL 45 | database: ars_test_shard1 46 | host: 127.0.0.1 47 | port: 13308 48 | 49 | test_shard_1_replica: 50 | <<: *MYSQL 51 | database: ars_test_shard1_replica 52 | 53 | test2: 54 | <<: *MYSQL 55 | database: ars_test2 56 | 57 | test2_replica: 58 | <<: *MYSQL 59 | database: ars_test2_replica 60 | 61 | test3: 62 | <<: *MYSQL 63 | database: ars_test3 64 | shard_names: [0] 65 | 66 | test3_shard_0: 67 | <<: *MYSQL 68 | database: ars_test3_shard0 69 | 70 | -------------------------------------------------------------------------------- /lib/active_record_shards/configuration_parser.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require 'active_support/core_ext' 4 | 5 | module ActiveRecordShards 6 | module ConfigurationParser 7 | module_function 8 | 9 | def explode(conf) 10 | conf = conf.to_h.deep_dup 11 | 12 | conf.to_a.each do |env_name, env_config| 13 | next unless shards = env_config.delete('shards') 14 | 15 | unless shards.keys.all? { |shard_name| shard_name.is_a?(Integer) } 16 | raise "All shard names must be integers: #{shards.keys.inspect}." 17 | end 18 | 19 | env_config['shard_names'] = shards.keys 20 | shards.each do |shard_name, shard_conf| 21 | expand_child!(env_config, shard_conf) 22 | conf["#{env_name}_shard_#{shard_name}"] = shard_conf 23 | end 24 | end 25 | 26 | conf.to_a.each do |env_name, env_config| 27 | if replica_conf = env_config.delete('replica') 28 | expand_child!(env_config, replica_conf) 29 | conf["#{env_name}_replica"] = replica_conf 30 | end 31 | end 32 | 33 | conf 34 | end 35 | 36 | def expand_child!(parent, child) 37 | parent.each do |key, value| 38 | unless ['replica', 'shards'].include?(key) || value.is_a?(Hash) 39 | child[key] ||= value 40 | end 41 | end 42 | end 43 | 44 | def configurations_with_shard_explosion=(conf) 45 | self.configurations_without_shard_explosion = explode(conf) 46 | end 47 | 48 | def self.extended(base) 49 | base.singleton_class.send(:alias_method, :configurations_without_shard_explosion=, :configurations=) 50 | base.singleton_class.send(:alias_method, :configurations=, :configurations_with_shard_explosion=) 51 | base.singleton_class.send(:public, :configurations=) 52 | 53 | base.configurations = base.configurations if base.configurations.present? 54 | end 55 | end 56 | end 57 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: push 4 | 5 | jobs: 6 | tests: 7 | runs-on: ubuntu-latest 8 | 9 | services: 10 | mysql: 11 | image: mysql:5.7 12 | ports: 13 | - 3306:3306 14 | env: 15 | MYSQL_ALLOW_EMPTY_PASSWORD: "yes" 16 | options: >- 17 | --health-cmd "mysql -uroot -e 'show databases'" 18 | --health-interval 2s 19 | --health-timeout 1s 20 | --health-retries 10 21 | 22 | strategy: 23 | fail-fast: false 24 | matrix: 25 | ruby-version: 26 | - "2.7" 27 | - "3.0" 28 | - "3.1" 29 | - "3.2" 30 | - "3.3" 31 | gemfile: 32 | - rails6.1 33 | - rails7.0 34 | include: 35 | - {ruby-version: "2.6", gemfile: rails5.1} 36 | - {ruby-version: "2.6", gemfile: rails5.2} 37 | - {ruby-version: "2.6", gemfile: rails6.0} 38 | - {ruby-version: "2.6", gemfile: rails6.1} 39 | - {ruby-version: "2.7", gemfile: rails6.0} 40 | env: 41 | BUNDLE_GEMFILE: gemfiles/${{ matrix.gemfile }}.gemfile 42 | steps: 43 | - uses: actions/checkout@v4 44 | - name: Set up Ruby 45 | uses: ruby/setup-ruby@v1 46 | with: 47 | ruby-version: ${{ matrix.ruby-version }} 48 | bundler-cache: true 49 | - run: bundle exec rake test 50 | 51 | tests_successful: 52 | name: Tests passing? 53 | needs: tests 54 | if: always() 55 | runs-on: ubuntu-latest 56 | steps: 57 | - run: | 58 | if ${{ needs.tests.result == 'success' }} 59 | then 60 | echo "All tests passed" 61 | else 62 | echo "Some tests failed" 63 | false 64 | fi 65 | 66 | lint: 67 | runs-on: ubuntu-latest 68 | steps: 69 | - uses: actions/checkout@v4 70 | - name: Set up Ruby 71 | uses: ruby/setup-ruby@v1 72 | with: 73 | ruby-version: "3.0" 74 | bundler-cache: true 75 | - run: bundle exec rake rubocop 76 | -------------------------------------------------------------------------------- /lib/active_record_shards/model.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module ActiveRecordShards 4 | module Model 5 | def not_sharded 6 | if self != ActiveRecord::Base && self != base_class 7 | raise "You should only call not_sharded on direct descendants of ActiveRecord::Base" 8 | end 9 | 10 | self.sharded = false 11 | end 12 | 13 | def is_sharded? # rubocop:disable Naming/PredicateName 14 | return @_ars_model_is_sharded unless @_ars_model_is_sharded.nil? 15 | 16 | @_ars_model_is_sharded = 17 | if self == ActiveRecord::Base 18 | sharded != false && supports_sharding? 19 | elsif self == base_class 20 | if sharded.nil? 21 | ActiveRecord::Base.is_sharded? 22 | else 23 | sharded != false 24 | end 25 | else 26 | base_class.is_sharded? 27 | end 28 | end 29 | 30 | def on_replica_by_default? 31 | if self == ActiveRecord::Base 32 | false 33 | else 34 | base = base_class 35 | if base.instance_variable_defined?(:@on_replica_by_default) 36 | base.instance_variable_get(:@on_replica_by_default) 37 | end 38 | end 39 | end 40 | 41 | def on_replica_by_default=(value) 42 | if self == ActiveRecord::Base 43 | raise ArgumentError, "Cannot set on_replica_by_default on ActiveRecord::Base" 44 | else 45 | base_class.instance_variable_set(:@on_replica_by_default, value) 46 | end 47 | end 48 | 49 | module InstanceMethods 50 | def initialize_shard_and_replica 51 | @from_replica = !!self.class.current_shard_selection.options[:replica] 52 | @from_shard = self.class.current_shard_selection.options[:shard] 53 | end 54 | 55 | def from_replica? 56 | @from_replica 57 | end 58 | 59 | def from_shard 60 | @from_shard 61 | end 62 | end 63 | 64 | def self.extended(base) 65 | base.send(:include, InstanceMethods) 66 | base.after_initialize :initialize_shard_and_replica 67 | end 68 | 69 | private 70 | 71 | attr_accessor :sharded 72 | end 73 | end 74 | -------------------------------------------------------------------------------- /lib/active_record_shards/shard_selection.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module ActiveRecordShards 4 | class ShardSelection 5 | NO_SHARD = :_no_shard 6 | cattr_accessor :ars_default_shard 7 | 8 | def initialize 9 | @on_replica = false 10 | @shard = nil 11 | end 12 | 13 | def shard 14 | if @shard.nil? || @shard == NO_SHARD 15 | nil 16 | else 17 | @shard || self.class.ars_default_shard 18 | end 19 | end 20 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 21 | when '6.1', '7.0' 22 | PRIMARY = "ActiveRecord::Base" 23 | else 24 | PRIMARY = "primary" 25 | end 26 | def resolve_connection_name(sharded:, configurations:) 27 | resolved_shard = sharded ? shard : nil 28 | env = ActiveRecordShards.app_env 29 | 30 | @connection_names ||= {} 31 | @connection_names[env] ||= {} 32 | @connection_names[env][resolved_shard] ||= {} 33 | @connection_names[env][resolved_shard][@on_replica] ||= begin 34 | name = env.dup 35 | name << "_shard_#{resolved_shard}" if resolved_shard 36 | replica_config = begin 37 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 38 | when '7.0' 39 | configurations.configs_for(env_name: "#{name}_replica", include_hidden: true).any? 40 | when '6.1' 41 | configurations.configs_for(env_name: "#{name}_replica", include_replicas: true).any? 42 | else 43 | configurations["#{name}_replica"] 44 | end 45 | end 46 | if @on_replica && replica_config 47 | "#{name}_replica" 48 | else 49 | # ActiveRecord always names its default connection pool 'primary' 50 | # while everything else is named by the configuration name 51 | resolved_shard ? name : PRIMARY 52 | end 53 | end 54 | end 55 | 56 | def shard=(new_shard) 57 | @shard = (new_shard || NO_SHARD) 58 | end 59 | 60 | def on_replica? 61 | @on_replica 62 | end 63 | 64 | def on_replica=(new_replica) 65 | @on_replica = (new_replica == true) 66 | end 67 | 68 | def options 69 | { shard: @shard, replica: @on_replica } 70 | end 71 | end 72 | end 73 | -------------------------------------------------------------------------------- /test/active_record_shards_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative 'helper' 4 | 5 | describe 'ActiveRecordShards' do 6 | describe '.app_env' do 7 | before do 8 | if defined?(Rails) || ENV['RAILS_ENV'] || defined?(APP_ENV) || ENV['APP_ENV'] 9 | raise "Tests in #{__FILE__} will overwrite environment constants, please update them to avoid conflicts" 10 | end 11 | 12 | Object.send(:remove_const, 'RAILS_ENV') 13 | ActiveRecordShards.instance_eval { @app_env = nil } 14 | end 15 | 16 | after do 17 | Object.const_set('RAILS_ENV', 'test') 18 | ActiveRecordShards.instance_eval { @app_env = nil } 19 | end 20 | 21 | describe 'Rails.env' do 22 | before do 23 | class Rails 24 | def self.env 25 | 'environment from Rails.env' 26 | end 27 | end 28 | end 29 | 30 | after { Object.send(:remove_const, 'Rails') } 31 | 32 | it 'looks for Rails.env' do 33 | assert_equal 'environment from Rails.env', ActiveRecordShards.app_env 34 | end 35 | end 36 | 37 | describe 'RAILS_ENV' do 38 | before { Object.const_set('RAILS_ENV', 'environment from RAILS_ENV') } 39 | after { Object.send(:remove_const, 'RAILS_ENV') } 40 | 41 | it 'looks for RAILS_ENV' do 42 | assert_equal 'environment from RAILS_ENV', ActiveRecordShards.app_env 43 | end 44 | end 45 | 46 | describe "ENV['RAILS_ENV']" do 47 | before { ENV['RAILS_ENV'] = "environment from ENV['RAILS_ENV']" } 48 | after { ENV.delete('RAILS_ENV') } 49 | 50 | it 'looks for RAILS_ENV' do 51 | assert_equal "environment from ENV['RAILS_ENV']", ActiveRecordShards.app_env 52 | end 53 | end 54 | 55 | describe 'APP_ENV' do 56 | before { Object.const_set('APP_ENV', 'environment from APP_ENV') } 57 | after { Object.send(:remove_const, 'APP_ENV') } 58 | 59 | it 'looks for APP_ENV' do 60 | assert_equal 'environment from APP_ENV', ActiveRecordShards.app_env 61 | end 62 | end 63 | 64 | describe "ENV['APP_ENV']" do 65 | before { ENV['APP_ENV'] = "environment from ENV['APP_ENV']" } 66 | after { ENV.delete('APP_ENV') } 67 | 68 | it 'looks for APP_ENV' do 69 | assert_equal "environment from ENV['APP_ENV']", ActiveRecordShards.app_env 70 | end 71 | end 72 | end 73 | end 74 | -------------------------------------------------------------------------------- /test/support/db_helper.rb: -------------------------------------------------------------------------------- 1 | module DbHelper 2 | class << self 3 | def client 4 | @client ||= begin 5 | config = URI(ENV["MYSQL_URL"] || "mysql://root@127.0.0.1:3306") 6 | Mysql2::Client.new( 7 | username: config.user, 8 | password: config.password, 9 | host: config.host, 10 | port: config.port 11 | ) 12 | end 13 | end 14 | 15 | def each_database(&_block) 16 | Dir.glob("test/schemas/*.sql").each do |schema_path| 17 | database_name = File.basename(schema_path, ".sql").to_s 18 | yield(database_name, schema_path) 19 | end 20 | end 21 | 22 | def mysql(commands) 23 | commands.split(/\s*;\s*/).reject(&:empty?).each do |command| 24 | client.query(command) 25 | end 26 | end 27 | 28 | def drop_databases 29 | each_database do |database_name, _schema_path| 30 | DbHelper.mysql("DROP DATABASE IF EXISTS #{database_name}") 31 | end 32 | end 33 | 34 | def create_databases 35 | each_database do |database_name, _schema_path| 36 | DbHelper.mysql("CREATE DATABASE IF NOT EXISTS #{database_name}") 37 | end 38 | end 39 | 40 | def load_database_schemas 41 | each_database do |database_name, schema_path| 42 | DbHelper.mysql("USE #{database_name}") 43 | DbHelper.mysql(File.read(schema_path)) 44 | end 45 | end 46 | 47 | # Load the database configuration into ActiveRecord 48 | def load_database_configuration(path_or_io = 'test/database.yml') 49 | erb_config = path_or_io.is_a?(String) ? IO.read(path_or_io) : path_or_io.read 50 | yaml_config = ERB.new(erb_config).result 51 | ActiveRecord::Base.configurations = begin 52 | YAML.load(yaml_config, aliases: true) # rubocop:disable Security/YAMLLoad 53 | rescue ArgumentError 54 | YAML.load(yaml_config) # rubocop:disable Security/YAMLLoad 55 | end 56 | end 57 | end 58 | 59 | # Create all databases and then tear them down after test 60 | def with_fresh_databases 61 | before do 62 | DbHelper.drop_databases 63 | DbHelper.create_databases 64 | DbHelper.load_database_schemas 65 | 66 | clear_global_connection_handler_state 67 | DbHelper.load_database_configuration 68 | end 69 | 70 | after do 71 | DbHelper.drop_databases 72 | end 73 | end 74 | end 75 | -------------------------------------------------------------------------------- /test/check_performance.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require "bundler/setup" 4 | 5 | $LOAD_PATH.unshift(File.join(__dir__, '..', 'lib')) 6 | $LOAD_PATH.unshift(__dir__) 7 | 8 | require "active_record" 9 | require "active_record_shards" 10 | require "mysql2" 11 | require "benchmark/ips" 12 | require "support/db_helper" 13 | 14 | RAILS_ENV = "test" 15 | NUMBER_OF_SHARDS = 300 16 | ActiveRecord::Base.logger = Logger.new("/dev/null") 17 | 18 | class Account < ActiveRecord::Base 19 | not_sharded 20 | 21 | has_many :tickets 22 | end 23 | 24 | class Ticket < ActiveRecord::Base 25 | belongs_to :account 26 | end 27 | 28 | erb_config = <<-ERB 29 | <% mysql = URI(ENV['MYSQL_URL'] || 'mysql://root@127.0.0.1:3306') %> 30 | 31 | mysql: &MYSQL 32 | encoding: utf8 33 | adapter: mysql2 34 | username: <%= mysql.user %> 35 | host: <%= mysql.host %> 36 | port: <%= mysql.port %> 37 | password: <%= mysql.password %> 38 | ssl_mode: :disabled 39 | 40 | test: 41 | <<: *MYSQL 42 | database: ars_test 43 | shard_names: <%= (0...NUMBER_OF_SHARDS).to_a %> 44 | test_replica: 45 | <<: *MYSQL 46 | database: ars_test_replica 47 | <% NUMBER_OF_SHARDS.times do |shard_id| %> 48 | test_shard_<%= shard_id %>: 49 | <<: *MYSQL 50 | database: ars_test_shard<%= shard_id %> 51 | test_shard_<%= shard_id %>_replica: 52 | <<: *MYSQL 53 | database: ars_test_shard<%= shard_id %>_replica 54 | <% end %> 55 | ERB 56 | 57 | config_io = StringIO.new(erb_config) 58 | DbHelper.drop_databases 59 | DbHelper.create_databases 60 | DbHelper.load_database_schemas 61 | DbHelper.load_database_configuration(config_io) 62 | 63 | ActiveRecord::Base.establish_connection(:test) 64 | 65 | def unsharded_primary 66 | ActiveRecord::Base.on_shard(nil) do 67 | ActiveRecord::Base.on_primary { Account.count } 68 | end 69 | end 70 | 71 | def unsharded_replica 72 | ActiveRecord::Base.on_shard(nil) do 73 | ActiveRecord::Base.on_replica { Account.count } 74 | end 75 | end 76 | 77 | def sharded_primary 78 | ActiveRecord::Base.on_shard(1) do 79 | ActiveRecord::Base.on_primary { Ticket.count } 80 | end 81 | end 82 | 83 | def sharded_replica 84 | ActiveRecord::Base.on_shard(1) do 85 | ActiveRecord::Base.on_replica { Ticket.count } 86 | end 87 | end 88 | 89 | def switch_around 90 | unsharded_primary 91 | unsharded_replica 92 | sharded_primary 93 | sharded_replica 94 | end 95 | 96 | Benchmark.ips do |x| 97 | x.report("#{ActiveRecord::VERSION::STRING} DB switching") { switch_around } 98 | end 99 | 100 | # Results using Ruby 2.7.5 101 | # 102 | # 5.1.7 DB switching 211.790 (± 3.8%) i/s 103 | # 5.2.5 DB switching 213.797 (± 5.1%) i/s 104 | # 6.0.4 DB switching 216.607 (± 4.2%) i/s 105 | 106 | DbHelper.drop_databases 107 | -------------------------------------------------------------------------------- /test/active_record_shards/configuration_parser_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative '../helper' 4 | 5 | describe "ActiveRecordShards::ConfigurationParser.explode" do 6 | let(:yaml) do 7 | <<~YAML 8 | test: 9 | adapter: mysql 10 | encoding: utf8 11 | database: ars_test 12 | port: 123 13 | username: root 14 | password: 15 | host: main_host 16 | replica: 17 | host: main_replica_host 18 | shards: 19 | 500: 20 | database: ars_test_shard_500 21 | host: shard_500_host 22 | replica: 23 | host: shard_500_replica_host 24 | 501: 25 | database: ars_test_shard_501 26 | host: shard_501_host 27 | replica: 28 | database: ars_test_shard_501_replica 29 | YAML 30 | end 31 | 32 | before do 33 | @exploded_conf = ActiveRecordShards::ConfigurationParser.explode(YAML.safe_load(yaml)) 34 | end 35 | 36 | it "expands configuration for the main primary" do 37 | config = @exploded_conf["test"] 38 | 39 | shared_assertions(config) 40 | assert_equal "main_host", config["host"] 41 | assert_equal "ars_test", config["database"] 42 | end 43 | 44 | it "expands configuration for the main replica" do 45 | config = @exploded_conf["test_replica"] 46 | 47 | shared_assertions(config) 48 | assert_equal "main_replica_host", config["host"] 49 | assert_equal "ars_test", config["database"] 50 | end 51 | 52 | it "expands configuration for shard 500's primary" do 53 | config = @exploded_conf["test_shard_500"] 54 | 55 | shared_assertions(config) 56 | assert_equal "shard_500_host", config["host"] 57 | assert_equal "ars_test_shard_500", config["database"] 58 | end 59 | 60 | it "expands configuration for shard 500's replica" do 61 | config = @exploded_conf["test_shard_500_replica"] 62 | 63 | shared_assertions(config) 64 | assert_equal "shard_500_replica_host", config["host"] 65 | assert_equal "ars_test_shard_500", config["database"] 66 | end 67 | 68 | it "expands configuration for shard 501's primary" do 69 | config = @exploded_conf["test_shard_501"] 70 | 71 | shared_assertions(config) 72 | assert_equal "shard_501_host", config["host"] 73 | assert_equal "ars_test_shard_501", config["database"] 74 | end 75 | 76 | it "expands configuration for shard 501's replica" do 77 | config = @exploded_conf["test_shard_501_replica"] 78 | 79 | shared_assertions(config) 80 | assert_equal "shard_501_host", config["host"] 81 | assert_equal "ars_test_shard_501_replica", config["database"] 82 | end 83 | 84 | def shared_assertions(config) 85 | assert_equal "mysql", config["adapter"] 86 | assert_equal "utf8", config["encoding"] 87 | assert_equal 123, config["port"] 88 | assert_equal "root", config["username"] 89 | assert_nil config["password"] 90 | assert_equal [500, 501], config["shard_names"] 91 | end 92 | end 93 | -------------------------------------------------------------------------------- /test/support/tcp_proxy.rb: -------------------------------------------------------------------------------- 1 | # TCPProxy binds `local_port` and forwards requests to `remote_host`:`remote_port` 2 | # 3 | # proxy = TCPProxy.start( 4 | # remote_host: '127.0.0.1', 5 | # remote_port: '3306', 6 | # local_port: '13306' 7 | # ) 8 | # 9 | # You can temporarily disable and re-enable the proxying: 10 | # 11 | # proxy.pause do 12 | # do_work_that_cannot_call_proxied_service 13 | # end 14 | # 15 | require 'socket' 16 | 17 | class TCPProxy 18 | THREAD_CHECK_INTERVAL = 0.001 19 | 20 | def self.start(remote_host:, remote_port:, local_port:) 21 | new( 22 | remote_host: remote_host, 23 | remote_port: remote_port, 24 | local_port: local_port 25 | ).tap(&:start) 26 | end 27 | 28 | def initialize(remote_host:, remote_port:, local_port:) 29 | @remote_host = remote_host 30 | @remote_port = remote_port 31 | @local_port = local_port 32 | 33 | @disabled = false 34 | end 35 | 36 | def start 37 | proxy_server = TCPServer.new('0.0.0.0', local_port) 38 | 39 | @thr = Thread.new do 40 | loop do 41 | requesting_socket = proxy_server.accept 42 | 43 | Thread.new do 44 | responding_socket = TCPSocket.new(remote_host, remote_port) 45 | 46 | requests = Thread.new { forward(requesting_socket, responding_socket, pause_behavior: :return) } 47 | requests.abort_on_exception = true 48 | 49 | responses = Thread.new { forward(responding_socket, requesting_socket) } 50 | responses.abort_on_exception = true 51 | 52 | # Either thread can be the first to finish - requests if the mysql2 client 53 | # closes the connection; responses if the MySQL server closes - so we 54 | # cannot do the more common `requests.join and responses.join`. 55 | sleep THREAD_CHECK_INTERVAL while requests.alive? && responses.alive? 56 | requests.kill 57 | responses.kill 58 | sleep THREAD_CHECK_INTERVAL until requests.stop? && responses.stop? 59 | ensure 60 | requesting_socket&.close 61 | responding_socket&.close 62 | end 63 | end 64 | ensure 65 | proxy_server.close 66 | end 67 | end 68 | 69 | def pause(&_block) 70 | # Give requests already sent to the socket a chance to be picked up before pausing. 71 | sleep 0.001 72 | @disabled = true 73 | yield 74 | ensure 75 | @disabled = false 76 | end 77 | 78 | private 79 | 80 | attr_reader :remote_host, :remote_port, :local_port 81 | 82 | def forward(src, dst, pause_behavior: :ignore) 83 | zero_counter = 0 84 | loop do 85 | data = src.recv(1024) 86 | 87 | if enabled? || pause_behavior == :ignore 88 | if data.nil? || data.empty? 89 | zero_counter += 1 90 | return if zero_counter >= 5 91 | else 92 | dst.send(data, 0) 93 | end 94 | elsif disabled? && pause_behavior == :return 95 | clean_data = data.gsub(/[^\w. ]/, '').strip 96 | 97 | warn "TCPProxy received a request while paused: `#{clean_data}`" 98 | return 99 | else 100 | raise "Invalid state" 101 | end 102 | end 103 | end 104 | 105 | def disabled? 106 | !enabled? 107 | end 108 | 109 | def enabled? 110 | !@disabled 111 | end 112 | end 113 | -------------------------------------------------------------------------------- /test/tasks_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative 'helper' 4 | 5 | # ActiveRecordShards overrides some of the ActiveRecord tasks, so 6 | # ActiveRecord needs to be loaded first. 7 | Rake::Application.new.rake_require("active_record/railties/databases") 8 | require 'active_record_shards/tasks' 9 | task :environment do 10 | # Only required as a dependency 11 | end 12 | 13 | describe "Database rake tasks" do 14 | include RakeSpecHelpers 15 | 16 | def capture_stderr 17 | $stderr = StringIO.new 18 | yield 19 | $stderr.string 20 | ensure 21 | $stderr = STDERR 22 | end 23 | 24 | let(:config) { DbHelper.load_database_configuration('test/support/database_tasks.yml') } 25 | let(:primary_name) { config['test']['database'] } 26 | let(:replica_name) { config['test']['replica']['database'] } 27 | let(:shard_names) { config['test']['shards'].values.map { |v| v['database'] } } 28 | let(:database_names) { shard_names + [primary_name, replica_name] } 29 | 30 | before do 31 | clear_global_connection_handler_state 32 | 33 | ActiveRecord::Tasks::DatabaseTasks.database_configuration = config 34 | ActiveRecord::Tasks::DatabaseTasks.env = RAILS_ENV 35 | ActiveRecord::Tasks::DatabaseTasks.migrations_paths = '/app/migrations' 36 | end 37 | 38 | after do 39 | DbHelper.load_database_configuration 40 | 41 | %w[ 42 | ars_tasks_test 43 | ars_tasks_test_replica 44 | ars_tasks_test_shard_a 45 | ars_tasks_test_shard_b 46 | ars_tasks_adapter_test 47 | ars_tasks_adapter_test_replica 48 | ars_tasks_adapter_test_shard_a 49 | ars_tasks_adapter_test_shard_b 50 | ].each do |database_name| 51 | DbHelper.mysql("DROP DATABASE IF EXISTS #{database_name}") 52 | end 53 | end 54 | 55 | describe "db:create" do 56 | it "creates the database and all shards" do 57 | rake('db:create') 58 | databases = show_databases(config) 59 | 60 | assert_includes databases, primary_name 61 | refute_includes databases, replica_name 62 | shard_names.each do |name| 63 | assert_includes databases, name 64 | end 65 | end 66 | end 67 | 68 | describe "db:drop" do 69 | it "drops the database and all shards" do 70 | rake('db:create') 71 | rake('db:drop') 72 | databases = show_databases(config) 73 | 74 | refute_includes databases, primary_name 75 | shard_names.each do |name| 76 | refute_includes databases, name 77 | end 78 | end 79 | 80 | it "does not fail when db is missing" do 81 | rake('db:create') 82 | rake('db:drop') 83 | refute_includes(show_databases(config), primary_name) 84 | end 85 | 86 | it "fails loudly when unknown error occurs" do 87 | out = ActiveRecordShards::Tasks.stub(:root_connection, -> { raise ArgumentError }) do 88 | capture_stderr { rake('db:drop') } 89 | end 90 | assert_includes(out, "Couldn't drop ") 91 | assert_includes(out, "test/helper.rb") 92 | end 93 | end 94 | 95 | describe "abort_if_pending_migrations" do 96 | before do 97 | rake('db:create') 98 | end 99 | 100 | it "passes when there is no pending migrations" do 101 | migrator = Struct.new(:pending_migrations).new([]) 102 | 103 | out = ActiveRecord::Migrator.stub(:new, migrator) do 104 | capture_stderr { rake('db:abort_if_pending_migrations') } 105 | end 106 | assert_empty out 107 | end 108 | 109 | it "fails when migrations are pending" do 110 | migration = ActiveRecord::Migration.new('Fake', 1) 111 | migrator = Struct.new(:pending_migrations).new([migration]) 112 | 113 | out = ActiveRecord::Migrator.stub(:new, migrator) do 114 | capture_stderr do 115 | rake('db:abort_if_pending_migrations') 116 | rescue SystemExit 117 | "" 118 | end 119 | end 120 | assert_match(/You have \d+ pending migrations:/, out) 121 | end 122 | end 123 | end 124 | -------------------------------------------------------------------------------- /lib/active_record_shards/migration.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module ActiveRecord 4 | class Migrator 5 | def self.shards_migration_context 6 | if ActiveRecord::VERSION::MAJOR >= 6 7 | ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths, ActiveRecord::SchemaMigration) 8 | elsif ActiveRecord::VERSION::STRING >= '5.2.0' 9 | ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths) 10 | else 11 | self 12 | end 13 | end 14 | 15 | def initialize_with_sharding(*args) 16 | initialize_without_sharding(*args) 17 | 18 | # Rails creates the internal tables on the unsharded DB. We make them 19 | # manually on the sharded DBs. 20 | ActiveRecord::Base.on_all_shards do 21 | ActiveRecord::SchemaMigration.create_table 22 | ActiveRecord::InternalMetadata.create_table 23 | end 24 | end 25 | ruby2_keywords(:initialize_with_sharding) if respond_to?(:ruby2_keywords, true) 26 | alias_method :initialize_without_sharding, :initialize 27 | alias_method :initialize, :initialize_with_sharding 28 | 29 | def run_with_sharding 30 | ActiveRecord::Base.on_shard(nil) { run_without_sharding } 31 | ActiveRecord::Base.on_all_shards { run_without_sharding } 32 | end 33 | alias_method :run_without_sharding, :run 34 | alias_method :run, :run_with_sharding 35 | 36 | def migrate_with_sharding 37 | ActiveRecord::Base.on_shard(nil) { migrate_without_sharding } 38 | ActiveRecord::Base.on_all_shards { migrate_without_sharding } 39 | end 40 | alias_method :migrate_without_sharding, :migrate 41 | alias_method :migrate, :migrate_with_sharding 42 | 43 | # don't allow Migrator class to cache versions 44 | undef migrated 45 | def migrated 46 | self.class.shards_migration_context.get_all_versions 47 | end 48 | 49 | # list of pending migrations is any migrations that haven't run on all shards. 50 | undef pending_migrations 51 | def pending_migrations 52 | pending, _missing = self.class.shard_status(migrations.map(&:version)) 53 | pending = pending.values.flatten 54 | migrations.select { |m| pending.include?(m.version) } 55 | end 56 | 57 | # public 58 | # list of pending and missing versions per shard 59 | # [{1 => [1234567]}, {1 => [2345678]}] 60 | def self.shard_status(versions) 61 | pending = {} 62 | missing = {} 63 | 64 | collect = lambda do |shard| 65 | migrated = shards_migration_context.get_all_versions 66 | 67 | p = versions - migrated 68 | pending[shard] = p if p.any? 69 | 70 | m = migrated - versions 71 | missing[shard] = m if m.any? 72 | end 73 | 74 | ActiveRecord::Base.on_shard(nil) { collect.call(nil) } 75 | ActiveRecord::Base.on_all_shards { |shard| collect.call(shard) } 76 | 77 | [pending, missing] 78 | end 79 | end 80 | end 81 | 82 | module ActiveRecordShards 83 | module MigrationClassExtension 84 | attr_accessor :migration_shard 85 | 86 | def shard(arg = nil) 87 | self.migration_shard = arg 88 | end 89 | end 90 | 91 | module ActualMigrationExtension 92 | def migrate_with_forced_shard(direction) 93 | if migration_shard.blank? 94 | raise "#{name}: Can't run migrations without a shard spec: this may be :all, :none, 95 | or a specific shard (for data-fixups). please call shard(arg) in your migration." 96 | end 97 | 98 | shard = ActiveRecord::Base.current_shard_selection.shard 99 | 100 | if shard.nil? 101 | return if migration_shard != :none 102 | else 103 | return if migration_shard == :none 104 | return if migration_shard != :all && migration_shard.to_s != shard.to_s 105 | end 106 | 107 | migrate_without_forced_shard(direction) 108 | end 109 | 110 | def migration_shard 111 | self.class.migration_shard 112 | end 113 | end 114 | end 115 | 116 | ActiveRecord::Migration.class_eval do 117 | extend ActiveRecordShards::MigrationClassExtension 118 | include ActiveRecordShards::ActualMigrationExtension 119 | 120 | alias_method :migrate_without_forced_shard, :migrate 121 | alias_method :migrate, :migrate_with_forced_shard 122 | end 123 | 124 | ActiveRecord::MigrationProxy.delegate :migration_shard, to: :migration 125 | -------------------------------------------------------------------------------- /test/migrator_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative 'helper' 4 | 5 | describe ActiveRecord::Migrator do 6 | with_fresh_databases 7 | 8 | before { ActiveRecord::Base.establish_connection(RAILS_ENV.to_sym) } 9 | 10 | describe "when DB is empty" do 11 | extend RailsEnvSwitch 12 | 13 | switch_app_env('test3') 14 | 15 | it "makes meta tables" do 16 | ActiveRecord::Base.on_shard(nil) do 17 | refute table_exists?(:unsharded_table) 18 | refute ActiveRecord::SchemaMigration.table_exists? 19 | end 20 | 21 | ActiveRecord::Base.on_all_shards do 22 | refute table_exists?(:sharded_table) 23 | refute ActiveRecord::SchemaMigration.table_exists? 24 | end 25 | 26 | migrator(:up, 'separate_migrations').migrate 27 | 28 | ActiveRecord::Base.on_shard(nil) do 29 | assert table_exists?(:unsharded_table) 30 | assert ActiveRecord::SchemaMigration.table_exists? 31 | assert ActiveRecord::InternalMetadata.table_exists? 32 | assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20190121112233'") 33 | assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20190121112234'") 34 | end 35 | 36 | ActiveRecord::Base.on_all_shards do 37 | assert table_exists?(:sharded_table) 38 | assert ActiveRecord::SchemaMigration.table_exists? 39 | assert ActiveRecord::InternalMetadata.table_exists? 40 | assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20190121112233'") 41 | assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20190121112234'") 42 | end 43 | end 44 | end 45 | 46 | it "migrates" do 47 | refute ActiveRecord::Base.current_shard_id 48 | 49 | migrator.migrate 50 | 51 | ActiveRecord::Base.on_all_shards do 52 | assert table_exists?(:schema_migrations), "Schema Migrations doesn't exist" 53 | assert table_exists?(:tickets) 54 | refute table_exists?(:accounts) 55 | assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20110824010216'") 56 | assert ActiveRecord::Base.connection.select_value("select version from schema_migrations where version = '20110829215912'") 57 | end 58 | 59 | ActiveRecord::Base.on_all_shards do 60 | assert table_has_column?("tickets", "sharded_column") 61 | end 62 | 63 | ActiveRecord::Base.on_shard(nil) do 64 | assert table_has_column?("accounts", "non_sharded_column") 65 | end 66 | 67 | # now test down/ up 68 | migrator(:down, 'migrations', 20110824010216).run 69 | ActiveRecord::Base.on_all_shards do 70 | assert !table_has_column?("tickets", "sharded_column") 71 | end 72 | 73 | migrator(:down, 'migrations', 20110829215912).run 74 | ActiveRecord::Base.on_shard(nil) do 75 | assert !table_has_column?("accounts", "non_sharded_column") 76 | end 77 | 78 | migrator(:up, 'migrations', 20110824010216).run 79 | ActiveRecord::Base.on_all_shards do 80 | assert table_has_column?("tickets", "sharded_column") 81 | end 82 | 83 | migrator(:up, 'migrations', 20110829215912).run 84 | ActiveRecord::Base.on_shard(nil) do 85 | assert table_has_column?("accounts", "non_sharded_column") 86 | end 87 | end 88 | 89 | it "does not migrate bad migrations" do 90 | assert_raises StandardError do 91 | migrator(:up, 'cowardly_migration').migrate 92 | end 93 | end 94 | 95 | it "fails with failing migrations" do 96 | # like, if you have to break a migration in the middle somewhere. 97 | assert failure_migration_pending?('failure_migration') 98 | begin 99 | migrator(:up, 'failure_migration').migrate 100 | rescue StandardError => e 101 | unless e.message.include?("ERROR_IN_MIGRATION") 102 | raise e 103 | end 104 | 105 | # after first fail, should still be pending 106 | assert failure_migration_pending?('failure_migration') 107 | retry 108 | end 109 | 110 | assert !failure_migration_pending?('failure_migration') 111 | ActiveRecord::Base.on_all_shards do 112 | assert table_has_column?("tickets", "sharded_column") 113 | end 114 | end 115 | 116 | describe "#shard_status" do 117 | it "shows nothing if everything is ok" do 118 | assert_equal [{}, {}], ActiveRecord::Migrator.shard_status([1]) 119 | end 120 | 121 | it "shows missing migrations" do 122 | assert_equal [{}, { nil => [1], 0 => [1], 1 => [1] }], ActiveRecord::Migrator.shard_status([]) 123 | end 124 | 125 | it "shows pending migrations" do 126 | assert_equal [{ nil => [2], 0 => [2], 1 => [2] }, {}], ActiveRecord::Migrator.shard_status([1, 2]) 127 | end 128 | end 129 | 130 | private 131 | 132 | def failure_migration_pending?(migration_path) 133 | migrator(:up, migration_path).pending_migrations.detect { |f| f.name == "FailureMigration" } 134 | end 135 | end 136 | -------------------------------------------------------------------------------- /lib/active_record_shards.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require 'active_record' 4 | require 'active_record/base' 5 | require 'active_record_shards/configuration_parser' 6 | require 'active_record_shards/model' 7 | require 'active_record_shards/shard_selection' 8 | require 'active_record_shards/connection_switcher' 9 | require 'active_record_shards/association_collection_connection_selection' 10 | require 'active_record_shards/migration' 11 | require 'active_record_shards/default_replica_patches' 12 | require 'active_record_shards/default_shard' 13 | require 'active_record_shards/schema_dumper_extension' 14 | 15 | module ActiveRecordShards 16 | class << self 17 | attr_accessor :disable_replica_readonly_records 18 | end 19 | 20 | def self.app_env 21 | @app_env ||= begin 22 | env = Rails.env if defined?(Rails.env) 23 | env ||= RAILS_ENV if Object.const_defined?(:RAILS_ENV) 24 | env ||= ENV['RAILS_ENV'] 25 | env ||= APP_ENV if Object.const_defined?(:APP_ENV) 26 | env ||= ENV['APP_ENV'] 27 | env || 'development' 28 | end 29 | end 30 | 31 | # Busts internal caches kept by active_record_shards, for things which are _supposed_ to be the 32 | # same for the life of the process. You shouldn't need to call this unless you're doing something 33 | # truly evil like changing RAILS_ENV after boot 34 | def self.reset_app_env! 35 | @app_env = nil 36 | models = [ActiveRecord::Base] + ActiveRecord::Base.descendants 37 | models.each do |model| 38 | model.remove_instance_variable(:@_ars_model_is_sharded) if model.instance_variable_defined?(:@_ars_model_is_sharded) 39 | end 40 | end 41 | end 42 | 43 | ActiveRecord::Base.extend(ActiveRecordShards::ConfigurationParser) 44 | ActiveRecord::Base.extend(ActiveRecordShards::Model) 45 | ActiveRecord::Base.extend(ActiveRecordShards::ConnectionSwitcher) 46 | ActiveRecord::Base.extend(ActiveRecordShards::DefaultReplicaPatches) 47 | ActiveRecord::Relation.include(ActiveRecordShards::DefaultReplicaPatches::ActiveRelationPatches) 48 | ActiveRecord::Associations::CollectionProxy.include(ActiveRecordShards::AssociationCollectionConnectionSelection) 49 | ActiveRecord::Associations::Builder::HasAndBelongsToMany.include(ActiveRecordShards::DefaultReplicaPatches::Rails41HasAndBelongsToManyBuilderExtension) 50 | ActiveRecord::SchemaDumper.prepend(ActiveRecordShards::SchemaDumperExtension) 51 | 52 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 53 | when '5.1' 54 | # https://github.com/rails/rails/blob/v5.1.7/activerecord/lib/active_record/associations/association.rb#L97 55 | ActiveRecord::Associations::Association.prepend(ActiveRecordShards::DefaultReplicaPatches::AssociationsAssociationAssociationScopePatch) 56 | 57 | # https://github.com/rails/rails/blob/v5.1.7/activerecord/lib/active_record/associations/singular_association.rb#L41 58 | ActiveRecord::Associations::SingularAssociation.prepend(ActiveRecordShards::DefaultReplicaPatches::AssociationsAssociationFindTargetPatch) 59 | 60 | # https://github.com/rails/rails/blob/v5.1.7/activerecord/lib/active_record/associations/collection_association.rb#L305 61 | ActiveRecord::Associations::CollectionAssociation.prepend(ActiveRecordShards::DefaultReplicaPatches::AssociationsAssociationFindTargetPatch) 62 | 63 | # https://github.com/rails/rails/blob/v5.1.7/activerecord/lib/active_record/associations/preloader/association.rb#L120 64 | ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecordShards::DefaultReplicaPatches::AssociationsPreloaderAssociationLoadRecordsPatch) 65 | when '5.2' 66 | # https://github.com/rails/rails/blob/v5.2.6/activerecord/lib/active_record/relation.rb#L530 67 | # But the #exec_queries method also calls #connection, and I don't know if we should patch that one, too... 68 | ActiveRecord::Relation.prepend(ActiveRecordShards::DefaultReplicaPatches::Rails52RelationPatches) 69 | 70 | # https://github.com/rails/rails/blob/v5.2.6/activerecord/lib/active_record/associations/singular_association.rb#L42 71 | ActiveRecord::Associations::SingularAssociation.prepend(ActiveRecordShards::DefaultReplicaPatches::AssociationsAssociationFindTargetPatch) 72 | 73 | # https://github.com/rails/rails/blob/v5.2.6/activerecord/lib/active_record/associations/collection_association.rb#L308 74 | ActiveRecord::Associations::CollectionAssociation.prepend(ActiveRecordShards::DefaultReplicaPatches::AssociationsAssociationFindTargetPatch) 75 | 76 | # https://github.com/rails/rails/blob/v5.2.6/activerecord/lib/active_record/associations/preloader/association.rb#L96 77 | ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecordShards::DefaultReplicaPatches::AssociationsPreloaderAssociationLoadRecordsPatch) 78 | when '6.0', '6.1', '7.0' 79 | # https://github.com/rails/rails/blob/v6.0.4/activerecord/lib/active_record/type_caster/connection.rb#L28 80 | ActiveRecord::TypeCaster::Connection.prepend(ActiveRecordShards::DefaultReplicaPatches::TypeCasterConnectionConnectionPatch) 81 | 82 | # https://github.com/rails/rails/blob/v6.0.4/activerecord/lib/active_record/schema.rb#L53-L54 83 | ActiveRecord::Schema.prepend(ActiveRecordShards::DefaultReplicaPatches::SchemaDefinePatch) 84 | 85 | # https://github.com/rails/rails/blob/v6.0.4/activerecord/lib/active_record/relation.rb#L739 86 | # But the #exec_queries and #compute_cache_version methods also call #connection, and I don't know if we should patch those, too... 87 | ActiveRecord::Relation.prepend(ActiveRecordShards::DefaultReplicaPatches::Rails52RelationPatches) 88 | 89 | # https://github.com/rails/rails/blob/v6.0.4/activerecord/lib/active_record/associations/association.rb#L213 90 | ActiveRecord::Associations::Association.prepend(ActiveRecordShards::DefaultReplicaPatches::AssociationsAssociationFindTargetPatch) 91 | else 92 | raise "ActiveRecordShards is not compatible with #{ActiveRecord::VERSION::STRING}" 93 | end 94 | -------------------------------------------------------------------------------- /Changelog.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | All notable changes to this project will be documented in this file. 3 | 4 | The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 5 | and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html) (from 3.17.0 onwards) 6 | 7 | ## [Unreleased] 8 | 9 | ## v5.5.1 10 | ### Fixed 11 | * Fixed `SqlComments` when using Rails 7.0. 12 | 13 | ## v5.5.0 14 | 15 | ### Changed 16 | * Use the default shard when resetting primary key 17 | 18 | ## v5.4.1 19 | 20 | ### Fixed 21 | * Changes all fiber-local variables to thread-local variables. Fixes any lingering issues where a fiber change could cause the state to be lost. 22 | 23 | ## v5.4.0 24 | 25 | ### Changed 26 | * Now raises an error in Rails 7.0 when attempting to use ARS when `ActiveSupport::IsolatedExecutionState.isolation_level` is not set to `:thread` 27 | 28 | ### Fixed 29 | * Stays on the correct database, even when using a new Fiber. Some Ruby methods, such as `to_enum` create a new Fiber for the block. `to_enum` is used by ActiveRecord when finding things in batches. This should resolve issues where ARS would connect to the unsharded database, even inside an `on_shard` block. The downside is we are now violating fiber concurrency and thus this breaks multi-fiber webservers. 30 | 31 | * Fixes an issue where ARS switches to the replica database in the middle of a transaction when it is supposed to remain on the Primary. 32 | 33 | ## v5.3.3 34 | 35 | ### Added 36 | * Added new 'ActiveRecordShards.reset_rails_env!' method to clear the cache within a process ([#315](https://github.com/zendesk/active_record_shards/pull/315)). 37 | 38 | ## v5.3.2 39 | 40 | ### Added 41 | * Run tests against Ruby 3.2 (the gem actually already worked with Ruby 3.2, but now we _promise_ that it does) ([#309](https://github.com/zendesk/active_record_shards/pull/309)). 42 | * Cache the results of `is_sharded?` and `app_env`, and ensure that the configuration is only validated once in `shard_names`. This should be a small performance improvement ([#311](https://github.com/zendesk/active_record_shards/pull/311), [#312](https://github.com/zendesk/active_record_shards/pull/312)). 43 | 44 | ## v5.3.1 45 | 46 | ### Added 47 | Raises a new `LegacyConnectionHandlingError` exception when using ActiveRecord >= 6.1 and `legacy_connection_handling` is set to `false`. 48 | 49 | ## v5.3.0 50 | 51 | ### Fixed 52 | 53 | Make connection switching thread safe, by fixing a thread safety issue caused by using a (class) instance variable instead of a thread-local variable. 54 | 55 | ## v5.2.0 56 | 57 | ### Added 58 | 59 | Support for Rails 7.0 when `legacy_connection_handling` is set to `true`. This is required to [opt-out of the native Rails 6.1+ sharding support](https://guides.rubyonrails.org/active_record_multiple_databases.html). 60 | 61 | ### Fixed 62 | 63 | Rails 6.1 deprecation warnings. 64 | 65 | ## v5.1.0 66 | 67 | ### Added 68 | 69 | Support for Rails 6.1 when `legacy_connection_handling` is set to `true`. This is required to [opt-out of the native Rails 6.1 sharding support](https://guides.rubyonrails.org/active_record_multiple_databases.html). 70 | 71 | ## v5.0.0 72 | 73 | ### Changed 74 | 75 | Rename `ActiveRecordShards.rails_env` to `ActiveRecordShards.app_env`, and include `APP_ENV` and `ENV['APP_ENV']` in the list of places it looks for environment information. 76 | 77 | Removed support for Ruby 2.3, 2.4, and 2.5. 78 | 79 | Removed support for Rails 4.2 and 5.0. 80 | 81 | [Deprecation] Removes all deprecated methods containing `master`/`slave`. Use the updated `primary`/`replica` methods instead. The main public methods changed: 82 | 83 | 1. `on_slave` => `on_replica` 84 | 1. `on_master` => `on_primary` 85 | 86 | other methods changed: 87 | 88 | 1. `on_master_if` => `on_primary_if` 89 | 1. `on_slave_if` => `on_replica_if` 90 | 1. `on_master_unless` => `on_primary_unless` 91 | 1. `on_slave_unless` => `on_replica_unless` 92 | 1. `on_master_or_slave` => `on_primary_or_replica` 93 | 1. `exists_with_default_slave` => `exists_with_default_replica` 94 | 1. `from_slave` => `from_replica` 95 | 1. `initialize_shard_and_slave` => `initialize_shard_and_replica` 96 | 1. `ShardSelection#options` no longer uses `:slave`, if this method was overridden ensure it returns `:replica` instead of `:slave`: `{ shard: .., replica: ... }` 97 | 98 | Also removes the class `ActiveRecordShards::Deprecation`. 99 | 100 | ### Added 101 | 102 | Add a global setting to disable marking instances from replicas as read-only. To enable: 103 | 104 | `ActiveRecordShards.disable_replica_readonly_records = true` 105 | 106 | ## v3.19.1 107 | 108 | ### Fixed 109 | 110 | Converts the `ActiveRecord::Base.configurations` object introduced in Rails 6 into a hash as expected. 111 | 112 | ## v3.19.0 113 | 114 | ### Changed / Fixed 115 | 116 | Lots of improvements to the `on_replica_by_default` logic, now covered by an improved test suite. Schema loading should now _always_ happen on the replica databases, and non-mutating queries will should now happen on the replica except when `on_replica_by_default` is not configured. 117 | 118 | ## v3.18.0 119 | 120 | ### Changed / Deprecated 121 | 122 | Adds deprecation warning for all methods containing `master`/`slave` which recommends using the updated `primary`/`replica` methods. The main public methods changed: 123 | 124 | 1. `on_slave` => `on_replica` 125 | 1. `on_master` => `on_primary` 126 | 127 | other methods changed: 128 | 129 | 1. `on_master_if` => `on_primary_if` 130 | 1. `on_slave_if` => `on_replica_if` 131 | 1. `on_master_unless` => `on_primary_unless` 132 | 1. `on_slave_unless` => `on_replica_unless` 133 | 1. `on_master_or_slave` => `on_primary_or_replica` 134 | 1. `exists_with_default_slave` => `exists_with_default_replica` 135 | 1. `from_slave` => `from_replica` 136 | 1. `initialize_shard_and_slave` => `initialize_shard_and_replica` 137 | 1. `ShardSelection#options` no longer uses `:slave`, if this method was overridden ensure it returns `:replica` instead of `:slave`: `{ shard: .., replica: ... }` 138 | 139 | SQL comments (see [debugging](/README.md#debugging)) will now log `... /* replica */` instead of `... /* slave */` 140 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | > [!CAUTION] 2 | > ActiveRecordShards is archived. Rails 6.1 introduced native support for multiple databases, including horizontal shards, please use the native functionality instead. 3 | 4 | [![Build Status](https://github.com/zendesk/active_record_shards/workflows/CI/badge.svg)](https://github.com/zendesk/active_record_shards/actions?query=workflow%3ACI) 5 | 6 | # ActiveRecord Shards 7 | 8 | ActiveRecord Shards is an extension for ActiveRecord that provides support for sharded database and replicas. Basically it is just a nice way to 9 | switch between database connections. We've made the implementation very small, and have tried not to reinvent any wheels already present in ActiveRecord. 10 | 11 | ActiveRecord Shards has been used and tested on Rails 5.x and 6.0, and has in some form or another been used in production on large Rails apps for several years. 12 | 13 | Rails 6.1 introduced new connection handling and support for sharding. Apps are encouraged to migrate to the native sharding logic but ActiveRecord Shards supports Rails 6.1 when `legacy_connection_handling` is set to `true`. For more information see [Rails 6.1 installation](#rails-61-installation) and Rails' [multiple databases guide](https://guides.rubyonrails.org/active_record_multiple_databases.html). 14 | 15 | - [Installation](#installation) 16 | - [Configuration](#configuration) 17 | - [Migrations](#migrations) 18 | - [Example](#example) 19 | - [Shared Model](#create-a-table-for-the-shared-not-sharded-model) 20 | - [Sharded Model](#create-a-table-for-the-sharded-model) 21 | - [Usage](#usage) 22 | - [Debugging](#debugging) 23 | 24 | ## Installation 25 | 26 | $ gem install active_record_shards 27 | 28 | and make sure to require 'active\_record\_shards' in some way. 29 | 30 | ### Rails 6.1 & 7.0 installation 31 | 32 | Rails 6.1 & 7.0 are **only** supported with `legacy_connection_handling` set to `true`. 33 | 34 | Enable the legacy handling in your configuration files e.g. `config/application.rb` by setting: 35 | 36 | ``` Ruby 37 | config.active_record.legacy_connection_handling = true 38 | ``` 39 | 40 | or 41 | 42 | ``` Ruby 43 | ActiveRecord::Base.legacy_connection_handling = true 44 | ``` 45 | 46 | ## Configuration 47 | 48 | Add the replica and shard configuration to config/database.yml: 49 | 50 | ```yaml 51 | production: 52 | adapter: mysql 53 | encoding: utf8 54 | database: my_app_main 55 | pool: 5 56 | host: db1 57 | username: root 58 | password: 59 | replica: 60 | host: db1_replica 61 | shards: 62 | 1: 63 | host: db_shard1 64 | database: my_app_shard 65 | replica: 66 | host: db_shard1_replica 67 | 2: 68 | host: db_shard2 69 | database: my_app_shard 70 | replica: 71 | host: db_shard2_replica 72 | ``` 73 | 74 | basically connections inherit configuration from the parent configuration file. 75 | 76 | ## Migrations 77 | 78 | ActiveRecord Shards also patches migrations to support running migrations on a shared (not sharded) or a sharded database. 79 | Each migration class has to specify a shard spec indicating where to run the migration. 80 | 81 | Valid shard specs: 82 | 83 | * `:none` - Run this migration on the shared database, not any shards 84 | * `:all` - Run this migration on all of the shards, not the shared database 85 | 86 | #### Example 87 | 88 | ###### Create a table for the shared (not sharded) model 89 | 90 | ```ruby 91 | class CreateAccounts < ActiveRecord::Migration 92 | shard :none 93 | 94 | def change 95 | create_table :accounts do |t| 96 | # This is NOT necessary for the gem to work, we just use it in the examples below demonstrating one way to switch shards 97 | t.integer :shard_id, null: false 98 | 99 | t.string :name 100 | end 101 | end 102 | end 103 | ``` 104 | 105 | ###### Create a table for the sharded model 106 | 107 | ```ruby 108 | class CreateProjects < ActiveRecord::Migration 109 | shard :all 110 | 111 | def change 112 | create_table :projects do |t| 113 | t.references :account 114 | t.string :name 115 | end 116 | end 117 | end 118 | ``` 119 | 120 | ## Usage 121 | 122 | Normally you have some models that live on a shared database, and you might need to query this data in order to know what shard to switch to. 123 | All the models that live on the shared database must be marked as not\_sharded: 124 | 125 | ```ruby 126 | class Account < ActiveRecord::Base 127 | not_sharded 128 | 129 | has_many :projects 130 | end 131 | 132 | class Project < ActiveRecord::Base 133 | belongs_to :account 134 | end 135 | ``` 136 | 137 | So in this setup the accounts live on the shared database, but the projects are sharded. If accounts have a shard\_id column, you could lookup the account 138 | in a rack middleware and switch to the right shard: 139 | 140 | ```ruby 141 | class AccountMiddleware 142 | def initialize(app) 143 | @app = app 144 | end 145 | 146 | def call(env) 147 | account = lookup_account(env) 148 | 149 | if account 150 | ActiveRecord::Base.on_shard(account.shard_id) do 151 | @app.call(env) 152 | end 153 | else 154 | @app.call(env) 155 | end 156 | end 157 | 158 | def lookup_account(env) 159 | # ... 160 | end 161 | end 162 | ``` 163 | 164 | You can switch to the replica databases at any point by wrapping your code in an on\_replica block: 165 | 166 | ```ruby 167 | ActiveRecord::Base.on_replica do 168 | Account.find_by_big_expensive_query 169 | end 170 | ``` 171 | 172 | This will perform the query on the replica, and mark the returned instances as read-only. There is also a shortcut for this: 173 | 174 | ```ruby 175 | Account.on_replica.find_by_big_expensive_query 176 | ``` 177 | 178 | If you do not want instances returned from replicas to be marked as read-only, this can be disabled globally: 179 | 180 | `ActiveRecordShards.disable_replica_readonly_records = true` 181 | 182 | ## Debugging 183 | 184 | Show if a query went to primary or replica in the logs: 185 | 186 | ```Ruby 187 | require 'active_record_shards/sql_comments' 188 | ActiveRecordShards::SqlComments.enable 189 | ``` 190 | 191 | ## Copyright 192 | 193 | Copyright (c) 2011 Zendesk. See LICENSE for details. 194 | 195 | ## Authors 196 | Mick Staugaard, Eric Chapweske 197 | -------------------------------------------------------------------------------- /lib/active_record_shards/tasks.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require 'active_record_shards' 4 | 5 | %w[db:drop db:create db:abort_if_pending_migrations db:reset db:test:purge].each do |name| 6 | Rake::Task[name].clear 7 | end 8 | 9 | namespace :db do 10 | desc 'Drops the database for the current RAILS_ENV including shards' 11 | task drop: :load_config do 12 | configurations = begin 13 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 14 | when '6.1', '7.0' 15 | ActiveRecord::Base.configurations.configurations.map { |configuration| [configuration.env_name, configuration.configuration_hash] } 16 | else 17 | ActiveRecord::Base.configurations.to_h 18 | end 19 | end 20 | 21 | configurations.each do |key, conf| 22 | next if !key.start_with?(ActiveRecordShards.app_env) || key.end_with?("_replica") 23 | 24 | begin 25 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 26 | when '6.1', '7.0' 27 | conf = conf.stringify_keys 28 | end 29 | 30 | ActiveRecordShards::Tasks.root_connection(conf).drop_database(conf['database']) 31 | # rescue ActiveRecord::NoDatabaseError # TODO: exists in AR but never is raised here ... 32 | # $stderr.puts "Database '#{conf['database']}' does not exist" 33 | rescue StandardError => e 34 | warn e, *e.backtrace 35 | warn "Couldn't drop #{conf['database']}" 36 | end 37 | end 38 | end 39 | 40 | task reset: :load_config do |t| 41 | Rake.application.lookup('db:drop', t.scope).invoke rescue nil # rubocop:disable Style/RescueModifier 42 | Rake.application.lookup('db:setup', t.scope).invoke 43 | end 44 | 45 | desc "Create the database defined in config/database.yml for the current RAILS_ENV including shards" 46 | task create: :load_config do 47 | configurations = begin 48 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 49 | when '6.1', '7.0' 50 | ActiveRecord::Base.configurations.configurations.map { |configuration| [configuration.env_name, configuration.configuration_hash] } 51 | else 52 | ActiveRecord::Base.configurations.to_h 53 | end 54 | end 55 | configurations.each do |key, conf| 56 | next if !key.start_with?(ActiveRecordShards.app_env) || key.end_with?("_replica") 57 | 58 | begin 59 | # MysqlAdapter takes charset instead of encoding in Rails 4.2 or greater 60 | # https://github.com/rails/rails/blob/4-2-stable/activerecord/lib/active_record/tasks/mysql_database_tasks.rb#L85-L96 61 | symbolized_configuration = conf.symbolize_keys 62 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 63 | when '6.1', '7.0' 64 | conf = conf.stringify_keys 65 | end 66 | 67 | symbolized_configuration[:charset] = symbolized_configuration[:encoding] 68 | 69 | ActiveRecordShards::Tasks.root_connection(conf).create_database(conf['database'], symbolized_configuration) 70 | rescue ActiveRecord::StatementInvalid => e 71 | if e.message.include?('database exists') 72 | puts "#{conf['database']} already exists" 73 | else 74 | raise e 75 | end 76 | end 77 | end 78 | ActiveRecord::Base.establish_connection(ActiveRecordShards.app_env.to_sym) 79 | end 80 | 81 | desc "Raises an error if there are pending migrations" 82 | task abort_if_pending_migrations: :environment do 83 | if defined? ActiveRecord 84 | pending_migrations = 85 | if ActiveRecord::VERSION::MAJOR >= 6 86 | migrations = ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths, ActiveRecord::SchemaMigration).migrations 87 | ActiveRecord::Migrator.new(:up, migrations, ActiveRecord::SchemaMigration).pending_migrations 88 | elsif ActiveRecord::VERSION::STRING >= "5.2.0" 89 | migrations = ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths).migrations 90 | ActiveRecord::Migrator.new(:up, migrations).pending_migrations 91 | else 92 | ActiveRecord::Base.on_shard(nil) { ActiveRecord::Migrator.open(ActiveRecord::Migrator.migrations_paths).pending_migrations } 93 | end 94 | 95 | if pending_migrations.any? 96 | warn "You have #{pending_migrations.size} pending migrations:" 97 | pending_migrations.each do |pending_migration| 98 | warn ' %4d %s' % [pending_migration.version, pending_migration.name] 99 | end 100 | abort %(Run "rake db:migrate" to update your database then try again.) 101 | end 102 | end 103 | end 104 | 105 | namespace :test do 106 | desc 'Purges the test databases by dropping and creating' 107 | task purge: :load_config do |t| 108 | saved_env = Rails.env 109 | Rails.env = 'test' 110 | Rake.application.lookup('db:drop', t.scope).execute 111 | Rake.application.lookup('db:create', t.scope).execute 112 | ensure 113 | Rails.env = saved_env 114 | end 115 | end 116 | end 117 | 118 | module ActiveRecordShards 119 | module Tasks 120 | class << self 121 | def root_connection(conf) 122 | conf = conf.merge('database' => nil) 123 | spec = spec_for(conf) 124 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 125 | when '6.1', '7.0' 126 | ActiveRecord::Base.send("#{conf['adapter']}_connection", spec.db_config.configuration_hash) 127 | else 128 | ActiveRecord::Base.send("#{conf['adapter']}_connection", spec.config) 129 | end 130 | end 131 | 132 | private 133 | 134 | def spec_for(conf) 135 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 136 | when '7.0' 137 | ActiveRecord::Base.connection_handler.send(:resolve_pool_config, conf, ActiveRecord::Base, ActiveRecord::Base.current_role, ActiveRecord::Base.current_shard) 138 | when '6.1' 139 | ActiveRecord::Base.connection_handler.send(:resolve_pool_config, conf, ActiveRecord::Base) 140 | else 141 | resolver = ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new(ActiveRecord::Base.configurations) 142 | resolver.spec(conf) 143 | end 144 | end 145 | end 146 | end 147 | end 148 | -------------------------------------------------------------------------------- /test/helper.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | # Stop Minitest creating threads we won't use. 4 | # They add noise to the thread safety tests when inspecting `Thread.list`. 5 | ENV["MT_CPU"] ||= "1" 6 | 7 | require 'bundler/setup' 8 | require 'minitest/autorun' 9 | require 'rake' 10 | 11 | Bundler.require 12 | 13 | $LOAD_PATH.unshift(File.join(__dir__, '..', 'lib')) 14 | $LOAD_PATH.unshift(__dir__) 15 | require 'active_support' 16 | require 'active_record_shards' 17 | require 'mysql2' 18 | require 'support/db_helper' 19 | require 'support/tcp_proxy' 20 | require 'logger' 21 | 22 | require 'pry-byebug' 23 | 24 | RAILS_ENV = "test" 25 | 26 | ActiveRecord::Base.logger = Logger.new(__dir__ + "/test.log") 27 | ActiveSupport.test_order = :sorted 28 | ActiveSupport::Deprecation.behavior = :raise 29 | 30 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 31 | when '7.0' 32 | ActiveRecord.legacy_connection_handling = true 33 | when '6.1' 34 | ActiveRecord::Base.legacy_connection_handling = true 35 | end 36 | 37 | BaseMigration = ActiveRecord::Migration[4.2] 38 | 39 | require 'active_support/test_case' 40 | require 'models' 41 | 42 | # support multiple before/after blocks per example 43 | module SpecDslPatch 44 | def before(_type = nil, &block) 45 | prepend( 46 | Module.new do 47 | define_method(:setup) do 48 | super() 49 | instance_exec(&block) 50 | end 51 | end 52 | ) 53 | end 54 | 55 | def after(_type = nil, &block) 56 | prepend( 57 | Module.new do 58 | define_method(:teardown) do 59 | instance_exec(&block) 60 | super() 61 | end 62 | end 63 | ) 64 | end 65 | end 66 | Minitest::Spec.singleton_class.prepend(SpecDslPatch) 67 | 68 | module RakeSpecHelpers 69 | def show_databases(config) 70 | client = Mysql2::Client.new( 71 | host: config['test']['host'], 72 | port: config['test']['port'], 73 | username: config['test']['username'], 74 | password: config['test']['password'] 75 | ) 76 | databases = client.query("SHOW DATABASES") 77 | databases.map { |d| d['Database'] } 78 | end 79 | 80 | def rake(name) 81 | Rake::Task[name].reenable 82 | Rake::Task[name].invoke 83 | end 84 | end 85 | 86 | module ConnectionSwitchingSpecHelpers 87 | def assert_using_primary_db 88 | assert_using_database('ars_test') 89 | end 90 | 91 | def assert_using_replica_db 92 | assert_using_database('ars_test_replica') 93 | end 94 | 95 | def assert_using_database(db_name, model = ActiveRecord::Base) 96 | assert_equal(db_name, model.connection.current_database) 97 | end 98 | end 99 | 100 | module SpecHelpers 101 | def self.mysql_url 102 | URI(ENV['MYSQL_URL'] || 'mysql://root@127.0.0.1:3306') 103 | end 104 | 105 | @@unsharded_primary_proxy ||= TCPProxy.start( 106 | remote_host: mysql_url.host, 107 | remote_port: mysql_url.port, 108 | local_port: '13306' 109 | ) 110 | 111 | @@shard_1_primary_proxy ||= TCPProxy.start( 112 | remote_host: mysql_url.host, 113 | remote_port: mysql_url.port, 114 | local_port: '13307' 115 | ) 116 | 117 | @@shard_2_primary_proxy ||= TCPProxy.start( 118 | remote_host: mysql_url.host, 119 | remote_port: mysql_url.port, 120 | local_port: '13308' 121 | ) 122 | 123 | # Verifies that a block of code is not using any of the the primaries by 124 | # pausing the TCP proxies between Ruby and MySQL. 125 | def with_all_primaries_unavailable 126 | with_unsharded_primary_unavailable do 127 | with_sharded_primaries_unavailable do 128 | yield 129 | end 130 | end 131 | end 132 | 133 | # Verifies that a block of code is not using the unsharded primary by pausing 134 | # the TCP proxy between Ruby and MySQL. 135 | def with_unsharded_primary_unavailable 136 | ActiveRecord::Base.connection_handler.clear_all_connections! 137 | @@unsharded_primary_proxy.pause do 138 | yield 139 | end 140 | end 141 | 142 | # Verifies that a block of code is not using the sharded primaries by pausing 143 | # the TCP proxies between Ruby and MySQL. 144 | def with_sharded_primaries_unavailable 145 | ActiveRecord::Base.connection_handler.clear_all_connections! 146 | @@shard_1_primary_proxy.pause do 147 | @@shard_2_primary_proxy.pause do 148 | yield 149 | end 150 | end 151 | end 152 | 153 | def clear_global_connection_handler_state 154 | # Close active connections 155 | ActiveRecord::Base.connection_handler.clear_all_connections! 156 | 157 | # Use a fresh connection handler 158 | ActiveRecord::Base.connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new 159 | end 160 | 161 | def table_exists?(name) 162 | ActiveRecord::Base.connection.data_source_exists?(name) 163 | end 164 | 165 | def table_has_column?(table, column) 166 | !ActiveRecord::Base.connection.select_values("desc #{table}").grep(column).empty? 167 | end 168 | 169 | def migrator(direction = :up, path = 'migrations', target_version = nil) 170 | migration_path = File.join(__dir__, "support", path) 171 | if ActiveRecord::VERSION::MAJOR >= 6 172 | migrations = ActiveRecord::MigrationContext.new(migration_path, ActiveRecord::SchemaMigration).migrations 173 | ActiveRecord::Migrator.new(direction, migrations, ActiveRecord::SchemaMigration, target_version) 174 | elsif ActiveRecord::VERSION::STRING >= "5.2.0" 175 | migrations = ActiveRecord::MigrationContext.new(migration_path).migrations 176 | ActiveRecord::Migrator.new(direction, migrations, target_version) 177 | else 178 | migrations = ActiveRecord::Migrator.migrations(migration_path) 179 | ActiveRecord::Migrator.new(direction, migrations, target_version) 180 | end 181 | end 182 | end 183 | Minitest::Spec.include(SpecHelpers) 184 | 185 | module RailsEnvSwitch 186 | def switch_app_env(env) 187 | before do 188 | silence_warnings { Object.const_set("RAILS_ENV", env) } 189 | ActiveRecordShards.reset_app_env! 190 | ActiveRecord::Base.establish_connection(::RAILS_ENV.to_sym) 191 | end 192 | after do 193 | silence_warnings { Object.const_set("RAILS_ENV", 'test') } 194 | ActiveRecordShards.reset_app_env! 195 | ActiveRecord::Base.establish_connection(::RAILS_ENV.to_sym) 196 | tmp_sharded_model = Class.new(ActiveRecord::Base) 197 | assert_equal('ars_test', tmp_sharded_model.connection.current_database) 198 | end 199 | end 200 | end 201 | 202 | Minitest::Spec.extend(DbHelper) 203 | -------------------------------------------------------------------------------- /lib/active_record_shards/default_replica_patches.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | module ActiveRecordShards 4 | module DefaultReplicaPatches 5 | def self.wrap_method_in_on_replica(class_method, base, method, force_on_replica: false) 6 | base_methods = 7 | if class_method 8 | base.methods + base.private_methods 9 | else 10 | base.instance_methods + base.private_instance_methods 11 | end 12 | 13 | return unless base_methods.include?(method) 14 | 15 | _, method, punctuation = method.to_s.match(/^(.*?)([\?\!]?)$/).to_a 16 | # _ALWAYS_ on replica, or only for on `on_replica_by_default = true` models? 17 | wrapper = force_on_replica ? 'force_on_replica' : 'on_replica_unless_tx' 18 | base.class_eval <<-RUBY, __FILE__, __LINE__ + 1 19 | #{class_method ? 'class << self' : ''} 20 | def #{method}_with_default_replica#{punctuation}(*args, &block) 21 | #{wrapper} do 22 | #{method}_without_default_replica#{punctuation}(*args, &block) 23 | end 24 | end 25 | ruby2_keywords(:#{method}_with_default_replica#{punctuation}) if respond_to?(:ruby2_keywords, true) 26 | alias_method :#{method}_without_default_replica#{punctuation}, :#{method}#{punctuation} 27 | alias_method :#{method}#{punctuation}, :#{method}_with_default_replica#{punctuation} 28 | #{class_method ? 'end' : ''} 29 | RUBY 30 | end 31 | 32 | module InstanceMethods 33 | def on_replica_unless_tx 34 | self.class.on_replica_unless_tx { yield } 35 | end 36 | end 37 | 38 | CLASS_REPLICA_METHODS = [ 39 | :calculate, 40 | :count_by_sql, 41 | :exists?, 42 | :find, 43 | :find_by, 44 | :find_by_sql, 45 | :find_every, 46 | :find_one, 47 | :find_some, 48 | :get_primary_key 49 | ].freeze 50 | 51 | CLASS_FORCE_REPLICA_METHODS = [ 52 | :replace_bind_variable, 53 | :replace_bind_variables, 54 | :sanitize_sql_array, 55 | :sanitize_sql_hash_for_assignment, 56 | :table_exists? 57 | ].freeze 58 | 59 | def self.extended(base) 60 | CLASS_REPLICA_METHODS.each { |m| ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(true, base, m) } 61 | CLASS_FORCE_REPLICA_METHODS.each { |m| ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(true, base, m, force_on_replica: true) } 62 | 63 | ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(true, base, :load_schema!, force_on_replica: true) 64 | ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(false, base, :reload) 65 | 66 | base.class_eval do 67 | include InstanceMethods 68 | end 69 | end 70 | 71 | def on_replica_unless_tx(&block) 72 | return yield if Thread.current._active_record_shards_in_migration 73 | return yield if _in_transaction? 74 | 75 | if on_replica_by_default? 76 | on_replica(&block) 77 | else 78 | yield 79 | end 80 | end 81 | 82 | def _in_transaction? 83 | connected? && connection.transaction_open? 84 | end 85 | 86 | def force_on_replica(&block) 87 | return yield if Thread.current._active_record_shards_in_migration 88 | 89 | on_cx_switch_block(:replica, construct_ro_scope: false, force: true, &block) 90 | end 91 | 92 | module ActiveRelationPatches 93 | def self.included(base) 94 | [:calculate, :exists?, :pluck, :load].each do |m| 95 | ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(false, base, m) 96 | end 97 | 98 | ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(false, base, :to_sql, force_on_replica: true) 99 | end 100 | 101 | def on_replica_unless_tx 102 | @klass.on_replica_unless_tx { yield } 103 | end 104 | end 105 | 106 | module Rails52RelationPatches 107 | def connection 108 | return super if Thread.current._active_record_shards_in_migration 109 | return super if _in_transaction? 110 | 111 | if @klass.on_replica_by_default? 112 | @klass.on_replica.connection 113 | else 114 | super 115 | end 116 | end 117 | end 118 | 119 | # in rails 4.1+, they create a join class that's used to pull in records for HABTM. 120 | # this simplifies the hell out of our existence, because all we have to do is inerit on-replica-by-default 121 | # down from the parent now. 122 | module Rails41HasAndBelongsToManyBuilderExtension 123 | def self.included(base) 124 | base.class_eval do 125 | alias_method :through_model_without_inherit_default_replica_from_lhs, :through_model 126 | alias_method :through_model, :through_model_with_inherit_default_replica_from_lhs 127 | end 128 | end 129 | 130 | def through_model_with_inherit_default_replica_from_lhs 131 | model = through_model_without_inherit_default_replica_from_lhs 132 | def model.on_replica_by_default? 133 | left_reflection.klass.on_replica_by_default? 134 | end 135 | 136 | # also transfer the sharded-ness of the left table to the join model 137 | model.not_sharded unless model.left_reflection.klass.is_sharded? 138 | model 139 | end 140 | end 141 | 142 | module AssociationsAssociationAssociationScopePatch 143 | def association_scope 144 | if klass 145 | on_replica_unless_tx { super } 146 | else 147 | super 148 | end 149 | end 150 | 151 | def on_replica_unless_tx 152 | klass.on_replica_unless_tx { yield } 153 | end 154 | end 155 | 156 | module AssociationsAssociationFindTargetPatch 157 | def find_target 158 | if klass 159 | on_replica_unless_tx { super } 160 | else 161 | super 162 | end 163 | end 164 | 165 | def on_replica_unless_tx 166 | klass.on_replica_unless_tx { yield } 167 | end 168 | end 169 | 170 | module AssociationsPreloaderAssociationAssociatedRecordsByOwnerPatch 171 | def associated_records_by_owner(preloader) 172 | if klass 173 | on_replica_unless_tx { super } 174 | else 175 | super 176 | end 177 | end 178 | 179 | def on_replica_unless_tx 180 | klass.on_replica_unless_tx { yield } 181 | end 182 | end 183 | 184 | module AssociationsPreloaderAssociationLoadRecordsPatch 185 | def load_records 186 | if klass 187 | on_replica_unless_tx { super } 188 | else 189 | super 190 | end 191 | end 192 | 193 | def on_replica_unless_tx 194 | klass.on_replica_unless_tx { yield } 195 | end 196 | end 197 | 198 | module TypeCasterConnectionConnectionPatch 199 | def connection 200 | return super if Thread.current._active_record_shards_in_migration 201 | return super if ActiveRecord::Base._in_transaction? 202 | 203 | if @klass.on_replica_by_default? 204 | @klass.on_replica.connection 205 | else 206 | super 207 | end 208 | end 209 | end 210 | 211 | module SchemaDefinePatch 212 | def define(info, &block) 213 | old_val = Thread.current._active_record_shards_in_migration 214 | Thread.current._active_record_shards_in_migration = true 215 | super 216 | ensure 217 | Thread.current._active_record_shards_in_migration = old_val 218 | end 219 | end 220 | end 221 | end 222 | -------------------------------------------------------------------------------- /lib/active_record_shards/connection_switcher.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require 'active_record_shards/shard_support' 4 | 5 | module ActiveRecordShards 6 | module ConnectionSwitcher 7 | class LegacyConnectionHandlingError < StandardError; end 8 | class IsolationLevelError < StandardError; end 9 | 10 | Thread.attr_accessor :_active_record_shards_disallow_replica_by_thread, 11 | :_active_record_shards_in_migration, 12 | :_active_record_shards_shard_selection 13 | 14 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 15 | when '6.1', '7.0' 16 | SHARD_NAMES_CONFIG_KEY = :shard_names 17 | else 18 | SHARD_NAMES_CONFIG_KEY = 'shard_names' 19 | end 20 | 21 | def self.extended(base) 22 | base.singleton_class.send(:alias_method, :load_schema_without_default_shard!, :load_schema!) 23 | base.singleton_class.send(:alias_method, :load_schema!, :load_schema_with_default_shard!) 24 | 25 | base.singleton_class.send(:alias_method, :table_exists_without_default_shard?, :table_exists?) 26 | base.singleton_class.send(:alias_method, :table_exists?, :table_exists_with_default_shard?) 27 | 28 | base.singleton_class.send(:alias_method, :reset_primary_key_without_default_shard, :reset_primary_key) 29 | base.singleton_class.send(:alias_method, :reset_primary_key, :reset_primary_key_with_default_shard) 30 | end 31 | 32 | def on_primary_db(&block) 33 | on_shard(nil, &block) 34 | end 35 | 36 | def on_shard(shard) 37 | old_options = current_shard_selection.options 38 | switch_connection(shard: shard) if supports_sharding? 39 | yield 40 | ensure 41 | switch_connection(old_options) 42 | end 43 | 44 | def on_first_shard(&block) 45 | shard_name = shard_names.first 46 | on_shard(shard_name, &block) 47 | end 48 | 49 | def shards 50 | ShardSupport.new(self == ActiveRecord::Base ? nil : where(nil)) 51 | end 52 | 53 | def on_all_shards 54 | old_options = current_shard_selection.options 55 | if supports_sharding? 56 | shard_names.map do |shard| 57 | switch_connection(shard: shard) 58 | yield(shard) 59 | end 60 | else 61 | [yield] 62 | end 63 | ensure 64 | switch_connection(old_options) 65 | end 66 | 67 | def on_replica_if(condition, &block) 68 | condition ? on_replica(&block) : yield 69 | end 70 | 71 | def on_replica_unless(condition, &block) 72 | on_replica_if(!condition, &block) 73 | end 74 | 75 | def on_primary_if(condition, &block) 76 | condition ? on_primary(&block) : yield 77 | end 78 | 79 | def on_primary_unless(condition, &block) 80 | on_primary_if(!condition, &block) 81 | end 82 | 83 | def on_primary_or_replica(which, &block) 84 | if block_given? 85 | on_cx_switch_block(which, &block) 86 | else 87 | PrimaryReplicaProxy.new(self, which) 88 | end 89 | end 90 | 91 | # Executes queries using the replica database. Fails over to primary if no replica is found. 92 | # if you want to execute a block of code on the replica you can go: 93 | # Account.on_replica do 94 | # Account.first 95 | # end 96 | # the first account will be found on the replica DB 97 | # 98 | # For one-liners you can simply do 99 | # Account.on_replica.first 100 | def on_replica(&block) 101 | on_primary_or_replica(:replica, &block) 102 | end 103 | 104 | def on_primary(&block) 105 | on_primary_or_replica(:primary, &block) 106 | end 107 | 108 | def on_cx_switch_block(which, force: false, construct_ro_scope: nil, &block) 109 | self.disallow_replica += 1 if which == :primary 110 | 111 | switch_to_replica = force || disallow_replica.zero? 112 | old_options = current_shard_selection.options 113 | 114 | switch_connection(replica: switch_to_replica) 115 | 116 | # we avoid_readonly_scope to prevent some stack overflow problems, like when 117 | # .columns calls .with_scope which calls .columns and onward, endlessly. 118 | if self == ActiveRecord::Base || !switch_to_replica || construct_ro_scope == false || ActiveRecordShards.disable_replica_readonly_records == true 119 | yield 120 | else 121 | readonly.scoping(&block) 122 | end 123 | ensure 124 | self.disallow_replica -= 1 if which == :primary 125 | switch_connection(old_options) if old_options 126 | end 127 | 128 | def disallow_replica=(value) 129 | Thread.current._active_record_shards_disallow_replica_by_thread = value 130 | end 131 | 132 | def disallow_replica 133 | Thread.current._active_record_shards_disallow_replica_by_thread ||= 0 134 | end 135 | 136 | def supports_sharding? 137 | shard_names.any? 138 | end 139 | 140 | def on_replica? 141 | current_shard_selection.on_replica? 142 | end 143 | 144 | def current_shard_selection 145 | Thread.current._active_record_shards_shard_selection ||= ShardSelection.new 146 | end 147 | 148 | def current_shard_id 149 | current_shard_selection.shard 150 | end 151 | 152 | def shard_names 153 | config_for_env[SHARD_NAMES_CONFIG_KEY] || [] 154 | end 155 | 156 | def reset_primary_key_with_default_shard 157 | with_default_shard { reset_primary_key_without_default_shard } 158 | end 159 | 160 | private 161 | 162 | def config_for_env 163 | @_ars_config_for_env ||= {} 164 | @_ars_config_for_env[shard_env] ||= begin 165 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 166 | when '7.0' 167 | config = configurations.configs_for(env_name: shard_env, include_hidden: true).first.configuration_hash 168 | when '6.1' 169 | config = configurations.configs_for(env_name: shard_env, include_replicas: true).first.configuration_hash 170 | else 171 | config = configurations[shard_env] 172 | end 173 | unless config 174 | raise "Did not find #{shard_env} in configurations, did you forget to add it to your database config? (configurations: #{configurations.to_h.keys.inspect})" 175 | end 176 | 177 | ensure_all_shard_names_are_integers(config) 178 | 179 | config 180 | end 181 | end 182 | alias_method :check_config_for_env, :config_for_env 183 | 184 | def ensure_all_shard_names_are_integers(config) 185 | unless config.fetch(SHARD_NAMES_CONFIG_KEY, []).all? { |shard_name| shard_name.is_a?(Integer) } 186 | raise "All shard names must be integers: #{config.inspect}." 187 | end 188 | end 189 | 190 | def switch_connection(options) 191 | ensure_legacy_connection_handling if ActiveRecord.version >= Gem::Version.new('6.1') 192 | ensure_thread_isolation_level if ActiveRecord.version >= Gem::Version.new('7.0') 193 | 194 | if options.any? 195 | if options.key?(:replica) 196 | current_shard_selection.on_replica = options[:replica] 197 | end 198 | 199 | if options.key?(:shard) 200 | check_config_for_env 201 | 202 | current_shard_selection.shard = options[:shard] 203 | end 204 | 205 | ensure_shard_connection 206 | end 207 | end 208 | 209 | def ensure_legacy_connection_handling 210 | unless legacy_connection_handling_owner.legacy_connection_handling 211 | raise LegacyConnectionHandlingError, "ActiveRecordShards is _only_ compatible with ActiveRecord `legacy_connection_handling` set to `true`." 212 | end 213 | end 214 | 215 | def ensure_thread_isolation_level 216 | unless ActiveSupport::IsolatedExecutionState.isolation_level == :thread 217 | raise IsolationLevelError, "ActiveRecordShards is _only_ compatible when ActiveSupport::IsolatedExecutionState's isolation_level is set to :thread" 218 | end 219 | end 220 | 221 | def legacy_connection_handling_owner 222 | @legacy_connection_handling_owner ||= 223 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 224 | when '7.0' 225 | ActiveRecord 226 | when '6.1' 227 | ActiveRecord::Base 228 | end 229 | end 230 | 231 | def shard_env 232 | ActiveRecordShards.app_env 233 | end 234 | 235 | # Make these few schema related methods available before having switched to 236 | # a shard. 237 | def with_default_shard(&block) 238 | if is_sharded? && current_shard_id.nil? && table_name != ActiveRecord::SchemaMigration.table_name 239 | on_first_shard(&block) 240 | else 241 | yield 242 | end 243 | end 244 | 245 | def load_schema_with_default_shard! 246 | with_default_shard { load_schema_without_default_shard! } 247 | end 248 | 249 | def table_exists_with_default_shard? 250 | with_default_shard { table_exists_without_default_shard? } 251 | end 252 | 253 | class PrimaryReplicaProxy 254 | def initialize(target, which) 255 | @target = target 256 | @which = which 257 | end 258 | 259 | def method_missing(method, *args, &block) # rubocop:disable Style/MethodMissingSuper, Style/MissingRespondToMissing 260 | @target.on_primary_or_replica(@which) { @target.send(method, *args, &block) } 261 | end 262 | ruby2_keywords(:method_missing) if respond_to?(:ruby2_keywords, true) 263 | end 264 | end 265 | end 266 | 267 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 268 | when '5.1', '5.2' 269 | require 'active_record_shards/connection_switcher-5-1' 270 | when '6.0' 271 | require 'active_record_shards/connection_switcher-6-0' 272 | when '6.1' 273 | require 'active_record_shards/connection_switcher-6-1' 274 | when '7.0' 275 | require 'active_record_shards/connection_switcher-7-0' 276 | else 277 | raise "ActiveRecordShards is not compatible with #{ActiveRecord::VERSION::STRING}" 278 | end 279 | -------------------------------------------------------------------------------- /test/thread_safety_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative 'helper' 4 | require_relative 'models' 5 | 6 | describe "connection switching thread safety" do 7 | with_fresh_databases 8 | 9 | before do 10 | ActiveRecord::Base.establish_connection(:test) 11 | use_same_connection_handler_for_all_theads 12 | create_seed_data 13 | end 14 | 15 | after do 16 | ActiveRecord::Base.connection_handler.clear_all_connections! 17 | end 18 | 19 | describe "isolation level" do 20 | it "uses the same connection for all fibers in a thread" do 21 | main_thread_connections = [] 22 | secondary_thread_connections = [] 23 | 24 | main_thread_connections << ActiveRecord::Base.connection.object_id 25 | Fiber.new { main_thread_connections << ActiveRecord::Base.connection.object_id }.resume 26 | 27 | Thread.new do 28 | secondary_thread_connections << ActiveRecord::Base.connection.object_id 29 | Fiber.new { secondary_thread_connections << ActiveRecord::Base.connection.object_id }.resume 30 | end.join 31 | 32 | assert_equal main_thread_connections.uniq.length, 1 33 | assert_equal secondary_thread_connections.uniq.length, 1 34 | refute_equal main_thread_connections.uniq, secondary_thread_connections.uniq 35 | end 36 | 37 | if ActiveRecord.version >= Gem::Version.new('7.0') 38 | it "raises an exception when isolation_level is not set to :thread" do 39 | assert_raises ActiveRecordShards::ConnectionSwitcher::IsolationLevelError do 40 | iso_level_was = ActiveSupport::IsolatedExecutionState.isolation_level 41 | ActiveSupport::IsolatedExecutionState.isolation_level = :fiber 42 | ActiveRecord::Base.on_primary_db do 43 | 1 44 | end 45 | ensure 46 | ActiveSupport::IsolatedExecutionState.isolation_level = iso_level_was 47 | end 48 | end 49 | end 50 | end 51 | 52 | it "can safely switch between all database connections in parallel" do 53 | new_thread("switches_through_all_1") do 54 | pause_and_mark_ready 55 | switch_through_all_databases 56 | end 57 | new_thread("switches_through_all_2") do 58 | pause_and_mark_ready 59 | switch_through_all_databases 60 | end 61 | new_thread("switches_through_all_3") do 62 | pause_and_mark_ready 63 | switch_through_all_databases 64 | end 65 | 66 | wait_for_threads_to_be_ready 67 | execute_and_wait_for_threads 68 | end 69 | 70 | describe "when multiple threads use different databases" do 71 | it "allows threads to parallelize their IO" do 72 | results = [] 73 | 74 | query_delay = { fast: "0.01", slow: "1", medium: "0.5" } 75 | new_thread("different_db_parallel_thread1") do 76 | ActiveRecord::Base.on_primary do 77 | pause_and_mark_ready 78 | result = execute_sql("SELECT name,'slower query',SLEEP(#{query_delay.fetch(:slow)}) FROM accounts") 79 | assert_equal('Primary account', result.first[0]) 80 | results.push(result) 81 | end 82 | end 83 | 84 | new_thread("different_db_parallel_thread2") do 85 | ActiveRecord::Base.on_replica do 86 | pause_and_mark_ready 87 | result = execute_sql("SELECT name, 'faster query',SLEEP(#{query_delay.fetch(:fast)}) FROM accounts") 88 | assert_equal('Replica account', result.first[0]) 89 | results.push(result) 90 | end 91 | end 92 | 93 | new_thread("different_db_parallel_thread3") do 94 | ActiveRecord::Base.on_shard(0) do 95 | pause_and_mark_ready 96 | result = execute_sql("SELECT title, 'medium query',SLEEP(#{query_delay.fetch(:medium)}) FROM tickets") 97 | assert_equal('Shard 0 Primary ticket', result.first[0]) 98 | results.push(result) 99 | end 100 | end 101 | 102 | wait_for_threads_to_be_ready 103 | 104 | thread_exection_time = Benchmark.realtime do 105 | execute_and_wait_for_threads 106 | end 107 | 108 | minimum_serial_query_exection_time = query_delay.values.map(&:to_f).sum 109 | # Arbitrarily faster time such that there must have been some parallelization 110 | max_parallel_time = minimum_serial_query_exection_time - 0.1 111 | assert_operator(max_parallel_time, :>, thread_exection_time) 112 | 113 | # This order cannot be guaranteed but it likely given the artificial delays 114 | rows = results.map(&:first) 115 | result_strings = rows.map { |r| r[1] } 116 | assert_equal( 117 | [ 118 | "faster query", 119 | "medium query", 120 | "slower query" 121 | ], 122 | result_strings 123 | ) 124 | end 125 | end 126 | 127 | describe "when multiple threads use the same database" do 128 | it "exposes a different connections to each thread" do 129 | connections = [] 130 | 131 | new_thread("connection_per_thread1") do 132 | ActiveRecord::Base.on_primary do 133 | pause_and_mark_ready 134 | connections << ActiveRecord::Base.connection 135 | end 136 | end 137 | 138 | new_thread("connection_per_thread2") do 139 | ActiveRecord::Base.on_primary do 140 | pause_and_mark_ready 141 | connections << ActiveRecord::Base.connection 142 | end 143 | end 144 | 145 | wait_for_threads_to_be_ready 146 | execute_and_wait_for_threads 147 | 148 | expect(connections.first).must_be_kind_of(ActiveRecord::ConnectionAdapters::Mysql2Adapter) 149 | assert_equal(2, connections.uniq.size) 150 | end 151 | 152 | it "allows threads to parallelize their IO" do 153 | results = [] 154 | 155 | query_delay = { fast: "0.01", slow: "1", medium: "0.5" } 156 | new_thread("same_db_parallel_thread1") do 157 | ActiveRecord::Base.on_primary do 158 | pause_and_mark_ready 159 | result = execute_sql("SELECT 'slower query',SLEEP(#{query_delay.fetch(:slow)})") 160 | results.push(result) 161 | end 162 | end 163 | 164 | new_thread("same_db_parallel_thread2") do 165 | ActiveRecord::Base.on_primary do 166 | pause_and_mark_ready 167 | result = execute_sql("SELECT 'faster query',SLEEP(#{query_delay.fetch(:fast)})") 168 | results.push(result) 169 | end 170 | end 171 | 172 | new_thread("same_db_parallel_thread3") do 173 | ActiveRecord::Base.on_primary do 174 | pause_and_mark_ready 175 | result = execute_sql("SELECT 'medium query',SLEEP(#{query_delay.fetch(:medium)})") 176 | results.push(result) 177 | end 178 | end 179 | 180 | wait_for_threads_to_be_ready 181 | 182 | thread_exection_time = Benchmark.realtime do 183 | execute_and_wait_for_threads 184 | end 185 | 186 | minimum_serial_query_exection_time = query_delay.values.map(&:to_f).sum 187 | # Arbitrarily faster time such that there must have been some parallelization 188 | max_parallel_time = minimum_serial_query_exection_time - 0.1 189 | assert_operator(max_parallel_time, :>, thread_exection_time) 190 | 191 | rows = results.map(&:first) 192 | result_strings = rows.map(&:first) 193 | # This order cannot be guaranteed but it likely given the artificial delays 194 | assert_equal( 195 | [ 196 | "faster query", 197 | "medium query", 198 | "slower query" 199 | ], 200 | result_strings 201 | ) 202 | end 203 | end 204 | 205 | def new_thread(name) 206 | thread = Thread.new do 207 | Thread.current.name = name 208 | yield 209 | end 210 | 211 | @test_threads ||= [] 212 | @test_threads.push(thread) 213 | end 214 | 215 | def switch_through_all_databases 216 | ActiveRecord::Base.on_primary do 217 | result = ActiveRecord::Base.connection.execute("SELECT * from accounts") 218 | assert_equal("Primary account", record_name(result)) 219 | end 220 | ActiveRecord::Base.on_replica do 221 | result = ActiveRecord::Base.connection.execute("SELECT * from accounts") 222 | assert_equal("Replica account", record_name(result)) 223 | end 224 | ActiveRecord::Base.on_shard(0) do 225 | result = ActiveRecord::Base.connection.execute("SELECT * from tickets") 226 | assert_equal("Shard 0 Primary ticket", record_name(result)) 227 | 228 | ActiveRecord::Base.on_replica do 229 | result = ActiveRecord::Base.connection.execute("SELECT * from tickets") 230 | assert_equal("Shard 0 Replica ticket", record_name(result)) 231 | end 232 | end 233 | ActiveRecord::Base.on_shard(1) do 234 | result = ActiveRecord::Base.connection.execute("SELECT * from tickets") 235 | assert_equal("Shard 1 Primary ticket", record_name(result)) 236 | 237 | ActiveRecord::Base.on_replica do 238 | result = ActiveRecord::Base.connection.execute("SELECT * from tickets") 239 | assert_equal("Shard 1 Replica ticket", record_name(result)) 240 | end 241 | end 242 | end 243 | 244 | # This allows us to get all of our threads into a prepared state by pausing 245 | # them at a 'ready' point so as there is as little overhead as possible 246 | # before the interesting code executes. 247 | # 248 | # Here we use 'ready' to mean the thread is spawned, has had its names set 249 | # and has established a database connection. 250 | def pause_and_mark_ready 251 | Thread.current[:ready] = true 252 | sleep 253 | end 254 | 255 | def execute_and_wait_for_threads 256 | @test_threads.each { |t| t.wakeup if t.alive? } 257 | @test_threads.each(&:join) 258 | end 259 | 260 | def wait_for_threads_to_be_ready 261 | sleep(0.01) until @test_threads.all? { |t| t[:ready] } 262 | end 263 | 264 | def use_same_connection_handler_for_all_theads 265 | ActiveRecord::Base.default_connection_handler = ActiveRecord::Base.connection_handler 266 | end 267 | 268 | def record_name(db_result) 269 | name_column_index = 1 270 | db_result.first[name_column_index] 271 | end 272 | 273 | def execute_sql(query) 274 | ActiveRecord::Base.connection.execute(query) 275 | end 276 | 277 | def create_seed_data 278 | ActiveRecord::Base.on_primary_db do 279 | Account.connection.execute(account_insert_sql(name: "Primary account")) 280 | 281 | Account.on_replica do 282 | Account.connection.execute(account_insert_sql(name: "Replica account")) 283 | end 284 | end 285 | 286 | [0, 1].each do |shard_id| 287 | ActiveRecord::Base.on_shard(shard_id) do 288 | Ticket.connection.execute(ticket_insert_sql(title: "Shard #{shard_id} Primary ticket")) 289 | 290 | Ticket.on_replica do 291 | Ticket.connection.execute(ticket_insert_sql(title: "Shard #{shard_id} Replica ticket")) 292 | end 293 | end 294 | end 295 | end 296 | 297 | def account_insert_sql(name:) 298 | "INSERT INTO accounts (id, name, created_at, updated_at)" \ 299 | " VALUES (1000, '#{name}', NOW(), NOW())" 300 | end 301 | 302 | def ticket_insert_sql(title:) 303 | "INSERT INTO tickets (id, title, account_id, created_at, updated_at)" \ 304 | " VALUES (1000, '#{title}', 5000, NOW(), NOW())" 305 | end 306 | end 307 | -------------------------------------------------------------------------------- /test/on_replica_by_default_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative 'helper' 4 | require 'models' 5 | 6 | describe ".on_replica_by_default" do 7 | with_fresh_databases 8 | 9 | before do 10 | ActiveRecord::Base.establish_connection(RAILS_ENV.to_sym) 11 | 12 | Account.on_replica_by_default = true 13 | Person.on_replica_by_default = true 14 | 15 | Account.connection.execute( 16 | "INSERT INTO accounts (id, name, created_at, updated_at) VALUES (1000, 'Primary account', NOW(), NOW())" 17 | ) 18 | Account.on_replica.connection.execute( 19 | "INSERT INTO accounts (id, name, created_at, updated_at) VALUES (1000, 'Replica account', NOW(), NOW())" 20 | ) 21 | Account.on_replica.connection.execute( 22 | "INSERT INTO accounts (id, name, created_at, updated_at) VALUES (1001, 'Replica account 2', NOW(), NOW())" 23 | ) 24 | 25 | Person.connection.execute( 26 | "REPLACE INTO people(id, name) VALUES (10, 'Primary person')" 27 | ) 28 | Person.on_replica.connection.execute( 29 | "REPLACE INTO people(id, name) VALUES (20, 'Replica person')" 30 | ) 31 | 32 | Account.connection.execute( 33 | "INSERT INTO account_people(account_id, person_id) VALUES (1000, 10)" 34 | ) 35 | Account.on_replica.connection.execute( 36 | "INSERT INTO account_people(account_id, person_id) VALUES (1001, 20)" 37 | ) 38 | end 39 | 40 | after do 41 | Account.on_replica_by_default = false 42 | Person.on_replica_by_default = false 43 | end 44 | 45 | describe "trying to access a primary DB connection" do 46 | before do 47 | clear_global_connection_handler_state 48 | end 49 | 50 | it "fails for the unsharded DB" do 51 | with_all_primaries_unavailable do 52 | assert_raises { Account.on_primary.connection } 53 | end 54 | end 55 | 56 | it "fails for the sharded DB" do 57 | with_all_primaries_unavailable do 58 | assert_raises { Ticket.on_primary.connection } 59 | end 60 | end 61 | end 62 | 63 | describe "on ActiveRecord::Base class" do 64 | it "reader is always false" do 65 | refute ActiveRecord::Base.on_replica_by_default? 66 | end 67 | 68 | it "setter does not work" do 69 | assert_raises ArgumentError do 70 | ActiveRecord::Base.on_replica_by_default = true 71 | end 72 | end 73 | end 74 | 75 | it "executes `find` on the replica" do 76 | with_all_primaries_unavailable do 77 | account = Account.find(1000) 78 | assert_equal "Replica account", account.name 79 | end 80 | end 81 | 82 | it "executes `find_by` on the replica" do 83 | with_all_primaries_unavailable do 84 | account = Account.find_by(id: 1000) 85 | assert_equal "Replica account", account.name 86 | end 87 | end 88 | 89 | it "executes `count` on the replica" do 90 | with_all_primaries_unavailable do 91 | count = Account.count 92 | assert_equal 2, count 93 | end 94 | end 95 | 96 | it "executes `reload` on the replica" do 97 | with_all_primaries_unavailable do 98 | account = Account.find(1000) 99 | assert_equal "Replica account", account.reload.name 100 | end 101 | end 102 | 103 | it "executes `exists?` on the replica" do 104 | with_all_primaries_unavailable do 105 | assert Account.exists?(1001) 106 | end 107 | end 108 | 109 | it "executes `exists?` on the replica with a named scope" do 110 | AccountThing.on_replica_by_default = true 111 | AccountThing.on_replica.connection.execute("INSERT INTO account_things (id, account_id) VALUES (123125, 1000)") 112 | 113 | with_all_primaries_unavailable do 114 | assert AccountThing.enabled.exists?(123125) 115 | end 116 | 117 | AccountThing.on_replica_by_default = false 118 | end 119 | 120 | it "counts associations on the replica" do 121 | AccountThing.on_replica_by_default = true 122 | AccountThing.on_replica.connection.execute("INSERT INTO account_things (id, account_id) VALUES (123123, 1000)") 123 | AccountThing.on_replica.connection.execute("INSERT INTO account_things (id, account_id) VALUES (123124, 1000)") 124 | 125 | with_all_primaries_unavailable do 126 | assert_equal 2, Account.find(1000).account_things.count 127 | end 128 | 129 | AccountThing.on_replica_by_default = false 130 | end 131 | 132 | it "`includes` things via has_and_belongs_to_many associations correctly" do 133 | with_all_primaries_unavailable do 134 | a = Account.where(id: 1001).includes(:people).first 135 | refute_empty(a.people) 136 | assert_equal "Replica person", a.people.first.name 137 | end 138 | end 139 | 140 | it "sets up has_and_belongs_to_many sharded-ness correctly" do 141 | with_all_primaries_unavailable do 142 | refute Account.const_get(:HABTM_People).is_sharded? 143 | end 144 | end 145 | 146 | it "executes `pluck` on the replica" do 147 | with_all_primaries_unavailable do 148 | assert_equal ["Replica account", "Replica account 2"], Account.pluck(:name) 149 | end 150 | end 151 | 152 | it "executes `first` on association on the replica" do 153 | with_all_primaries_unavailable do 154 | account = Account.find(1001) 155 | person = account.people.first 156 | assert_equal "Replica person", person.name 157 | end 158 | end 159 | 160 | it "executes `map` on preloaded relation on the primary" do 161 | Ticket.on_shard(1) do 162 | Ticket.connection.execute( 163 | "INSERT INTO tickets (id, title, account_id, created_at, updated_at) VALUES (50000, 'Primary ticket', 1001, NOW(), NOW())" 164 | ) 165 | Ticket.on_replica.connection.execute( 166 | "INSERT INTO tickets (id, title, account_id, created_at, updated_at) VALUES (50001, 'Replica ticket', 1001, NOW(), NOW())" 167 | ) 168 | 169 | with_unsharded_primary_unavailable do 170 | ticket_rel = Ticket.preload(:account).where(id: 50000) 171 | ticket_titles = ticket_rel.map(&:title) 172 | assert_equal ["Primary ticket"], ticket_titles 173 | end 174 | end 175 | end 176 | 177 | it "executes `all` on association on the replica" do 178 | with_all_primaries_unavailable do 179 | account = Account.find(1001) 180 | all_people = account.people.all 181 | assert_equal ["Replica person"], all_people.map(&:name) 182 | end 183 | end 184 | 185 | it "executes `count` on association on the replica" do 186 | Person.on_replica.connection.execute( 187 | "INSERT INTO people(id, name) VALUES (30, 'Replica person 2')" 188 | ) 189 | Account.on_replica.connection.execute( 190 | "INSERT INTO account_people(account_id, person_id) VALUES (1001, 30)" 191 | ) 192 | 193 | with_all_primaries_unavailable do 194 | account = Account.find(1001) 195 | count = account.people.count 196 | assert_equal 2, count 197 | end 198 | end 199 | 200 | it "can call preload from sharded model to unsharded model" do 201 | Ticket.on_shard(1) do 202 | Ticket.connection.execute( 203 | "INSERT INTO tickets (id, title, account_id, created_at, updated_at) VALUES (50000, 'Primary ticket', 1000, NOW(), NOW())" 204 | ) 205 | Ticket.on_replica.connection.execute( 206 | "INSERT INTO tickets (id, title, account_id, created_at, updated_at) VALUES (50001, 'Replica ticket', 1001, NOW(), NOW())" 207 | ) 208 | end 209 | 210 | begin 211 | Ticket.on_replica_by_default = true 212 | 213 | Ticket.on_shard(1) do 214 | with_all_primaries_unavailable do 215 | tickets = Ticket.preload(:account) 216 | ticket = tickets.first 217 | 218 | assert_equal "Replica ticket", ticket.title 219 | assert_equal "Replica account 2", ticket.account.name 220 | end 221 | end 222 | ensure 223 | Ticket.on_replica_by_default = false 224 | end 225 | end 226 | 227 | it "can handle association from sharded model to unsharded model" do 228 | Ticket.on_shard(1) do 229 | Ticket.connection.execute( 230 | "INSERT INTO tickets (id, title, account_id, created_at, updated_at) VALUES (50000, 'Primary ticket', 1000, NOW(), NOW())" 231 | ) 232 | Ticket.on_replica.connection.execute( 233 | "INSERT INTO tickets (id, title, account_id, created_at, updated_at) VALUES (50001, 'Replica ticket', 1001, NOW(), NOW())" 234 | ) 235 | end 236 | 237 | begin 238 | Ticket.on_replica_by_default = true 239 | 240 | Ticket.on_shard(1) do 241 | with_all_primaries_unavailable do 242 | ticket = Ticket.find(50001) 243 | account_name = ticket.account.name 244 | assert_equal "Replica account 2", account_name 245 | end 246 | end 247 | ensure 248 | Ticket.on_replica_by_default = false 249 | end 250 | end 251 | 252 | it "can instantiate a new record whose model defines an ordered default_scope" do 253 | with_all_primaries_unavailable do 254 | User.new 255 | end 256 | end 257 | 258 | it "loads schema from replica" do 259 | Account.reset_column_information 260 | 261 | # Verify that the schema hasn't been loaded yet 262 | assert_nil Account.instance_variable_get :@columns 263 | assert_nil Account.instance_variable_get :@columns_hash 264 | assert_nil Account.instance_variable_get :@column_names 265 | 266 | with_all_primaries_unavailable do 267 | assert_equal ["id", "name", "created_at", "updated_at"], Account.column_names 268 | end 269 | end 270 | 271 | it "loads primary key column from replica" do 272 | with_all_primaries_unavailable do 273 | Account.reset_primary_key 274 | 275 | assert_equal "id", Account.primary_key 276 | end 277 | end 278 | 279 | describe "joins" do 280 | it "supports implicit joins" do 281 | with_all_primaries_unavailable do 282 | accounts = Account.includes(:account_things).references(:account_things) 283 | account_names = accounts.order("account_things.id").map(&:name).sort 284 | assert_equal ["Replica account", "Replica account 2"], account_names 285 | end 286 | end 287 | 288 | it "supports explicit joins" do 289 | with_all_primaries_unavailable do 290 | accounts = Account.joins("LEFT OUTER JOIN account_things ON account_things.account_id = accounts.id") 291 | account_names = accounts.map(&:name).sort 292 | assert_equal ["Replica account", "Replica account 2"], account_names 293 | end 294 | end 295 | 296 | it "does not support implicit joins between an unsharded and a sharded table" do 297 | with_all_primaries_unavailable do 298 | accounts = Account.includes(:tickets).references(:tickets).order("tickets.id") 299 | assert_raises(ActiveRecord::StatementInvalid) { accounts.first } 300 | end 301 | end 302 | 303 | it "does not support explicit joins between an unsharded and a sharded table" do 304 | with_all_primaries_unavailable do 305 | accounts = Account.joins("LEFT OUTER JOIN tickets ON tickets.account_id = accounts.id") 306 | assert_raises(ActiveRecord::StatementInvalid) { accounts.first } 307 | end 308 | end 309 | end 310 | 311 | describe "overriding with `on_primary`" do 312 | it "allows overriding with `on_primary`" do 313 | model = Account.on_primary.find(1000) 314 | assert_equal "Primary account", model.name 315 | end 316 | 317 | it "does not allow overriding `on_primary` with `on_replica`" do 318 | model = Account.on_primary { Account.on_replica.find(1000) } 319 | assert_equal "Primary account", model.name 320 | end 321 | 322 | it "allows overriding `on_replica` with `on_primary`" do 323 | model = Account.on_replica { Account.on_primary.find(1000) } 324 | assert_equal "Primary account", model.name 325 | end 326 | end 327 | 328 | describe "inheritance" do 329 | it "propagates the `on_replica_by_default?` reader to inherited classes" do 330 | assert AccountInherited.on_replica_by_default? 331 | end 332 | 333 | it "propagates the `on_replica_by_default` writer to inherited classes" do 334 | AccountInherited.on_replica_by_default = false 335 | refute AccountInherited.on_replica_by_default? 336 | refute Account.on_replica_by_default? 337 | ensure 338 | AccountInherited.on_replica_by_default = true 339 | end 340 | end 341 | 342 | describe 'in transactions' do 343 | it 'performs reads on the primary in an #after_save hook' do 344 | model = Account.on_primary.find(1000) 345 | model.name = 'bartfoo' 346 | model.singleton_class.after_save do 347 | @fetched_person = Person.find(10) 348 | end 349 | model.save! 350 | assert_equal "Primary person", model.instance_variable_get(:@fetched_person).name 351 | end 352 | 353 | it 'stays on Primary when calling .transaction on AR::Base' do 354 | ActiveRecord::Base.transaction do 355 | assert_equal Account.find(1000).name, 'Primary account' 356 | assert ActiveRecord::Base._in_transaction? 357 | end 358 | 359 | refute ActiveRecord::Base._in_transaction? 360 | end 361 | 362 | it 'stays on Primary when calling #transaction on connection' do 363 | ActiveRecord::Base.connection.transaction do 364 | assert_equal Account.find(1000).name, 'Primary account' 365 | assert ActiveRecord::Base._in_transaction? 366 | end 367 | 368 | refute ActiveRecord::Base._in_transaction? 369 | end 370 | 371 | it 'stays on Primary when calling .transaction on a model' do 372 | Account.transaction do 373 | assert_equal Account.find(1000).name, 'Primary account' 374 | assert Account._in_transaction? 375 | end 376 | 377 | refute ActiveRecord::Base._in_transaction? 378 | end 379 | end 380 | end 381 | -------------------------------------------------------------------------------- /test/connection_switching_test.rb: -------------------------------------------------------------------------------- 1 | # frozen_string_literal: true 2 | 3 | require_relative 'helper' 4 | 5 | describe "connection switching" do 6 | include ConnectionSwitchingSpecHelpers 7 | extend RailsEnvSwitch 8 | 9 | def clear_connection_pool 10 | ActiveRecord::Base.connection_handler.connection_pool_list.clear 11 | end 12 | 13 | with_fresh_databases 14 | 15 | before do 16 | ActiveRecord::Base.establish_connection(RAILS_ENV.to_sym) 17 | end 18 | 19 | describe "legacy_connection_handling" do 20 | if ActiveRecord.version >= Gem::Version.new('6.1') 21 | it "raises an exception when legacy_connection_handling is false" do 22 | legacy_connection_handling_owner ||= 23 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 24 | when '7.0' 25 | ActiveRecord 26 | when '6.1' 27 | ActiveRecord::Base 28 | end 29 | 30 | begin 31 | legacy_connection_handling_owner.legacy_connection_handling = false 32 | 33 | assert_raises ActiveRecordShards::ConnectionSwitcher::LegacyConnectionHandlingError do 34 | ActiveRecord::Base.on_primary_db do 35 | 1 36 | end 37 | end 38 | 39 | assert_raises ActiveRecordShards::ConnectionSwitcher::LegacyConnectionHandlingError do 40 | ActiveRecord::Base.on_shard(1) do 41 | 1 42 | end 43 | end 44 | ensure 45 | legacy_connection_handling_owner.legacy_connection_handling = true 46 | end 47 | end 48 | end 49 | end 50 | 51 | describe "on_primary_db" do 52 | after do 53 | ActiveRecord::Base.default_shard = nil 54 | end 55 | 56 | it "switches to the primary database" do 57 | ActiveRecord::Base.on_primary_db do 58 | assert_using_database('ars_test') 59 | end 60 | end 61 | 62 | it "switches to the primary database and back when a default shard is set" do 63 | ActiveRecord::Base.default_shard = 0 64 | assert_using_database('ars_test_shard0') 65 | 66 | ActiveRecord::Base.on_primary_db do 67 | assert_using_database('ars_test') 68 | end 69 | 70 | assert_using_database('ars_test_shard0') 71 | end 72 | 73 | it "switches to the primary datatbase and back when nested of an on_shard block" do 74 | ActiveRecord::Base.on_shard(1) do 75 | assert_using_database('ars_test_shard1') 76 | 77 | ActiveRecord::Base.on_primary_db do 78 | assert_using_database('ars_test') 79 | end 80 | 81 | assert_using_database('ars_test_shard1') 82 | end 83 | end 84 | end 85 | 86 | describe "shard switching" do 87 | it "only switch connection on sharded models" do 88 | assert_using_database('ars_test', Ticket) 89 | assert_using_database('ars_test', Account) 90 | 91 | ActiveRecord::Base.on_shard(0) do 92 | assert_using_database('ars_test_shard0', Ticket) 93 | assert_using_database('ars_test', Account) 94 | end 95 | end 96 | 97 | it "switch to shard and back" do 98 | assert_using_database('ars_test') 99 | ActiveRecord::Base.on_replica { assert_using_database('ars_test_replica') } 100 | 101 | ActiveRecord::Base.on_shard(0) do 102 | assert_using_database('ars_test_shard0') 103 | ActiveRecord::Base.on_replica { assert_using_database('ars_test_shard0_replica') } 104 | 105 | ActiveRecord::Base.on_shard(nil) do 106 | assert_using_database('ars_test') 107 | ActiveRecord::Base.on_replica { assert_using_database('ars_test_replica') } 108 | end 109 | 110 | assert_using_database('ars_test_shard0') 111 | ActiveRecord::Base.on_replica { assert_using_database('ars_test_shard0_replica') } 112 | end 113 | 114 | assert_using_database('ars_test') 115 | ActiveRecord::Base.on_replica { assert_using_database('ars_test_replica') } 116 | end 117 | 118 | it "does not fail on unrelated ensure error when current_shard_selection fails" do 119 | ActiveRecord::Base.stub(:current_shard_selection, -> { raise ArgumentError }) do 120 | assert_raises ArgumentError do 121 | ActiveRecord::Base.on_replica { 1 } 122 | end 123 | end 124 | end 125 | 126 | describe "on_first_shard" do 127 | it "use the first shard" do 128 | ActiveRecord::Base.on_first_shard do 129 | assert_using_database('ars_test_shard0') 130 | end 131 | end 132 | end 133 | 134 | describe "on_all_shards" do 135 | before do 136 | @shard_0_primary = ActiveRecord::Base.on_shard(0) { ActiveRecord::Base.connection } 137 | @shard_1_primary = ActiveRecord::Base.on_shard(1) { ActiveRecord::Base.connection } 138 | refute_equal(@shard_0_primary.select_value("SELECT DATABASE()"), @shard_1_primary.select_value("SELECT DATABASE()")) 139 | end 140 | 141 | it "execute the block on all shard primaries" do 142 | result = ActiveRecord::Base.on_all_shards do |shard| 143 | [ActiveRecord::Base.connection.select_value("SELECT DATABASE()"), shard] 144 | end 145 | database_names = result.map(&:first) 146 | database_shards = result.map(&:last) 147 | 148 | assert_equal(2, database_names.size) 149 | assert_includes(database_names, @shard_0_primary.select_value("SELECT DATABASE()")) 150 | assert_includes(database_names, @shard_1_primary.select_value("SELECT DATABASE()")) 151 | 152 | assert_equal(2, database_shards.size) 153 | assert_includes(database_shards, 0) 154 | assert_includes(database_shards, 1) 155 | end 156 | 157 | it "execute the block unsharded" do 158 | result = ActiveRecord::Base.stub(:supports_sharding?, false) do 159 | ActiveRecord::Base.on_all_shards do |shard| 160 | [ActiveRecord::Base.connection.select_value("SELECT DATABASE()"), shard] 161 | end 162 | end 163 | assert_equal [["ars_test", nil]], result 164 | end 165 | end 166 | 167 | describe ".shards.enum" do 168 | it "works like this:" do 169 | ActiveRecord::Base.on_all_shards { Ticket.create! } 170 | count = ActiveRecord::Base.shards.enum.map { Ticket.count }.inject(&:+) 171 | assert_equal(ActiveRecord::Base.shard_names.size, count) 172 | end 173 | end 174 | 175 | describe ".shards.find" do 176 | it "works like this:" do 177 | t = ActiveRecord::Base.on_shard(1) { Ticket.create! } 178 | assert(Ticket.shards.find(t.id)) 179 | 180 | assert_raises(ActiveRecord::RecordNotFound) do 181 | Ticket.shards.find(123123123) 182 | end 183 | end 184 | end 185 | 186 | describe ".shards.count" do 187 | before do 188 | ActiveRecord::Base.on_shard(0) { 2.times { Ticket.create!(title: "0") } } 189 | ActiveRecord::Base.on_shard(1) { Ticket.create!(title: "1") } 190 | end 191 | 192 | it "works like this:" do 193 | count = Ticket.shards.count 194 | assert_equal(3, count) 195 | end 196 | 197 | it "works with scopes" do 198 | assert_equal(1, Ticket.where(title: "1").shards.count) 199 | end 200 | end 201 | 202 | describe ".shards.to_a" do 203 | it "works like this" do 204 | ActiveRecord::Base.on_all_shards { |s| Ticket.create!(title: s.to_s) } 205 | 206 | res = Ticket.where(title: "1").shards.to_a 207 | assert_equal 1, res.size 208 | end 209 | end 210 | 211 | describe "when loading primary keys" do 212 | it "uses the default shard" do 213 | refute_nil Ticket.primary_key 214 | end 215 | end 216 | end 217 | 218 | describe "fibers" do 219 | it "switches to the correct shard, even when changing fibers" do 220 | # Rails utilizes Object#to_enum in a few places such as #find_in_batches. 221 | # to_enum, in the C code, is creating a new _Fiber_ and executing the block 222 | # in the context of that Fiber. We want to ensure that, when we call on_shard 223 | # that we remain on the shard we intend, even if the fiber changes. 224 | ActiveRecord::Base.on_shard(0) do 225 | Fiber.new { assert_using_database('ars_test_shard0', Ticket) }.resume 226 | end 227 | 228 | assert_using_primary_db 229 | end 230 | end 231 | 232 | describe "default shard selection" do 233 | describe "of nil" do 234 | before do 235 | ActiveRecord::Base.default_shard = nil 236 | end 237 | 238 | it "use unsharded db for sharded models" do 239 | assert_using_database('ars_test', Ticket) 240 | assert_using_database('ars_test', Account) 241 | end 242 | end 243 | 244 | describe "value" do 245 | before do 246 | ActiveRecord::Base.default_shard = 0 247 | end 248 | 249 | after do 250 | ActiveRecord::Base.default_shard = nil 251 | end 252 | 253 | it "use default shard db for sharded models" do 254 | assert_using_database('ars_test_shard0', Ticket) 255 | assert_using_database('ars_test', Account) 256 | end 257 | 258 | it "still be able to switch to shard nil" do 259 | ActiveRecord::Base.on_shard(nil) do 260 | assert_using_database('ars_test', Ticket) 261 | assert_using_database('ars_test', Account) 262 | end 263 | end 264 | end 265 | end 266 | 267 | describe "ActiveRecord::Base.columns" do 268 | before do 269 | ActiveRecord::Base.default_shard = nil 270 | end 271 | 272 | describe "for unsharded models" do 273 | before do 274 | Account.on_replica do 275 | Account.connection.execute("alter table accounts add column foo int") 276 | Account.reset_column_information 277 | end 278 | end 279 | 280 | after do 281 | Account.on_replica do 282 | ActiveRecord::Base.connection.execute("alter table accounts drop column foo") 283 | Account.reset_column_information 284 | refute_includes(Account.column_names, 'foo') 285 | end 286 | end 287 | 288 | it "use the non-sharded replica connection" do 289 | assert_using_database('ars_test', Account) 290 | assert_includes(Account.column_names, 'foo') 291 | end 292 | 293 | it "ignores primary/transactions" do 294 | assert_using_database('ars_test', Account) 295 | Account.on_primary { assert_includes(Account.column_names, 'foo') } 296 | end 297 | end 298 | 299 | describe "for sharded models" do 300 | before do 301 | ActiveRecord::Base.on_first_shard do 302 | ActiveRecord::Base.on_replica do 303 | ActiveRecord::Base.connection.execute("alter table tickets add column foo int") 304 | Ticket.reset_column_information 305 | end 306 | end 307 | end 308 | 309 | after do 310 | ActiveRecord::Base.on_first_shard do 311 | ActiveRecord::Base.on_replica do 312 | ActiveRecord::Base.connection.execute("alter table tickets drop column foo") 313 | Ticket.reset_column_information 314 | end 315 | end 316 | end 317 | 318 | it "gets columns from the replica shard" do 319 | assert_includes(Ticket.column_names, 'foo') 320 | end 321 | 322 | it "have correct from_shard" do 323 | ActiveRecord::Base.on_all_shards do |shard| 324 | assert_equal shard, Ticket.new.from_shard 325 | end 326 | end 327 | end 328 | 329 | describe "for SchemaMigration" do 330 | before do 331 | ActiveRecord::Base.on_shard(nil) do 332 | ActiveRecord::Base.connection.execute("alter table schema_migrations add column foo int") 333 | end 334 | end 335 | 336 | after do 337 | ActiveRecord::Base.on_shard(nil) do 338 | ActiveRecord::Base.connection.execute("alter table schema_migrations drop column foo") 339 | end 340 | end 341 | 342 | it "doesn't switch to shard" do 343 | table_has_column?('schema_migrations', 'foo') 344 | end 345 | end 346 | end 347 | 348 | describe ".where.to_sql" do 349 | it "doesn't use the primary (for escaping)" do 350 | with_unsharded_primary_unavailable do 351 | Account.all.to_sql 352 | Account.where('id = 1').to_sql 353 | Account.where('id = ?', 1).to_sql 354 | Account.where(id: 1).to_sql 355 | end 356 | end 357 | end 358 | 359 | describe "ActiveRecord::Base.table_exists?" do 360 | before do 361 | ActiveRecord::Base.default_shard = nil 362 | end 363 | 364 | describe "for unsharded models" do 365 | it "use the unsharded replica connection" do 366 | class UnshardedModel < ActiveRecord::Base 367 | not_sharded 368 | end 369 | 370 | UnshardedModel.on_replica { UnshardedModel.connection.execute("create table unsharded_models (id int)") } 371 | assert UnshardedModel.table_exists? 372 | 373 | ActiveRecord::Base.on_all_shards do 374 | refute table_exists?("unsharded_models") 375 | end 376 | end 377 | end 378 | 379 | describe "for sharded models" do 380 | it "uses the first shard replica" do 381 | class ShardedModel < ActiveRecord::Base 382 | end 383 | 384 | ActiveRecord::Base.on_first_shard do 385 | ShardedModel.on_replica do 386 | ShardedModel.connection.execute("create table sharded_models (id int)") 387 | end 388 | end 389 | 390 | assert ShardedModel.table_exists? 391 | end 392 | end 393 | end 394 | 395 | describe "in an environment without replica" do 396 | switch_app_env('test3') 397 | def spec_name 398 | case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" 399 | when '6.1', '7.0' 400 | ActiveRecord::Base.connection_pool.db_config.name 401 | else 402 | ActiveRecord::Base.connection_pool.spec.name 403 | end 404 | end 405 | 406 | describe "shard switching" do 407 | it "just stay on the primary db" do 408 | main_spec_name = spec_name 409 | shard_spec_name = ActiveRecord::Base.on_shard(0) { spec_name } 410 | 411 | ActiveRecord::Base.on_replica do 412 | assert_using_database('ars_test3', Account) 413 | assert_equal main_spec_name, spec_name 414 | 415 | ActiveRecord::Base.on_shard(0) do 416 | assert_using_database('ars_test3_shard0', Ticket) 417 | assert_equal shard_spec_name, spec_name 418 | end 419 | end 420 | end 421 | end 422 | end 423 | 424 | describe "in an unsharded environment" do 425 | switch_app_env('test2') 426 | 427 | describe "shard switching" do 428 | it "just stay on the main db" do 429 | assert_using_database('ars_test2', Ticket) 430 | assert_using_database('ars_test2', Account) 431 | 432 | ActiveRecord::Base.on_shard(0) do 433 | assert_using_database('ars_test2', Ticket) 434 | assert_using_database('ars_test2', Account) 435 | end 436 | end 437 | end 438 | 439 | describe "on_all_shards" do 440 | before do 441 | @database_names = [] 442 | ActiveRecord::Base.on_all_shards do 443 | @database_names << ActiveRecord::Base.connection.select_value("SELECT DATABASE()") 444 | end 445 | end 446 | 447 | it "execute the block on all shard primaries" do 448 | assert_equal([ActiveRecord::Base.connection.select_value("SELECT DATABASE()")], @database_names) 449 | end 450 | end 451 | end 452 | 453 | describe "replica driving" do 454 | describe "without replica configuration" do 455 | before do 456 | if ActiveRecord::VERSION::MAJOR >= 6 457 | @saved_config = ActiveRecord::Base.configurations.find_db_config('test_replica') 458 | ActiveRecord::Base.configurations.configurations.delete(@saved_config) 459 | else 460 | @saved_config = ActiveRecord::Base.configurations.delete('test_replica') 461 | end 462 | Thread.current._active_record_shards_shard_selection = nil # drop caches 463 | clear_connection_pool 464 | ActiveRecord::Base.establish_connection(:test) 465 | end 466 | 467 | after do 468 | if ActiveRecord::VERSION::MAJOR >= 6 469 | ActiveRecord::Base.configurations.configurations << @saved_config 470 | else 471 | ActiveRecord::Base.configurations['test_replica'] = @saved_config 472 | end 473 | Thread.current._active_record_shards_shard_selection = nil # drop caches 474 | end 475 | 476 | it "default to the primary database" do 477 | Account.create! 478 | 479 | ActiveRecord::Base.on_replica { assert_using_primary_db } 480 | Account.on_replica { assert_using_primary_db } 481 | Ticket.on_replica { assert_using_primary_db } 482 | end 483 | 484 | it "successfully execute queries" do 485 | Account.create! 486 | assert_using_primary_db 487 | 488 | assert_equal(Account.count, ActiveRecord::Base.on_replica { Account.count }) 489 | assert_equal(Account.count, Account.on_replica { Account.count }) 490 | end 491 | end 492 | 493 | describe "with replica configuration" do 494 | it "successfully execute queries" do 495 | assert_using_primary_db 496 | Account.create! 497 | 498 | assert_equal(1, Account.count) 499 | assert_equal(0, ActiveRecord::Base.on_replica { Account.count }) 500 | end 501 | 502 | it "support global on_replica blocks" do 503 | assert_using_primary_db 504 | assert_using_primary_db 505 | 506 | ActiveRecord::Base.on_replica do 507 | assert_using_replica_db 508 | assert_using_replica_db 509 | end 510 | 511 | assert_using_primary_db 512 | assert_using_primary_db 513 | end 514 | 515 | it "support conditional methods" do 516 | assert_using_primary_db 517 | 518 | Account.on_replica_if(true) do 519 | assert_using_replica_db 520 | end 521 | 522 | assert_using_primary_db 523 | 524 | Account.on_replica_if(false) do 525 | assert_using_primary_db 526 | end 527 | 528 | Account.on_replica_unless(true) do 529 | assert_using_primary_db 530 | end 531 | 532 | Account.on_replica_unless(false) do 533 | assert_using_replica_db 534 | end 535 | end 536 | 537 | describe "a model loaded with the replica" do 538 | before do 539 | Account.on_replica_by_default = true 540 | 541 | Account.on_primary.connection.execute("INSERT INTO accounts (id, name, created_at, updated_at) VALUES(1000, 'primary_name', '2009-12-04 20:18:48', '2009-12-04 20:18:48')") 542 | assert(Account.on_primary.find(1000)) 543 | assert_equal('primary_name', Account.on_primary.find(1000).name) 544 | 545 | Account.on_replica.connection.execute("INSERT INTO accounts (id, name, created_at, updated_at) VALUES(1000, 'replica_name', '2009-12-04 20:18:48', '2009-12-04 20:18:48')") 546 | 547 | @model = Account.on_replica.find(1000) 548 | assert(@model) 549 | assert_equal('replica_name', @model.name) 550 | end 551 | 552 | it "read from replica on reload" do 553 | @model.reload 554 | assert_equal('replica_name', @model.name) 555 | end 556 | 557 | it "be marked as read-only" do 558 | assert(@model.readonly?) 559 | end 560 | 561 | it "be marked as comming from the replica" do 562 | assert(@model.from_replica?) 563 | end 564 | 565 | describe "when ActiveRecordShards.disable_replica_readonly_records is true" do 566 | before do 567 | @original_disable_replica_readonly_records = ActiveRecordShards.disable_replica_readonly_records 568 | ActiveRecordShards.disable_replica_readonly_records = true 569 | 570 | @model = Account.on_replica.find(1000) 571 | assert(@model) 572 | assert_equal('replica_name', @model.name) 573 | end 574 | 575 | it "read from replica on reload" do 576 | @model.reload 577 | assert_equal('replica_name', @model.name) 578 | end 579 | 580 | it "not be marked as read-only" do 581 | refute(@model.readonly?) 582 | end 583 | 584 | it "be marked as comming from the replica" do 585 | assert(@model.from_replica?) 586 | end 587 | 588 | after do 589 | ActiveRecordShards.disable_replica_readonly_records = @original_disable_replica_readonly_records 590 | end 591 | end 592 | 593 | after do 594 | Account.on_replica_by_default = false 595 | end 596 | end 597 | 598 | describe "a inherited model without cached columns hash" do 599 | # before columns -> with_scope -> type-condition -> columns == loop 600 | it "not loop when on replica by default" do 601 | Person.on_replica_by_default = true 602 | assert User.on_replica_by_default? 603 | assert User.finder_needs_type_condition? 604 | 605 | User.reset_column_information 606 | User.columns_hash 607 | ensure 608 | Person.on_replica_by_default = false 609 | end 610 | end 611 | 612 | describe "a model loaded with the primary" do 613 | before do 614 | Account.connection.execute("INSERT INTO accounts (id, name, created_at, updated_at) VALUES(1000, 'primary_name', '2009-12-04 20:18:48', '2009-12-04 20:18:48')") 615 | @model = Account.first 616 | assert(@model) 617 | assert_equal('primary_name', @model.name) 618 | end 619 | 620 | it "not unset readonly" do 621 | @model = Account.on_primary.readonly.first 622 | assert(@model.readonly?) 623 | end 624 | 625 | it "not be marked as read-only" do 626 | assert(!@model.readonly?) 627 | end 628 | 629 | it "not be marked as comming from the replica" do 630 | assert(!@model.from_replica?) 631 | end 632 | end 633 | end 634 | 635 | describe "replica proxy" do 636 | it "successfully execute queries" do 637 | assert_using_primary_db 638 | Account.create! 639 | 640 | refute_equal Account.count, Account.on_replica.count 641 | end 642 | 643 | it "work on association collections" do 644 | assert_using_primary_db 645 | account = Account.create! 646 | 647 | ActiveRecord::Base.on_shard(0) do 648 | account.tickets.create! title: 'primary ticket' 649 | 650 | Ticket.on_replica do 651 | account.tickets.create! title: 'replica ticket' 652 | end 653 | 654 | assert_equal "primary ticket", account.tickets.first.title 655 | assert_equal "replica ticket", account.tickets.on_replica.first.title 656 | end 657 | end 658 | end 659 | end 660 | end 661 | --------------------------------------------------------------------------------