railtie: Add relation extension (#41)

This commit adds two methods to `ActiveRecord::Relation`:
- `Relation#create_readyset_cache!`, which creates a new cache based on
the query represented by the relation; and
- `Relation#drop_cache!`, which drops the cache associated with the
query represented by the relation (if one exists)

---------

Co-authored-by: Paul Lemus <paullemus@protonmail.com>

lib/readyset.rb

@@ -5,10 +5,82 @@
 require 'readyset/middleware'
 require 'readyset/query'
 require 'readyset/railtie' if defined?(Rails::Railtie)
+require 'readyset/relation_extension'
 
 require 'active_record'
 
 module Readyset
+  attr_writer :configuration
+
+  def self.configuration
+    @configuration ||= Configuration.new
+  end
+
+  def self.configure
+    yield(configuration)
+  end
+
+  # Creates a new cache on ReadySet using the given ReadySet query ID or SQL query. Exactly one of
+  # the `id` or `sql` keyword arguments must be provided.
+  #
+  # This method is a no-op if a cache for the given ID/query already exists.
+  #
+  # @param [String] id the ReadySet query ID of the query from which a cache should be created
+  # @param [String] sql the SQL string from which a cache should be created
+  # @param [String] name the name for the cache being created
+  # @param [Boolean] always whether the cache should always be used. if this is true, queries to
+  # these caches will never fall back to the database
+  # @return [void]
+  # @raise [ArgumentError] raised if exactly one of the `id` or `sql` arguments was not provided
+  def self.create_cache!(id: nil, sql: nil, name: nil, always: false)
+    if (sql.nil? && id.nil?) || (!sql.nil? && !id.nil?)
+      raise ArgumentError, 'Exactly one of the `id` and `sql` parameters must be provided'
+    end
+
+    suffix = sql ? '%s' : '?'
+    from = (id || sql)
+
+    if always && name
+      Readyset.raw_query('CREATE CACHE ALWAYS ? FROM ' + suffix, name, from)
+    elsif always
+      Readyset.raw_query('CREATE CACHE ALWAYS FROM ' + suffix, from)
+    elsif name
+      Readyset.raw_query('CREATE CACHE ? FROM ' + suffix, name, from)
+    else
+      Readyset.raw_query('CREATE CACHE FROM ' + suffix, from)
+    end
+
+    nil
+  end
+
+  def self.current_config
+    configuration.inspect
+  end
+
+  # Creates a new cache on ReadySet using the given SQL query or ReadySet query ID. Exactly one of
+  # the `name_or_id` or `sql` keyword arguments must be provided.
+  #
+  # This method is a no-op if a cache for the given ID/query already doesn't exist.
+  #
+  # @param [String] name_or_id the name or the ReadySet query ID of the cache that should be dropped
+  # @param [String] sql a SQL string for a query whose associated cache should be dropped
+  # @return [void]
+  # @raise [ArgumentError] raised if exactly one of the `name_or_id` or `sql` arguments was not
+  # provided
+  def self.drop_cache!(name_or_id: nil, sql: nil)
+    if (sql.nil? && name_or_id.nil?) || (!sql.nil? && !name_or_id.nil?)
+      raise ArgumentError, 'Exactly one of the `name_or_id` and `sql` parameters must be provided'
+    end
+
+    if sql
+      Readyset.raw_query('DROP CACHE %s', sql)
+    else
+      Readyset.raw_query('DROP CACHE ?', name_or_id)
+    end
+
+    nil
+  end
+
   # Executes a raw SQL query against ReadySet. The query is sanitized prior to being executed.
   #
   # @param [Array<Object>] *sql_array the SQL array to be executed against ReadySet

lib/readyset/query.rb

@@ -164,23 +164,7 @@ def cache!(name: nil, always: false)
       elsif supported == :unsupported
         raise UnsupportedError, id
       else
-        query = 'CREATE CACHE '
-        params = []
-
-        if always
-          query += 'ALWAYS '
-        end
-
-        unless name.nil?
-          query += '? '
-          params.push(name)
-        end
-
-        query += 'FROM %s'
-        params.push(id)
-
-        Readyset.raw_query(query, *params)
-
+        Readyset.create_cache!(id: id, name: name, always: always)
         reload
       end
     end
@@ -199,7 +183,7 @@ def cached?
     # doesn't have a cache
     def drop_cache!
       if cached?
-        Readyset.raw_query('DROP CACHE %s', id)
+        Readyset.drop_cache!(id: id)
         reload
       else
         raise NotCachedError, id

lib/readyset/railtie.rb

@@ -7,5 +7,11 @@ class Railtie < Rails::Railtie
         prepend Readyset::ControllerExtension
       end
     end
+
+    initializer 'readyset.active_record' do |app|
+      ActiveSupport.on_load(:active_record) do
+        ActiveRecord::Relation.prepend(Readyset::RelationExtension)
+      end
+    end
   end
 end

lib/readyset/relation_extension.rb

@@ -0,0 +1,31 @@
+module Readyset
+  module RelationExtension
+    extend ActiveSupport::Concern
+
+    prepended do
+      # Creates a new cache on ReadySet for this query. This method is a no-op if a cache for the
+      # query already exists.
+      #
+      # NOTE: If the ActiveRecord query eager loads associations (e.g. via `#includes`), the
+      # the queries issues to do the eager loading will not have caches created. Those queries must
+      # have their caches created separately.
+      #
+      # @return [void]
+      def create_readyset_cache!
+        Readyset.create_cache!(sql: to_sql)
+      end
+
+      # Drops the cache on ReadySet associated with this query. This method is a no-op if a cache
+      # for the query already doesn't exist.
+      #
+      # NOTE: If the ActiveRecord query eager loads associations (e.g. via `#includes`), the
+      # the queries issues to do the eager loading will not have caches dropped. Those queries must
+      # have their caches dropped separately.
+      #
+      # @return [void]
+      def drop_readyset_cache!
+        Readyset.drop_cache!(sql: to_sql)
+      end
+    end
+  end
+end

spec/internal/app/models/cat.rb

@@ -0,0 +1 @@
+class Cat < ActiveRecord::Base; end

spec/internal/db/schema.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 ActiveRecord::Schema.define do
-  # Set up any tables you need to exist for your test suite that don't belong
-  # in migrations.
+  create_table(:cats, :force => true) do |t|
+    t.string :name
+    t.timestamps
+  end
 end

spec/query_spec.rb

@@ -309,62 +309,24 @@
     end
 
     context 'when the query is supported and not cached' do
-      let(:query) { build(:seen_but_not_cached_query) }
-
-      before { allow(query).to receive(:reload) }
+      subject { query.cache!(**args) }
 
-      context 'when only the "always" parameter is passed' do
-        subject { query.cache!(always: true) }
+      let(:args) { { always: true, name: 'test name' } }
+      let(:query) { build(:seen_but_not_cached_query) }
 
-        it_behaves_like 'a wrapper around a ReadySet SQL extension',
-            'CREATE CACHE ALWAYS FROM %s' do
-          let(:args) { [query.id] }
+      before do
+        allow(query).to receive(:reload)
+        allow(Readyset).to receive(:create_cache!).with(id: query.id, **args)
 
-          it 'invokes Readyset::Query#reload' do
-            expect(query).to have_received(:reload)
-          end
-        end
+        subject
       end
 
-      context 'when only the "name" parameter is passed' do
-        subject { query.cache!(name: name) }
-
-        let(:name) { 'test cache' }
-
-        it_behaves_like 'a wrapper around a ReadySet SQL extension', 'CREATE CACHE ? FROM %s' do
-          let(:args) { [name, query.id] }
-
-          it 'invokes Readyset::Query#reload' do
-            expect(query).to have_received(:reload)
-          end
-        end
+      it 'invokes Readyset.create_cache! with the correct arguments' do
+        expect(Readyset).to have_received(:create_cache!).with(id: query.id, **args)
       end
 
-      context 'when both the "always" and "name" parameters are passed' do
-        subject { query.cache!(always: true, name: name) }
-
-        let(:name) { 'test cache' }
-
-        it_behaves_like 'a wrapper around a ReadySet SQL extension',
-'CREATE CACHE ALWAYS ? FROM %s' do
-          let(:args) { [name, query.id] }
-
-          it 'invokes Readyset::Query#reload' do
-            expect(query).to have_received(:reload)
-          end
-        end
-      end
-
-      context 'when neither the "always" nor the "name" parameters are passed' do
-        subject { query.cache! }
-
-        it_behaves_like 'a wrapper around a ReadySet SQL extension', 'CREATE CACHE FROM %s' do
-          let(:args) { [query.id] }
-
-          it 'invokes Readyset::Query#reload' do
-            expect(query).to have_received(:reload)
-          end
-        end
+      it 'invokes Readyset::Query#reload' do
+        expect(query).to have_received(:reload)
       end
     end
   end
@@ -393,14 +355,19 @@
     context 'when the query is cached' do
       let(:query) { build(:cached_query) }
 
-      before { allow(query).to receive(:reload) }
+      before do
+        allow(query).to receive(:reload)
+        allow(Readyset).to receive(:drop_cache!).with(id: query.id)
+
+        subject
+      end
 
-      it_behaves_like 'a wrapper around a ReadySet SQL extension', 'DROP CACHE %s' do
-        let(:args) { [query.id] }
+      it 'invokes Readyset.drop_cache!' do
+        expect(Readyset).to have_received(:drop_cache!).with(id: query.id)
+      end
 
-        it 'invokes Readyset::Query#reload' do
-          expect(query).to have_received(:reload)
-        end
+      it 'invokes Readyset::Query#reload' do
+        expect(query).to have_received(:reload)
       end
     end
 

spec/railtie_spec.rb

@@ -2,21 +2,29 @@
 
 require 'rails_helper'
 
-RSpec.describe Readyset::Railtie, type: :controller do
-  controller(ActionController::Base) do
-    # Define a test action
-    def index
-      render plain: 'Test'
+RSpec.describe Readyset::Railtie do
+  describe 'readyset.action_controller', type: :controller do
+    controller(ActionController::Base) do
+      # Define a test action
+      def index
+        render plain: 'Test'
+      end
+    end
+
+    before do
+      routes.draw do
+        get 'index' => 'anonymous#index'
+      end
     end
-  end
 
-  before do
-    routes.draw do
-      get 'index' => 'anonymous#index'
+    it 'includes ControllerExtension into ActionController::Base' do
+      expect(controller.class.ancestors).to include(Readyset::ControllerExtension)
     end
   end
 
-  it 'includes ControllerExtension into ActionController::Base' do
-    expect(controller.class.ancestors).to include(Readyset::ControllerExtension)
+  describe 'readyset.active_record' do
+    it 'includes RelationExtension into ActiveRecord::Relation' do
+      expect(ActiveRecord::Relation.ancestors).to include(Readyset::RelationExtension)
+    end
   end
 end