├── .document ├── .github ├── CODEOWNERS ├── term-check.yaml └── workflows │ ├── ci.yml │ └── publish.yml ├── .gitignore ├── .rubocop.yml ├── Changelog.md ├── Gemfile ├── LICENSE ├── README.md ├── Rakefile ├── active_record_shards.gemspec ├── gemfiles ├── common.rb ├── rails5.1.gemfile ├── rails5.2.gemfile ├── rails6.0.gemfile ├── rails6.1.gemfile └── rails7.0.gemfile ├── lib ├── active_record_shards.rb └── active_record_shards │ ├── association_collection_connection_selection.rb │ ├── configuration_parser.rb │ ├── connection_switcher-5-1.rb │ ├── connection_switcher-6-0.rb │ ├── connection_switcher-6-1.rb │ ├── connection_switcher-7-0.rb │ ├── connection_switcher.rb │ ├── default_replica_patches.rb │ ├── default_shard.rb │ ├── migration.rb │ ├── model.rb │ ├── schema_dumper_extension.rb │ ├── shard_selection.rb │ ├── shard_support.rb │ ├── sql_comments.rb │ └── tasks.rb └── test ├── active_record └── connection_adapters │ └── mysql2_fake_adapter.rb ├── active_record_shards ├── configuration_parser_test.rb ├── schema_dumper_extension_test.rb └── sql_comments_test.rb ├── active_record_shards_test.rb ├── check_performance.rb ├── connection_switching_test.rb ├── database.yml ├── helper.rb ├── migrator_test.rb ├── models.rb ├── on_replica_by_default_test.rb ├── schemas ├── ars_test.sql ├── ars_test2.sql ├── ars_test2_replica.sql ├── ars_test3.sql ├── ars_test3_shard0.sql ├── ars_test_replica.sql ├── ars_test_shard0.sql ├── ars_test_shard0_replica.sql ├── ars_test_shard1.sql └── ars_test_shard1_replica.sql ├── support ├── cowardly_migration │ └── 20110824010215_cowardly_migration.rb ├── database_tasks.yml ├── db_helper.rb ├── failure_migration │ └── 20110824010215_failure_migration.rb ├── migrations │ ├── 20110824010216_shard_migration.rb │ └── 20110829215912_account_migration.rb ├── separate_migrations │ ├── 20190121112233_separate_unsharded_migration.rb │ └── 20190121112234_separate_sharded_migration.rb ├── sharded_schema.rb ├── tcp_proxy.rb └── unsharded_schema.rb ├── tasks_test.rb └── thread_safety_test.rb /.document: -------------------------------------------------------------------------------- 1 | README.rdoc 2 | lib/**/*.rb 3 | bin/* 4 | features/**/*.feature 5 | LICENSE 6 | -------------------------------------------------------------------------------- /.github/CODEOWNERS: -------------------------------------------------------------------------------- 1 | * @zendesk/ruby-core @zendesk/database-gem-owners 2 | -------------------------------------------------------------------------------- /.github/term-check.yaml: -------------------------------------------------------------------------------- 1 | ignore: 2 | - Changelog.md 3 | -------------------------------------------------------------------------------- /.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 | -------------------------------------------------------------------------------- /.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 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | pkg/* 2 | *.gem 3 | .bundle 4 | *.log 5 | Gemfile.lock 6 | gemfiles/*.lock 7 | .tags* 8 | tmp/* 9 | -------------------------------------------------------------------------------- /.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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | eval_gemfile 'gemfiles/rails6.0.gemfile' 2 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![Build Status](https://github.com/zendesk/active_record_shards/workflows/CI/badge.svg)](https://github.com/zendesk/active_record_shards/actions?query=workflow%3ACI) 2 | 3 | # ActiveRecord Shards 4 | 5 | ActiveRecord Shards is an extension for ActiveRecord that provides support for sharded database and replicas. Basically it is just a nice way to 6 | switch between database connections. We've made the implementation very small, and have tried not to reinvent any wheels already present in ActiveRecord. 7 | 8 | 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. 9 | 10 | 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). 11 | 12 | - [Installation](#installation) 13 | - [Configuration](#configuration) 14 | - [Migrations](#migrations) 15 | - [Example](#example) 16 | - [Shared Model](#create-a-table-for-the-shared-not-sharded-model) 17 | - [Sharded Model](#create-a-table-for-the-sharded-model) 18 | - [Usage](#usage) 19 | - [Debugging](#debugging) 20 | 21 | ## Installation 22 | 23 | $ gem install active_record_shards 24 | 25 | and make sure to require 'active\_record\_shards' in some way. 26 | 27 | ### Rails 6.1 & 7.0 installation 28 | 29 | Rails 6.1 & 7.0 are **only** supported with `legacy_connection_handling` set to `true`. 30 | 31 | Enable the legacy handling in your configuration files e.g. `config/application.rb` by setting: 32 | 33 | ``` Ruby 34 | config.active_record.legacy_connection_handling = true 35 | ``` 36 | 37 | or 38 | 39 | ``` Ruby 40 | ActiveRecord::Base.legacy_connection_handling = true 41 | ``` 42 | 43 | ## Configuration 44 | 45 | Add the replica and shard configuration to config/database.yml: 46 | 47 | ```yaml 48 | production: 49 | adapter: mysql 50 | encoding: utf8 51 | database: my_app_main 52 | pool: 5 53 | host: db1 54 | username: root 55 | password: 56 | replica: 57 | host: db1_replica 58 | shards: 59 | 1: 60 | host: db_shard1 61 | database: my_app_shard 62 | replica: 63 | host: db_shard1_replica 64 | 2: 65 | host: db_shard2 66 | database: my_app_shard 67 | replica: 68 | host: db_shard2_replica 69 | ``` 70 | 71 | basically connections inherit configuration from the parent configuration file. 72 | 73 | ## Migrations 74 | 75 | ActiveRecord Shards also patches migrations to support running migrations on a shared (not sharded) or a sharded database. 76 | Each migration class has to specify a shard spec indicating where to run the migration. 77 | 78 | Valid shard specs: 79 | 80 | * `:none` - Run this migration on the shared database, not any shards 81 | * `:all` - Run this migration on all of the shards, not the shared database 82 | 83 | #### Example 84 | 85 | ###### Create a table for the shared (not sharded) model 86 | 87 | ```ruby 88 | class CreateAccounts < ActiveRecord::Migration 89 | shard :none 90 | 91 | def change 92 | create_table :accounts do |t| 93 | # This is NOT necessary for the gem to work, we just use it in the examples below demonstrating one way to switch shards 94 | t.integer :shard_id, null: false 95 | 96 | t.string :name 97 | end 98 | end 99 | end 100 | ``` 101 | 102 | ###### Create a table for the sharded model 103 | 104 | ```ruby 105 | class CreateProjects < ActiveRecord::Migration 106 | shard :all 107 | 108 | def change 109 | create_table :projects do |t| 110 | t.references :account 111 | t.string :name 112 | end 113 | end 114 | end 115 | ``` 116 | 117 | ## Usage 118 | 119 | 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. 120 | All the models that live on the shared database must be marked as not\_sharded: 121 | 122 | ```ruby 123 | class Account < ActiveRecord::Base 124 | not_sharded 125 | 126 | has_many :projects 127 | end 128 | 129 | class Project < ActiveRecord::Base 130 | belongs_to :account 131 | end 132 | ``` 133 | 134 | 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 135 | in a rack middleware and switch to the right shard: 136 | 137 | ```ruby 138 | class AccountMiddleware 139 | def initialize(app) 140 | @app = app 141 | end 142 | 143 | def call(env) 144 | account = lookup_account(env) 145 | 146 | if account 147 | ActiveRecord::Base.on_shard(account.shard_id) do 148 | @app.call(env) 149 | end 150 | else 151 | @app.call(env) 152 | end 153 | end 154 | 155 | def lookup_account(env) 156 | # ... 157 | end 158 | end 159 | ``` 160 | 161 | You can switch to the replica databases at any point by wrapping your code in an on\_replica block: 162 | 163 | ```ruby 164 | ActiveRecord::Base.on_replica do 165 | Account.find_by_big_expensive_query 166 | end 167 | ``` 168 | 169 | This will perform the query on the replica, and mark the returned instances as read-only. There is also a shortcut for this: 170 | 171 | ```ruby 172 | Account.on_replica.find_by_big_expensive_query 173 | ``` 174 | 175 | If you do not want instances returned from replicas to be marked as read-only, this can be disabled globally: 176 | 177 | `ActiveRecordShards.disable_replica_readonly_records = true` 178 | 179 | ## Debugging 180 | 181 | Show if a query went to primary or replica in the logs: 182 | 183 | ```Ruby 184 | require 'active_record_shards/sql_comments' 185 | ActiveRecordShards::SqlComments.enable 186 | ``` 187 | 188 | ## Copyright 189 | 190 | Copyright (c) 2011 Zendesk. See LICENSE for details. 191 | 192 | ## Authors 193 | Mick Staugaard, Eric Chapweske 194 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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/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.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 | -------------------------------------------------------------------------------- /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/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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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/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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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/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 | -------------------------------------------------------------------------------- /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/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 | -------------------------------------------------------------------------------- /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/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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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/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/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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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/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_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_test3.sql: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zendesk/active_record_shards/3c30e44f0585c523536478d2a269d9d812472435/test/schemas/ars_test3.sql -------------------------------------------------------------------------------- /test/schemas/ars_test3_shard0.sql: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zendesk/active_record_shards/3c30e44f0585c523536478d2a269d9d812472435/test/schemas/ars_test3_shard0.sql -------------------------------------------------------------------------------- /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/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_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.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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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/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/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 | -------------------------------------------------------------------------------- /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/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 | -------------------------------------------------------------------------------- /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/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/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/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 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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 | --------------------------------------------------------------------------------