routing: Add Readyset.route

Adds a `Readyset.route` method, which takes a block and routes to
ReadySet all the queries that occur within that block.

lib/readyset.rb

@@ -6,8 +6,6 @@
 require 'readyset/railtie' if defined?(Rails::Railtie)
 require 'readyset/relation_extension'
 
-require 'active_record'
-
 module Readyset
   attr_writer :configuration
 
@@ -56,7 +54,7 @@ def self.create_cache!(id: nil, sql: nil, name: nil, always: false)
       params.push(id)
     end
 
-    Readyset.raw_query(query, *params)
+    raw_query(query, *params)
 
     nil
   end
@@ -81,9 +79,9 @@ def self.drop_cache!(name_or_id: nil, sql: nil)
     end
 
     if sql
-      Readyset.raw_query('DROP CACHE %s', sql)
+      raw_query('DROP CACHE %s', sql)
     else
-      Readyset.raw_query('DROP CACHE ?', name_or_id)
+      raw_query('DROP CACHE ?', name_or_id)
     end
 
     nil
@@ -93,8 +91,43 @@ def self.drop_cache!(name_or_id: nil, sql: nil)
   #
   # @param [Array<Object>] *sql_array the SQL array to be executed against ReadySet
   # @return [PG::Result]
-  def self.raw_query(*sql_array)
-    ActiveRecord::Base.establish_connection(Readyset.configuration.connection_url)
-    ActiveRecord::Base.connection.execute(ActiveRecord::Base.sanitize_sql_array(sql_array))
+  def self.raw_query(*sql_array) # :nodoc:
+    ActiveRecord::Base.connected_to(role: reading_role, shard: shard, prevent_writes: false) do
+      ActiveRecord::Base.connection.execute(ActiveRecord::Base.sanitize_sql_array(sql_array))
+    end
+  end
+
+  # Routes to ReadySet any queries that occur in the given block. If `prevent_writes` is false, an
+  # attempt to execute a write within the given block will raise an error. Keep in mind that if
+  # `prevent_writes` is true, any writes that occur within the given block will be proxied through
+  # ReadySet to the database.
+  #
+  # @param [Boolean] prevent_writes prevent writes from being executed on the connection to ReadySet
+  # @yield a block whose queries should be routed to ReadySet
+  # @return the value of the last line of the block
+  def self.route(prevent_writes: true, &block)
+    if prevent_writes
+      ActiveRecord::Base.
+        connected_to(role: reading_role, shard: shard, prevent_writes: true, &block)
+    else
+      ActiveRecord::Base.
+        connected_to(role: writing_role, shard: shard, prevent_writes: false, &block)
+    end
+  end
+
+  private
+
+  class << self
+    private delegate(:shard, to: :configuration)
+  end
+
+  def self.reading_role
+    ActiveRecord.reading_role
+  end
+  private_class_method :reading_role
+
+  def self.writing_role
+    ActiveRecord.writing_role
   end
+  private_class_method :writing_role
 end

lib/readyset/configuration.rb

@@ -1,11 +1,13 @@
 # lib/readyset/configuration.rb
+require 'active_record'
 
 module Readyset
   class Configuration
-    attr_accessor :connection_url
+    attr_accessor :connection_url, :shard
 
     def initialize
       @connection_url = ENV['READYSET_URL'] || default_connection_url
+      @shard = :readyset
     end
 
     def default_connection_url

readyset.gemspec

@@ -42,4 +42,5 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency 'rspec-rails', '~> 6.0'
   spec.add_development_dependency 'rubocop-airbnb'
   spec.add_development_dependency 'sqlite3', '~> 1.6'
+  spec.add_development_dependency 'pry'
 end

spec/internal/app/models/application_record.rb

@@ -0,0 +1,11 @@
+class ApplicationRecord < ActiveRecord::Base
+  self.abstract_class = true
+
+  connects_to shards: {
+    primary: { reading: :primary, writing: :primary },
+    Readyset.configuration.shard => {
+      reading: Readyset.configuration.shard,
+      writing: Readyset.configuration.shard,
+    },
+  }
+end

spec/internal/app/models/cat.rb

@@ -1 +1,5 @@
-class Cat < ActiveRecord::Base; end
+class Cat < ApplicationRecord
+  def ==(other)
+    name == other.name
+  end
+end

spec/internal/config/database.yml

@@ -1,3 +1,8 @@
 test:
-  adapter:  sqlite3
-  database: db/combustion_test.sqlite
+  primary:
+    adapter:  sqlite3
+    database: db/combustion_test.sqlite
+  readyset:
+    adapter:  sqlite3
+    database: db/combustion_test_readyset.sqlite
+

spec/ready_set_spec.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'spec_helper'
+
 RSpec.describe Readyset do
   it 'has a version number' do
     expect(Readyset::VERSION).not_to be nil
@@ -116,35 +118,146 @@
   describe '.raw_query' do
     subject { Readyset.raw_query(*query) }
 
-    let(:connection) { instance_double(ActiveRecord::ConnectionAdapters::AbstractAdapter) }
-    let(:connection_url) { 'postgres://postgres:readyset@127.0.0.1:5432/test' }
-    let(:query) { ['SELECT * FROM t WHERE x = ?', 0] }
-    let(:results) { Object.new }
-    let(:sanitized_query) { 'SELECT * FROM t WHERE x = 0' }
+    let(:query) { ['SELECT * FROM cats WHERE name = ?', 'whiskers'] }
 
     before do
