Add migration file (#61)

This commit adds two rake tasks:
- `readyset:caches:dump`, which reads the caches that exist on the
  connected ReadySet instance and dumps them to a Ruby DSL file
- `readyset:caches:migrate`, which reads the aforementioned Ruby file,
  create the caches from the file that don't exist on ReadySet, and
  drops the caches on ReadySet that don't exist in the file

Closes #9

.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,5 +1,6 @@
 # lib/readyset.rb
 
+require 'readyset/caches'
 require 'readyset/configuration'
 require 'readyset/controller_extension'
 require 'readyset/model_extension'
@@ -101,7 +102,7 @@ def self.explain(query)
   # @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,19 @@
+module Readyset
+  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/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/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

spec/factories/cached_query.rb

@@ -8,4 +8,14 @@
 
     initialize_with { new(**attributes) }
   end
+
+  factory :cached_query_2, class: 'Readyset::Query::CachedQuery' do
+    id { 'q_8892818e62c34ecd' }
+    count { 5 }
+    text { 'SELECT * FROM "t" WHERE ("y" = $1)' }
+    name { 'q_8892818e62c34ecd' }
+    always { true }
+
+    initialize_with { new(**attributes) }
+  end
 end

spec/factories/proxied_query.rb

@@ -6,11 +6,11 @@
     text { 'SELECT * FROM "t" WHERE ("x" = $1)' }
     supported { :yes }
 
-    factory :pending_query do
+    factory :pending_proxied_query do
       supported { :pending }
     end
 
-    factory :unsupported_query do
+    factory :unsupported_proxied_query do
       id { 'q_f9bfc11a043b2f75' }
       text { 'SHOW TIME ZONE' }
       supported { :unsupported }

spec/proxied_query_spec.rb

@@ -46,8 +46,8 @@
     end
     let(:unsupported_or_pending_queries) do
       [
-        build(:unsupported_query),
-        build(:pending_query),
+        build(:unsupported_proxied_query),
+        build(:pending_proxied_query),
       ]
     end
 
@@ -179,7 +179,7 @@
     context 'when the query is unsupported' do
       subject { query.cache! }
 
-      let(:query) { build(:unsupported_query) }
+      let(:query) { build(:unsupported_proxied_query) }
 
       it 'raises a ProxiedQuery::UnsupportedError' do
         expect { subject }.to raise_error(Readyset::Query::ProxiedQuery::UnsupportedError)

spec/rake_spec.rb

@@ -0,0 +1,120 @@
+require 'colorize'
+require 'rake'
+require 'spec_helper'
+
+load './lib/tasks/readyset.rake'
+
+RSpec.describe 'readyset.rake' do
+  before do
+    Rake::Task.define_task(:environment)
+  end
+
+  describe 'readyset' do
+    describe 'caches' do
+      describe 'dump' do
+        it 'dumps the current set of caches to a migration file' do
+          # Setup
+          allow(Readyset::Query::CachedQuery).to receive(:all).
+            and_return([build(:cached_query), build(:cached_query_2)])
+
+          # Execute
+          Rake::Task['readyset:caches:dump'].execute
+
+          # Verify
+          load './spec/internal/db/readyset_caches.rb'
+          subclasses = Readyset::Caches.subclasses
+          expect(subclasses.size).to eq(1)
+
+          caches = subclasses.first.caches
+          expect(caches.size).to eq(2)
+          expect(caches).to include(build(:cached_query, count: nil, name: nil))
+          expect(caches).to include(build(:cached_query_2, count: nil, name: nil))
+        end
+      end
+
+      describe 'migrate' do
+        after(:each) do
+          if File.exist?('./spec/internal/db/readyset_caches.rb')
+            File.delete('./spec/internal/db/readyset_caches.rb')
+          end
+        end
+
+        context "when the migration file contains caches that don't exist on ReadySet" do
+          it "creates the caches in the migration file that don't exist on ReadySet" do
+            # Setup
+            cache_to_create = build(:cached_query_2)
+            generate_migration_file([build(:cached_query), cache_to_create])
+
+            allow(Readyset::Query::CachedQuery).to receive(:all).and_return([build(:cached_query)])
+            allow(Readyset).to receive(:create_cache!).with(id: cache_to_create.id)
+            allow(STDIN).to receive(:gets).and_return("y\n")
+
+            # Execute
+            Rake::Task['readyset:caches:migrate'].execute
+
+            # Verify
+            expect(Readyset).to have_received(:create_cache!).with(id: cache_to_create.id)
+          end
+
+          it 'prints the expected output' do
+            # Setup
+            cache_to_create = build(:cached_query_2)
+            generate_migration_file([build(:cached_query), cache_to_create])
+
+            allow(Readyset::Query::CachedQuery).to receive(:all).and_return([build(:cached_query)])
+            allow(Readyset).to receive(:create_cache!).with(id: cache_to_create.id)
+            allow(STDIN).to receive(:gets).and_return("y\n")
+
+            # Execute + Verify
+            expected_message = "#{'Dropping'.red} 0 caches and #{'creating'.green} 1 caches. " \
+              'Continue? (y/n) '
+            expect { Rake::Task['readyset:caches:migrate'].execute }.to output(expected_message).
+              to_stdout
+          end
+        end
+
+        context "when ReadySet has caches that don't exist in the migration file" do
+          it 'drops the caches that exist on ReadySet that are not in the migration file' do
+            # Setup
+            generate_migration_file([build(:cached_query)])
+
+            cache_to_drop = build(:cached_query_2)
+            allow(Readyset::Query::CachedQuery).to receive(:all).
+              and_return([build(:cached_query), cache_to_drop])
+            allow(Readyset).to receive(:drop_cache!).with(name_or_id: cache_to_drop.id)
+            allow(STDIN).to receive(:gets).and_return("y\n")
+
+            # Execute
+            Rake::Task['readyset:caches:migrate'].execute
+
+            # Verify
+            expect(Readyset).to have_received(:drop_cache!).with(name_or_id: cache_to_drop.id)
+          end
+
+          it 'prints the expected output' do
+            # Setup
+            generate_migration_file([build(:cached_query)])
+
+            cache_to_drop = build(:cached_query_2)
+            allow(Readyset::Query::CachedQuery).to receive(:all).
+              and_return([build(:cached_query), cache_to_drop])
+            allow(Readyset).to receive(:drop_cache!).with(name_or_id: cache_to_drop.id)
+            allow(STDIN).to receive(:gets).and_return("y\n")
+
+            # Execute + Verify
+            expected_message = "#{'Dropping'.red} 1 caches and #{'creating'.green} 0 caches. " \
+              'Continue? (y/n) '
+            expect { Rake::Task['readyset:caches:migrate'].execute }.to output(expected_message).
+              to_stdout
+          end
+        end
+
+        def generate_migration_file(caches)
+          allow(Readyset::Query::CachedQuery).to receive(:all).and_return(caches)
+          Rake::Task['readyset:caches:dump'].execute
+          allow(Readyset::Query::CachedQuery).to receive(:all).and_call_original
+        end
+      end
+    end
+  end
+end