relation: Add .readyset_explain (#62)

This commit adds a `.readyset_explain` method to our relation extension.
This method invokes `EXPLAIN CREATE CACHE` upstream on ReadySet and
returns information about the query from ReadySet.

Closes #44

.gitignore

@@ -13,3 +13,5 @@ Gemfile.lock
 spec/internal/tmp/
 spec/internal/config/storage.yml
 spec/internal/db/*.sqlite
+spec/internal/db/*.sqlite-shm
+spec/internal/db/*.sqlite-wal

lib/readyset.rb

@@ -1,8 +1,10 @@
 # lib/readyset.rb
 
+require 'readyset/caches'
 require 'readyset/configuration'
 require 'readyset/controller_extension'
 require 'readyset/model_extension'
+require 'readyset/explain'
 require 'readyset/query'
 require 'readyset/query/cached_query'
 require 'readyset/query/proxied_query'
@@ -82,12 +84,24 @@ def self.drop_cache!(name)
     nil
   end
 
+  # Gets information about the given query from ReadySet, including whether it's supported to be
+  # cached, its current status, the rewritten query text, and the query ID.
+  #
+  # The information about the given query is retrieved by invoking `EXPLAIN CREATE CACHE FROM` on
+  # ReadySet.
+  #
+  # @param [String] a query about which information should be retrieved
+  # @return [Explain]
+  def self.explain(query)
+    Explain.call(query)
+  end
+
   # Executes a raw SQL query against ReadySet. The query is sanitized prior to being executed.
   # @note This method is not part of the public API.
   # @param sql_array [Array<Object>] the SQL array to be executed against ReadySet.
   # @return [PG::Result] the result of executing the SQL query.
   def self.raw_query(*sql_array) # :nodoc:
-    ActiveRecord::Base.connected_to(role: reading_role, shard: shard, prevent_writes: false) do
+    ActiveRecord::Base.connected_to(role: writing_role, shard: shard, prevent_writes: false) do
       ActiveRecord::Base.connection.execute(ActiveRecord::Base.sanitize_sql_array(sql_array))
     end
   end

lib/readyset/caches.rb

@@ -0,0 +1,21 @@
+module Readyset
+  # Defines the DSL used in the gem's "migration" files. The DSL should be used by inheriting
+  # from this class and invoking the `.cache` class method to define new caches.
+  class Caches
+    class << self
+      attr_reader :caches
+    end
+
+    def self.cache(id:, always: false)
+      @caches ||= Set.new
+
+      query = yield
+
+      @caches << Query::CachedQuery.new(
+        id: id,
+        text: query.strip,
+        always: always,
+      )
+    end
+  end
+end

lib/readyset/configuration.rb

@@ -3,9 +3,10 @@
 
 module Readyset
   class Configuration
-    attr_accessor :shard
+    attr_accessor :migration_path, :shard
 
     def initialize
+      @migration_path = File.join(Rails.root, 'db/readyset_caches.rb')
       @shard = :readyset
     end
   end

lib/readyset/explain.rb

@@ -0,0 +1,60 @@
+module Readyset
+  # Represents the result of an `EXPLAIN CREATE CACHE` invocation on ReadySet.
+  class Explain
+    attr_reader :id, :text, :supported
+
+    # Gets information about the given query from ReadySet, including whether it's supported to be
+    # cached, its current status, the rewritten query text, and the query ID.
+    #
+    # The information about the given query is retrieved by invoking `EXPLAIN CREATE CACHE FROM` on
+    # ReadySet.
+    #
+    # @param [String] a query about which information should be retrieved
+    # @return [Explain]
+    def self.call(query)
+      raw_results = Readyset.raw_query('EXPLAIN CREATE CACHE FROM %s', query)
+      from_readyset_results(**raw_results.first.to_h.symbolize_keys)
+    end
+
+    # Creates a new `Explain` with the given attributes.
+    #
+    # @param [String] id the ID of the query
+    # @param [String] text the query text
+    # @param [Symbol] supported the supported status of the query
+    # @return [Explain]
+    def initialize(id:, text:, supported:) # :nodoc:
+      @id = id
+      @text = text
+      @supported = supported
+    end
+
+    # Compares `self` with another `Explain` by comparing them attribute-wise.
+    #
+    # @param [Explain] other the `Explain` to which `self` should be compared
+    # @return [Boolean]
+    def ==(other)
+      id == other.id &&
+        text == other.text &&
+        supported == other.supported
+    end
+
+    # Returns true if the explain information returned by ReadySet indicates that the query is
+    # unsupported.
+    #
+    # @return [Boolean]
+    def unsupported?
+      supported == :unsupported
+    end
+
+    private
+
+    def self.from_readyset_results(**attributes)
+      new(
+        id: attributes[:'query id'],
+        text: attributes[:query],
+        supported: attributes[:'readyset supported'].to_sym,
+      )
+    end
+    private_class_method :from_readyset_results
+  end
+end

lib/readyset/query/cached_query.rb

@@ -45,7 +45,7 @@ def self.find(id)
       # @param [Hash] attributes the attributes from which the `CachedQuery` should be
       # constructed
       # @return [CachedQuery]
-      def initialize(id:, text:, name:, always:, count:)
+      def initialize(id: nil, text:, name: nil, always: nil, count: nil)
         @id = id
         @text = text
         @name = name

lib/readyset/railtie.rb

@@ -14,5 +14,9 @@ class Railtie < Rails::Railtie
         ActiveRecord::Relation.prepend(Readyset::RelationExtension)
       end
     end
+
+    rake_tasks do
+      Dir[File.join(File.dirname(__FILE__), '../tasks/*.rake')].each { |f| load f }
+    end
   end
 end

lib/readyset/relation_extension.rb

@@ -7,7 +7,7 @@ module RelationExtension
       # 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
+      # the queries issued to do the eager loading will not have caches created. Those queries must
       # have their caches created separately.
       #
       # @return [void]
@@ -19,13 +19,21 @@ def create_readyset_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
+      # the queries issued 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
+
+      # Gets information about this query from ReadySet, including the query's ID, the normalized
+      # query text, and whether the query is supported by ReadySet.
+      #
+      # @return [Readyset::Explain]
+      def readyset_explain
+        Readyset.explain(to_sql)
+      end
     end
   end
 end

lib/tasks/readyset.rake

@@ -0,0 +1,75 @@
+require 'colorize'
+require 'erb'
+require 'progressbar'
+
+namespace :readyset do
+  namespace :caches do
+    desc 'Dumps the set of caches that currently exist on ReadySet to a file'
+    task dump: :environment do
+      Rails.application.eager_load!
+
+      template = File.read(File.join(File.dirname(__FILE__), '../templates/caches.rb.tt'))
+
+      queries = Readyset::Query::CachedQuery.all
+      f = File.new(Readyset.configuration.migration_path, 'w')
+      f.write(ERB.new(template, trim_mode: '-').result(binding))
+      f.close
+    end
+
+    desc 'Synchronizes the caches on ReadySet such that the caches on ReadySet match those ' \
+      'listed in db/readyset_caches.rb'
+    task migrate: :environment do
+      Rails.application.eager_load!
+
+      file = Readyset.configuration.migration_path
+
+      # We load the definition of the `Readyset::Caches` subclass in the context of a
+      # container object so we can be sure that we are never re-opening a previously-defined
+      # subclass of `Readyset::Caches`. When the container object is garbage collected, the
+      # definition of the `Readyset::Caches` subclass is garbage collected too
+      container = Object.new
+      container.instance_eval(File.read(file))
+      caches = container.singleton_class::ReadysetCaches.caches
+
+      caches_on_readyset = Readyset::Query::CachedQuery.all.index_by(&:id)
+      caches_on_readyset_ids = caches_on_readyset.keys.to_set
+
+      caches_in_migration_file = caches.index_by(&:id)
+      caches_in_migration_file_ids = caches_in_migration_file.keys.to_set
+
+      to_drop_ids = caches_on_readyset_ids - caches_in_migration_file_ids
+      to_create_ids = caches_in_migration_file_ids - caches_on_readyset_ids
+
+      if to_drop_ids.size.positive? || to_create_ids.size.positive?
+        dropping = 'Dropping'.red
+        creating = 'creating'.green
+        print "#{dropping} #{to_drop_ids.size} caches and #{creating} #{to_create_ids.size} " \
+          'caches. Continue? (y/n) '
+        $stdout.flush
+        y_or_n = STDIN.gets.strip
+
+        if y_or_n == 'y'
+          if to_drop_ids.size.positive?
+            bar = ProgressBar.create(title: 'Dropping caches', total: to_drop_ids.size)
+
+            to_drop_ids.each do |id|
+              bar.increment
+              Readyset.drop_cache!(name_or_id: id)
+            end
+          end
+
+          if to_create_ids.size.positive?
+            bar = ProgressBar.create(title: 'Creating caches', total: to_create_ids.size)
+
+            to_create_ids.each do |id|
+              bar.increment
+              Readyset.create_cache!(id: id)
+            end
+          end
+        end
+      else
+        puts 'Nothing to do'
+      end
+    end
+  end
+end

lib/templates/caches.rb.tt

@@ -0,0 +1,11 @@
+class ReadysetCaches < Readyset::Caches
+<% queries.each do |query| -%>
+  cache id: <%= query.id.dump %>, always: <%= query.always %> do
+    <<~SQL
+      <%= query.text.gsub("\n", "\n      ") %>
+    SQL
+  end
+
+<%- end -%>
+end
+

readyset.gemspec

@@ -34,6 +34,8 @@ Gem::Specification.new do |spec|
   spec.add_dependency 'actionpack', '>= 6.1'
   spec.add_dependency 'activerecord', '>= 6.1'
   spec.add_dependency 'activesupport', '>= 6.1'
+  spec.add_dependency 'colorize', '~> 1.1'
+  spec.add_dependency 'progressbar', '~> 1.13'
   spec.add_dependency 'rake', '~> 13.0'
 
   spec.add_development_dependency 'combustion', '~> 1.3'

spec/caches_spec.rb

@@ -0,0 +1,40 @@
+RSpec.describe Readyset::Caches do
+  describe '.cache' do
+    after(:each) do
+      Readyset::Caches.instance_variable_set(:@caches, nil)
+    end
+
+    it 'adds a cache with the given attributes to the @caches ivar' do
+      query = build(:cached_query, always: true, count: nil, name: nil)
+
+      Readyset::Caches.cache(always: true, id: query.id) { query.text }
+
+      caches = Readyset::Caches.instance_variable_get(:@caches)
+      expect(caches.size).to eq(1)
+      expect(caches.first).to eq(query)
+    end
+
+    context 'when no always parameter is passed' do
+      it 'defaults the always parameter to false' do
+        query = build(:cached_query, count: nil, name: nil)
+
+        Readyset::Caches.cache(id: query.id) { query.text }
+
+        always = Readyset::Caches.instance_variable_get(:@caches).first.always
+        expect(always).to eq(false)
+      end
+    end
+  end
+
+  describe '.caches' do
+    it 'returns the caches stored in the @caches ivar' do
+      query = build(:cached_query, count: nil, name: nil)
+      Readyset::Caches.cache(always: query.always, id: query.id) { query.text }
+
+      result = Readyset::Caches.caches
+
+      expect(result.size).to eq(1)
+      expect(result.first).to eq(query)
+    end
+  end
+end

spec/configuration_spec.rb

@@ -9,5 +9,12 @@
       config = Readyset::Configuration.new
       expect(config.shard).to eq(:readyset)
     end
+
+    it 'initializes migration_path to be db/readyset_caches.rb' do
+      config = Readyset::Configuration.new
+
+      expected = File.join(Rails.root, 'db/readyset_caches.rb')
+      expect(config.migration_path).to eq(expected)
+    end
   end
 end