-      Readyset.configuration.connection_url = connection_url
-
-      allow(ActiveRecord::Base).to receive(:establish_connection).with(connection_url)
-      allow(ActiveRecord::Base).to receive(:connection).and_return(connection)
-      allow(ActiveRecord::Base).to receive(:sanitize_sql_array).with(query).
-        and_return(sanitized_query)
-      allow(connection).to receive(:execute).with(sanitized_query).and_return(results)
+      ActiveRecord::Base.connected_to(shard: Readyset.configuration.shard) do
+        Cat.create!(name: 'whiskers')
+        Cat.create!(name: 'fluffy')
+        Cat.create!(name: 'tails')
+      end
 
       subject
     end
 
-    it 'establishes a connection to ReadySet via the configured URL' do
-      expect(ActiveRecord::Base).to have_received(:establish_connection).with(connection_url)
+    it 'returns the results from the query' do
+      expect(subject.size).to eq(1)
+      expect(subject.first['name']).to eq('whiskers')
     end
+  end
+
+  describe '.route' do
+    subject { Readyset.route(prevent_writes: prevent_writes, &block) }
+
+    let(:query_results) { instance_double(Cat) }
+
+    RSpec.shared_examples 'uses the expected connection parameters' do |role, shard|
+      it "sets the role to be #{role}" do
+        begin
+          subject
+        rescue ActiveRecord::ReadOnlyError
+        end
+
+        expect(@role).to eq(ActiveRecord.writing_role)
+      end
+
+      it "sets the shard to be #{shard}" do
+        begin
+          subject
+        rescue ActiveRecord::ReadOnlyError
+        end
+
+        expect(@shard).to eq(Readyset.configuration.shard)
+      end
+    end
+
+    context 'when prevent_writes is true' do
+      let(:prevent_writes) { true }
+
+      context 'when the block contains a write query' do
+        let(:block) do
+          Proc.new do
+            @role = ActiveRecord::Base.connection.role
+            @shard = ActiveRecord::Base.connection.shard
+            Cat.create!(name: 'whiskers')
+          end
+        end
 
-    it 'sanitizes query input' do
-      expect(ActiveRecord::Base).to have_received(:sanitize_sql_array).with(query)
+        it 'raises an ActiveRecord::ReadOnlyError' do
+          expect { subject }.to raise_error(ActiveRecord::ReadOnlyError)
+        end
+
+        include_examples 'uses the expected connection parameters', ActiveRecord.writing_role,
+          Readyset.configuration.shard
+      end
+
+      context 'when the block contains a read query' do
+        let(:block) do
+          Proc.new do
+            @role = ActiveRecord::Base.connection.role
+            @shard = ActiveRecord::Base.connection.shard
+            'test return value'
+          end
+        end
+
+        it 'returns the result of the block' do
+          expect(subject).to eq('test return value')
+        end
+
+        include_examples 'uses the expected connection parameters', ActiveRecord.writing_role,
+          Readyset.configuration.shard
+      end
     end
 
-    it 'executes the query on the connection and returns the results' do
-      expect(connection).to have_received(:execute).with(sanitized_query)
-      is_expected.to eq(results)
+    context 'when prevent_writes is false' do
+      let(:prevent_writes) { false }
+
+      context 'when the block contains a write query' do
+        let(:block) do
+          Proc.new do
+            @role = ActiveRecord::Base.connection.role
+            @shard = ActiveRecord::Base.connection.shard
+            Cat.create!(name: 'whiskers')
+            'test return value'
+          end
+        end
+
+        it 'returns the result of the block' do
+          expect(subject).to eq('test return value')
+        end
+
+        it 'executes the write against ReadySet' do
+          subject
+
+          exists = ActiveRecord::Base.connected_to(shard: Readyset.configuration.shard) do
+            Cat.where(name: 'whiskers').exists?
+          end
+
+          expect(exists).to eq(true)
+
+          exists = ActiveRecord::Base.connected_to(shard: :primary) do
+            Cat.where(name: 'whiskers').exists?
+          end
+
+          expect(exists).to eq(false)
+        end
+
+        include_examples 'uses the expected connection parameters', ActiveRecord.writing_role,
+          Readyset.configuration.shard
+      end
+
+      context 'when the block contains a read query' do
+        let(:block) do
+          Proc.new do
+            @role = ActiveRecord::Base.connection.role
+            @shard = ActiveRecord::Base.connection.shard
+            Cat.where(name: 'whiskers').exists?
+          end
+        end
+
+        before do
+          ActiveRecord::Base.connected_to(shard: Readyset.configuration.shard) do
+            Cat.create!(name: 'whiskers')
+          end
+        end
+
+        it 'executes the read against ReadySet' do
+          expect(subject).to eq(true)
+        end
+
+        include_examples 'uses the expected connection parameters', ActiveRecord.writing_role,
+          Readyset.configuration.shard
+      end
     end
   end
 end

spec/shared_examples.rb

@@ -4,7 +4,8 @@
   let(:raw_query_result) { [] }
 
   before do
-    allow(Readyset).to receive(:raw_query).with(sql_command, *args).and_return(raw_query_result)
+    allow(Readyset).to receive(:raw_query).with(sql_command, *args).
+      and_return(raw_query_result)
 
     subject
   end

spec/spec_helper.rb

@@ -7,7 +7,16 @@
 require 'factory_bot'
 require 'readyset'
 require_relative 'shared_examples'
-Combustion.initialize! :action_controller, :active_record
+
+Combustion.initialize! :action_controller, :active_record do
+  config.eager_load = true
+end
+
+# This is a bit of a hack. Combustion doesn't appear to support migrating multiple databases, so we
+# just copy the primary database file to serve as the database for our fake ReadySet instance
+primary_db_file = Rails.configuration.database_configuration['test']['primary']['database']
+readyset_db_file = Rails.configuration.database_configuration['test']['readyset']['database']
+FileUtils.cp("spec/internal/#{primary_db_file}", "spec/internal/#{readyset_db_file}")
 
 RSpec.configure do |config|
   # Enable flags like --only-failures and --next-failure