diff --git a/.gitignore b/.gitignore index d38669e..0aae463 100644 --- a/.gitignore +++ b/.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 diff --git a/lib/readyset.rb b/lib/readyset.rb index ce2af87..9ac971b 100644 --- a/lib/readyset.rb +++ b/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] 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 diff --git a/lib/readyset/caches.rb b/lib/readyset/caches.rb new file mode 100644 index 0000000..4aa3434 --- /dev/null +++ b/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 diff --git a/lib/readyset/configuration.rb b/lib/readyset/configuration.rb index aec2d73..038a952 100644 --- a/lib/readyset/configuration.rb +++ b/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 diff --git a/lib/readyset/query/cached_query.rb b/lib/readyset/query/cached_query.rb index a0f9f18..6db21d5 100644 --- a/lib/readyset/query/cached_query.rb +++ b/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 diff --git a/lib/readyset/railtie.rb b/lib/readyset/railtie.rb index 3e3e043..4bacec8 100644 --- a/lib/readyset/railtie.rb +++ b/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 diff --git a/lib/tasks/readyset.rake b/lib/tasks/readyset.rake new file mode 100644 index 0000000..bdf110f --- /dev/null +++ b/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 diff --git a/lib/templates/caches.rb.tt b/lib/templates/caches.rb.tt new file mode 100644 index 0000000..78a8c13 --- /dev/null +++ b/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 + diff --git a/readyset.gemspec b/readyset.gemspec index 43f12b7..6459e36 100644 --- a/readyset.gemspec +++ b/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' diff --git a/spec/caches_spec.rb b/spec/caches_spec.rb new file mode 100644 index 0000000..e9a80d0 --- /dev/null +++ b/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 diff --git a/spec/configuration_spec.rb b/spec/configuration_spec.rb index 71a8c71..e20d512 100644 --- a/spec/configuration_spec.rb +++ b/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 diff --git a/spec/factories/cached_query.rb b/spec/factories/cached_query.rb index b66950f..06677ba 100644 --- a/spec/factories/cached_query.rb +++ b/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 diff --git a/spec/factories/proxied_query.rb b/spec/factories/proxied_query.rb index 5e296fd..6799487 100644 --- a/spec/factories/proxied_query.rb +++ b/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 } diff --git a/spec/proxied_query_spec.rb b/spec/proxied_query_spec.rb index be2fff1..e381562 100644 --- a/spec/proxied_query_spec.rb +++ b/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) diff --git a/spec/rake_spec.rb b/spec/rake_spec.rb new file mode 100644 index 0000000..4f0abaf --- /dev/null +++ b/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