|
1 |
| -require "active_record" |
| 1 | +# frozen_string_literal: true |
| 2 | + |
| 3 | +require 'active_record' |
2 | 4 |
|
3 | 5 | # Uncomment these two to see all queries.
|
4 | 6 | # ActiveRecord.verbose_query_logs = true
|
5 | 7 | # ActiveRecord::Base.logger = Logger.new(STDOUT)
|
6 | 8 |
|
7 | 9 | ActiveRecord::Base.establish_connection(
|
8 |
| - adapter: "postgresql", |
9 |
| - host: "127.0.0.1", |
| 10 | + adapter: 'postgresql', |
| 11 | + host: '127.0.0.1', |
10 | 12 | port: 6432,
|
11 |
| - username: "sharding_user", |
12 |
| - password: "sharding_user", |
13 |
| - database: "rails_dev", |
| 13 | + username: 'sharding_user', |
| 14 | + password: 'sharding_user', |
| 15 | + database: 'rails_dev', |
14 | 16 | prepared_statements: false, # Transaction mode
|
15 |
| - advisory_locks: false, # Same |
| 17 | + advisory_locks: false # Same |
16 | 18 | )
|
17 | 19 |
|
18 |
| -class TestTable < ActiveRecord::Base |
19 |
| - self.table_name = "test_table" |
20 |
| -end |
21 |
| - |
22 | 20 | class TestSafeTable < ActiveRecord::Base
|
23 |
| - self.table_name = "test_safe_table" |
| 21 | + self.table_name = 'test_safe_table' |
24 | 22 | end
|
25 | 23 |
|
26 |
| -class ShouldNeverHappenException < Exception |
27 |
| -end |
28 |
| - |
29 |
| -# # Create the table. |
30 |
| -class CreateTestTable < ActiveRecord::Migration[7.0] |
31 |
| - # Disable transasctions or things will fly out of order! |
32 |
| - disable_ddl_transaction! |
33 |
| - |
34 |
| - SHARDS = 3 |
35 |
| - |
36 |
| - def change |
37 |
| - SHARDS.times do |x| |
38 |
| - # This will make this migration reversible! |
39 |
| - reversible do |
40 |
| - connection.execute "SET SHARD TO '#{x.to_i}'" |
41 |
| - connection.execute "SET SERVER ROLE TO 'primary'" |
42 |
| - end |
43 |
| - |
44 |
| - # Always wrap the entire migration inside a transaction. If that's not possible, |
45 |
| - # execute a `SET SHARD` command before every statement and make sure AR doesn't need |
46 |
| - # to load database information beforehand (i.e. it's not the first query in the migration). |
47 |
| - connection.transaction do |
48 |
| - create_table :test_table, if_not_exists: true do |t| |
49 |
| - t.string :name |
50 |
| - t.string :description |
51 |
| - |
52 |
| - t.timestamps |
53 |
| - end |
54 |
| - end |
55 |
| - end |
56 |
| - end |
| 24 | +class ShouldNeverHappenException < RuntimeError |
57 | 25 | end
|
58 | 26 |
|
59 | 27 | class CreateSafeShardedTable < ActiveRecord::Migration[7.0]
|
@@ -85,53 +53,60 @@ def down
|
85 | 53 | SHARDS.times do |x|
|
86 | 54 | connection.execute "SET SHARD TO '#{x.to_i}'"
|
87 | 55 | connection.execute "SET SERVER ROLE TO 'primary'"
|
88 |
| - connection.execute "DROP TABLE test_safe_table CASCADE" |
| 56 | + connection.execute 'DROP TABLE test_safe_table CASCADE' |
89 | 57 | end
|
90 | 58 | end
|
91 | 59 | end
|
92 | 60 |
|
93 |
| -20.times do |
94 |
| - begin |
95 |
| - CreateTestTable.migrate(:down) |
96 |
| - rescue Exception |
97 |
| - puts "Tables don't exist yet" |
98 |
| - end |
| 61 | +SHARDS = 3 |
99 | 62 |
|
| 63 | +2.times do |
100 | 64 | begin
|
101 | 65 | CreateSafeShardedTable.migrate(:down)
|
102 | 66 | rescue Exception
|
103 | 67 | puts "Tables don't exist yet"
|
104 | 68 | end
|
105 | 69 |
|
106 |
| - CreateTestTable.migrate(:up) |
107 | 70 | CreateSafeShardedTable.migrate(:up)
|
108 | 71 |
|
109 |
| - 3.times do |x| |
| 72 | + SHARDS.times do |x| |
110 | 73 | TestSafeTable.connection.execute "SET SHARD TO '#{x.to_i}'"
|
111 | 74 | TestSafeTable.connection.execute "SET SERVER ROLE TO 'primary'"
|
112 |
| - TestSafeTable.connection.execute "TRUNCATE #{TestTable.table_name}" |
| 75 | + TestSafeTable.connection.execute "TRUNCATE #{TestSafeTable.table_name}" |
113 | 76 | end
|
114 | 77 |
|
115 |
| - 10.times do |x| |
| 78 | + # Equivalent to Makara's stick_to_master! except it sticks until it's changed. |
| 79 | + TestSafeTable.connection.execute "SET SERVER ROLE TO 'primary'" |
| 80 | + |
| 81 | + 200.times do |x| |
116 | 82 | x += 1 # Postgres ids start at 1
|
117 | 83 | TestSafeTable.connection.execute "SET SHARDING KEY TO '#{x.to_i}'"
|
118 |
| - TestSafeTable.connection.execute "SET SERVER ROLE TO 'primary'" |
119 | 84 | TestSafeTable.create(id: x, name: "something_special_#{x.to_i}", description: "It's a surprise!")
|
120 | 85 | end
|
121 | 86 |
|
122 |
| - 10.times do |x| |
| 87 | + TestSafeTable.connection.execute "SET SERVER ROLE TO 'replica'" |
| 88 | + |
| 89 | + 100.times do |x| |
123 | 90 | x += 1 # 0 confuses our sharding function
|
124 | 91 | TestSafeTable.connection.execute "SET SHARDING KEY TO '#{x.to_i}'"
|
125 |
| - TestSafeTable.connection.execute "SET SERVER ROLE TO 'replica'" |
| 92 | + TestSafeTable.find_by_id(x).id |
| 93 | + end |
| 94 | + |
| 95 | + # Will use the query parser to direct reads to replicas |
| 96 | + TestSafeTable.connection.execute "SET SERVER ROLE TO 'auto'" |
| 97 | + |
| 98 | + 100.times do |x| |
| 99 | + x += 101 |
| 100 | + TestSafeTable.connection.execute "SET SHARDING KEY TO '#{x.to_i}'" |
126 | 101 | TestSafeTable.find_by_id(x).id
|
127 | 102 | end
|
128 | 103 | end
|
129 | 104 |
|
130 | 105 | # Test wrong shard
|
131 | 106 | TestSafeTable.connection.execute "SET SHARD TO '1'"
|
132 | 107 | begin
|
133 |
| - TestSafeTable.create(id: 5, name: "test", description: "test description") |
134 |
| - raise ShouldNeverHappenException("Uh oh") |
| 108 | + TestSafeTable.create(id: 5, name: 'test', description: 'test description') |
| 109 | + raise ShouldNeverHappenException('Uh oh') |
135 | 110 | rescue ActiveRecord::StatementInvalid
|
136 |
| - puts "OK" |
| 111 | + puts 'OK' |
137 | 112 | end
|
0 commit comments