More query router commands; settings last until changed again; docs (#25)

* readme

* touch up docs

* stuff

* refactor query router

* remove unused

* less verbose

* docs

* no link

* method rename

tests/ruby/tests.rb

@@ -1,7 +1,8 @@
 require "active_record"
 
-ActiveRecord.verbose_query_logs = true
-ActiveRecord::Base.logger = Logger.new(STDOUT)
+# Uncomment these two to see all queries.
+# ActiveRecord.verbose_query_logs = true
+# ActiveRecord::Base.logger = Logger.new(STDOUT)
 
 ActiveRecord::Base.establish_connection(
   adapter: "postgresql",
@@ -18,6 +19,13 @@ class TestTable < ActiveRecord::Base
   self.table_name = "test_table"
 end
 
+class TestSafeTable < ActiveRecord::Base
+  self.table_name = "test_safe_table"
+end
+
+class ShouldNeverHappenException < Exception
+end
+
 # # Create the table.
 class CreateTestTable < ActiveRecord::Migration[7.0]
   # Disable transasctions or things will fly out of order!
@@ -30,6 +38,7 @@ class CreateTestTable < ActiveRecord::Migration[7.0]
       # This will make this migration reversible!
       reversible do
         connection.execute "SET SHARD TO '#{x.to_i}'"
+        connection.execute "SET SERVER ROLE TO 'primary'"
       end
 
       # Always wrap the entire migration inside a transaction. If that's not possible,
@@ -47,28 +56,82 @@ class CreateTestTable < ActiveRecord::Migration[7.0]
   end
 end
 
-begin
-  CreateTestTable.migrate(:down)
-rescue Exception
-  puts "Tables don't exist yet"
-end
+class CreateSafeShardedTable < ActiveRecord::Migration[7.0]
+  # Disable transasctions or things will fly out of order!
+  disable_ddl_transaction!
 
-CreateTestTable.migrate(:up)
+  SHARDS = 3
 
-10.times do |x|
-  x += 1 # Postgres ids start at 1
-  r = TestTable.connection.execute "SET SHARDING KEY TO '#{x.to_i}'"
+  def up
+    SHARDS.times do |x|
+      # This will make this migration reversible!
+      connection.execute "SET SHARD TO '#{x.to_i}'"
+      connection.execute "SET SERVER ROLE TO 'primary'"
 
-  # Always wrap writes inside explicit transactions like these because ActiveRecord may fetch table info
-  # before actually issuing the `INSERT` statement. This ensures that that happens inside a transaction
-  # and the write goes to the correct shard.
-  TestTable.connection.transaction do
-    TestTable.create(id: x, name: "something_special_#{x.to_i}", description: "It's a surprise!")
+      connection.execute <<-SQL
+        CREATE TABLE test_safe_table (
+          id BIGINT PRIMARY KEY,
+          name VARCHAR,
+          description TEXT
+        ) PARTITION BY HASH (id);
+
+        CREATE TABLE test_safe_table_data PARTITION OF test_safe_table
+        FOR VALUES WITH (MODULUS #{SHARDS.to_i}, REMAINDER #{x.to_i});
+      SQL
+    end
+  end
+
+  def down
+    SHARDS.times do |x|
+      connection.execute "SET SHARD TO '#{x.to_i}'"
+      connection.execute "SET SERVER ROLE TO 'primary'"
+      connection.execute "DROP TABLE test_safe_table CASCADE"
+    end
   end
 end
 
-10.times do |x|
-  x += 1 # 0 confuses our sharding function
-  TestTable.connection.execute "SET SHARDING KEY TO '#{x.to_i}'"
-  puts TestTable.find_by_id(x).id
+20.times do
+  begin
+    CreateTestTable.migrate(:down)
+  rescue Exception
+    puts "Tables don't exist yet"
+  end
+
+  begin
+    CreateSafeShardedTable.migrate(:down)
+  rescue Exception
+    puts "Tables don't exist yet"
+  end
+
+  CreateTestTable.migrate(:up)
+  CreateSafeShardedTable.migrate(:up)
+
+  3.times do |x|
+    TestSafeTable.connection.execute "SET SHARD TO '#{x.to_i}'"
+    TestSafeTable.connection.execute "SET SERVER ROLE TO 'primary'"
+    TestSafeTable.connection.execute "TRUNCATE #{TestTable.table_name}"
+  end
+
+  10.times do |x|
+    x += 1 # Postgres ids start at 1
+    TestSafeTable.connection.execute "SET SHARDING KEY TO '#{x.to_i}'"
+    TestSafeTable.connection.execute "SET SERVER ROLE TO 'primary'"
+    TestSafeTable.create(id: x, name: "something_special_#{x.to_i}", description: "It's a surprise!")
+  end
+
+  10.times do |x|
+    x += 1 # 0 confuses our sharding function
+    TestSafeTable.connection.execute "SET SHARDING KEY TO '#{x.to_i}'"
+    TestSafeTable.connection.execute "SET SERVER ROLE TO 'replica'"
+    TestSafeTable.find_by_id(x).id
+  end
+end
+
+# Test wrong shard
+TestSafeTable.connection.execute "SET SHARD TO '1'"
+begin
+  TestSafeTable.create(id: 5, name: "test", description: "test description")
+  raise ShouldNeverHappenException("Uh oh")
+rescue ActiveRecord::StatementInvalid
+  puts "OK"
 end