Explicit shard selection; Rails tests (#24)

* Explicit shard selection; Rails tests

* try running ruby tests

* try without lockfile

* aha

* ok

tests/ruby/.ruby-version

@@ -0,0 +1 @@
+2.7.1

tests/ruby/Gemfile

@@ -0,0 +1,4 @@
+source "https://rubygems.org"
+
+gem "pg"
+gem "activerecord"

tests/ruby/Gemfile.lock

@@ -0,0 +1,30 @@
+GEM
+  remote: https://rubygems.org/
+  specs:
+    activemodel (7.0.2.2)
+      activesupport (= 7.0.2.2)
+    activerecord (7.0.2.2)
+      activemodel (= 7.0.2.2)
+      activesupport (= 7.0.2.2)
+    activesupport (7.0.2.2)
+      concurrent-ruby (~> 1.0, >= 1.0.2)
+      i18n (>= 1.6, < 2)
+      minitest (>= 5.1)
+      tzinfo (~> 2.0)
+    concurrent-ruby (1.1.9)
+    i18n (1.10.0)
+      concurrent-ruby (~> 1.0)
+    minitest (5.15.0)
+    pg (1.3.2)
+    tzinfo (2.0.4)
+      concurrent-ruby (~> 1.0)
+
+PLATFORMS
+  x86_64-linux
+
+DEPENDENCIES
+  activerecord
+  pg
+
+BUNDLED WITH
+   2.3.7

tests/ruby/tests.rb

@@ -1,11 +1,74 @@
-require 'pg'
+require "active_record"
 
-conn = PG.connect(host: '127.0.0.1', port: 5433, dbname: 'test')
+ActiveRecord.verbose_query_logs = true
+ActiveRecord::Base.logger = Logger.new(STDOUT)
 
-conn.exec( "SELECT * FROM pg_stat_activity" ) do |result|
-  puts "     PID | User             | Query"
-  result.each do |row|
-    puts " %7d | %-16s | %s " %
-      row.values_at('pid', 'usename', 'query')
+ActiveRecord::Base.establish_connection(
+  adapter: "postgresql",
+  host: "127.0.0.1",
+  port: 6432,
+  username: "sharding_user",
+  password: "sharding_user",
+  database: "rails_dev",
+  prepared_statements: false, # Transaction mode
+  advisory_locks: false, # Same
+)
+
+class TestTable < ActiveRecord::Base
+  self.table_name = "test_table"
+end
+
+# # Create the table.
+class CreateTestTable < ActiveRecord::Migration[7.0]
+  # Disable transasctions or things will fly out of order!
+  disable_ddl_transaction!
+
+  SHARDS = 3
+
+  def change
+    SHARDS.times do |x|
+      # This will make this migration reversible!
+      reversible do
+        connection.execute "SET SHARD TO '#{x.to_i}'"
+      end
+
+      # Always wrap the entire migration inside a transaction. If that's not possible,
+      # execute a `SET SHARD` command before every statement and make sure AR doesn't need
+      # to load database information beforehand (i.e. it's not the first query in the migration).
+      connection.transaction do
+        create_table :test_table, if_not_exists: true do |t|
+          t.string :name
+          t.string :description
+
+          t.timestamps
+        end
+      end
+    end
   end
-end
+end
+
+begin
+  CreateTestTable.migrate(:down)
+rescue Exception
+  puts "Tables don't exist yet"
+end
+
+CreateTestTable.migrate(:up)
+
+10.times do |x|
+  x += 1 # Postgres ids start at 1
+  r = TestTable.connection.execute "SET SHARDING KEY TO '#{x.to_i}'"
+
+  # 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!")
+  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
+end