readyset 0.1.1

data/lib/readyset/controller_extension.rb

@@ -0,0 +1,34 @@
+module Readyset
+  # Module providing controller extensions for routing ActiveRecord queries to a replica database.
+  module ControllerExtension
+    extend ActiveSupport::Concern
+
+    prepended do
+      # Sets up an `around_action` for a specified set of controller actions.
+      # This method is used to route the specified actions through Readyset,
+      # allowing ActiveRecord queries within those actions to be handled by a replica database.
+      #
+      # @example
+      #   route_to_readyset only: [:index, :show]
+      #   route_to_readyset :index
+      #   route_to_readyset except: :index
+      #   route_to_readyset :show, only: [:index, :show], if: -> { some_condition }
+      #
+      # @param args [Array<Symbol, Hash>] A list of actions and/or options dictating when the
+      # around_action should apply.
+      #   The options can include Rails' standard `:only`, `:except`, and conditionals like `:if`.
+      # @yield [_controller, action_block] An optional block that will execute around the actions.
+      #   Yields the block from the controller action.
+      # @yieldparam _controller [ActionController::Base] Param is unused.
+      # @yieldparam action_block [Proc] The block passed along with the action.
+      #
+      def self.route_to_readyset(*args, &block)
+        around_action(*args, *block) do |_controller, action_block|
+          Readyset.route do
+            action_block.call
+          end
+        end
+      end
+    end
+  end
+end

data/lib/readyset/error.rb

@@ -0,0 +1,3 @@
+module Readyset
+  module Error; end
+end

data/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 #{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

data/lib/readyset/health/healthchecker.rb

@@ -0,0 +1,127 @@
+require 'net/http'
+require 'uri'
+
+require 'readyset/health/healthchecks'
+
+module Readyset
+  module Health
+    # Processes the given exceptions to determine whether ReadySet is currently unhealthy. If
+    # ReadySet is indeed unhealthy, a background task is spawned that periodically checks
+    # ReadySet's health directly until a healthy state has been restored. While ReadySet is in an
+    # unhealthy state, `Healthchecker#healthy?` will return false.
+    class Healthchecker
+      UNHEALTHY_ERRORS = [::PG::UnableToSend, ::PG::ConnectionBad].freeze
+
+      def initialize(config, shard:)
+        @healthy = Concurrent::AtomicBoolean.new(true)
+        @healthcheck_interval = config.healthcheck_interval!
+        @healthchecks = Health::Healthchecks.new(shard: shard)
+        @lock = Mutex.new
+        @shard = shard
+        @window_counter = Readyset::Utils::WindowCounter.new(
+          window_size: config.error_window_size!,
+          time_period: config.error_window_period!,
+        )
+      end
+
+      # Returns true only if the connection to ReadySet is healthy. ReadySet's health is gauged by
+      # keeping track of the number of connection errors that have occurred over a given time
+      # period. If the number of errors in that time period exceeds the preconfigured threshold,
+      # ReadySet is considered to be unhealthy.
+      #
+      # @return [Boolean] whether ReadySet is healthy
+      def healthy?
+        healthy.true?
+      end
+
+      # Checks if the given exception is a connection error that occurred on a ReadySet connection,
+      # and if so, logs the error internally. If ReadySet is unhealthy, a background task is
+      # spawned that periodically tries to connect to ReadySet and check its status. When this task
+      # determines that ReadySet is healthy again, the task is shut down and the state of the
+      # healthchecker is switched back to "healthy".
+      #
+      # @param [Exception] the exception to be processed
+      def process_exception(exception)
+        is_readyset_connection_error = is_readyset_connection_error?(exception)
+        window_counter.log if is_readyset_connection_error
+
+        # We lock here to ensure that only one thread starts the healthcheck task
+        lock.lock
+        if healthy.true? && window_counter.threshold_crossed?
+          healthy.make_false
+          lock.unlock
+
+          logger.warn('ReadySet unhealthy: Routing queries to their original destination until ' \
+            'ReadySet becomes healthy again')
+
+          disconnect_readyset_pool!
+          task.execute
+        end
+      ensure
+        lock.unlock if lock.locked?
+      end
+
+      private
+
+      attr_reader :healthcheck_interval, :healthchecks, :healthy, :lock, :shard, :window_counter
+
+      def build_task
+        @task ||= Concurrent::TimerTask.new(execution_interval: healthcheck_interval) do |t|
+          if healthchecks.healthy?
+            # We disconnect the ReadySet connection pool here to ensure that any pre-existing
+            # connections to ReadySet are re-established. This fixes an issue where connections
+            # return "PQsocket() can't get socket descriptor" errors even after ReadySet comes
+            # back up. See this stackoverflow post for more details:
+            # https://stackoverflow.com/q/36582380
+            disconnect_readyset_pool!
+
+            # We need to disconnect the pool before making `healthy` true to ensure that, once we
+            # start routing queries back to ReadySet, they are using fresh connections
+            lock.synchronize { healthy.make_true }
+
+            logger.info('ReadySet healthy again')
+
+            # We clear out the window counter here to ensure that errors from ReadySet's previous
+            # unhealthy state don't bias the healthchecker towards determining that ReadySet is
+            # unhealthy after only a small number of new errors
+            window_counter.clear
+
+            t.shutdown
+          end
+        end
+
+        observer = Object.new.instance_eval do
+          def update(_time, _result, e)
+            logger.debug("ReadySet still unhealthy: #{e}") if e
+          end
+        end
+        task.add_observer(observer)
+
+        task
+      end
+
+      def disconnect_readyset_pool!
+        ActiveRecord::Base.connected_to(shard: shard) do
+          ActiveRecord::Base.connection_pool.disconnect!
+        end
+      end
+
+      def is_readyset_connection_error?(exception)
+        if exception.cause
+          is_readyset_connection_error?(exception.cause)
+        else
+          UNHEALTHY_ERRORS.any? { |e| exception.is_a?(e) } &&
+            exception.is_a?(Readyset::Error)
+        end
+      end
+
+      def logger
+        @logger ||= Rails.logger
+      end
+
+      def task
+        @task ||= build_task
+      end
+    end
+  end
+end

data/lib/readyset/health/healthchecks.rb

@@ -0,0 +1,41 @@
+module Readyset
+  module Health
+    # Represents healthchecks that are run against ReadySet to determine whether ReadySet is in a
+    # state where it can serve queries.
+    class Healthchecks
+      def initialize(shard:)
+        @shard = shard
+      end
+
+      # Checks if ReadySet is healthy by invoking `SHOW READYSET STATUS` and checking if
+      # ReadySet is connected to the upstream database.
+      #
+      # @return [Boolean] whether ReadySet is healthy
+      def healthy?
+        connection.execute('SHOW READYSET STATUS').any? do |row|
+          row['name'] == 'Database Connection' && row['value'] == 'Connected'
+        end
+      rescue
+        false
+      end
+
+      private
+
+      attr_reader :shard
+
+      def connection
+        @connection ||= ActiveRecord::Base.connected_to(shard: shard) do
+          ActiveRecord::Base.retrieve_connection
+        end
+
+        # We reconnect with each healthcheck to ensure that connection state is not cached across
+        # uses
+        @connection.reconnect!
+
+        @connection
+      rescue
+        false
+      end
+    end
+  end
+end

data/lib/readyset/model_extension.rb

@@ -0,0 +1,40 @@
+module Readyset
+  module ModelExtension
+    extend ActiveSupport::Concern
+
+    prepended do
+      require 'active_support'
+
+      # Defines a new class method on a model that wraps the ActiveRecord query in `body` in a call
+      # to `Readyset.route`.
+      #
+      # NOTE: `body` should consist of nothing other than an ActiveRecord query! If you need to run
+      # actions before or after the query execution, you should wrap the invocation of the named
+      # query in another method.
+      #
+      # @param [Symbol] name the name of the method that will be defined
+      # @param [Proc] body a lambda that wraps an ActiveRecord query
+      def self.readyset_query(name, body)
+        unless body.respond_to?(:call)
+          raise ArgumentError, 'The query body needs to be callable.'
+        end
+
+        if dangerous_class_method?(name)
+          raise ArgumentError, "You tried to define a ReadySet query named \"#{name}\" " \
+            "on the model \"#{self.name}\", but Active Record already defined " \
+            'a class method with the same name.'
+        end
+
+        if method_defined_within?(name, ActiveRecord::Relation)
+          raise ArgumentError, "You tried to define a ReadySet query named \"#{name}\" " \
+            "on the model \"#{self.name}\", but ActiveRecord::Relation already defined " \
+            'an instance method with the same name.'
+        end
+
+        singleton_class.define_method(name) do |*args|
+          Readyset.route { body.call(*args) }
+        end
+      end
+    end
+  end
+end

data/lib/readyset/query/cached_query.rb

@@ -0,0 +1,104 @@
+require 'active_model'
+
+require 'readyset/query/queryable'
+
+module Readyset
+  module Query
+    # Represents a query that is cached by ReadySet.
+    class CachedQuery
+      include ActiveModel::AttributeMethods
+      extend Queryable
+
+      attr_reader :id, :text, :name, :count, :always
+
+      # Returns all of the queries currently cached on ReadySet by invoking the `SHOW CACHES` SQL
+      # extension on ReadySet.
+      #
+      # @return [Array<CachedQuery>]
+      def self.all
+        super('SHOW CACHES')
+      end
+
+      # Drops all the caches that exist on ReadySet.
+      #
+      # @return [void]
+      def self.drop_all!
+        Readyset.raw_query('DROP ALL CACHES')
+
+        nil
+      end
+
+      # Returns the cached query with the given query ID by directly querying ReadySet. If a cached
+      # query with the given ID doesn't exist, this method raises a
+      # `Readyset::Query::NotFoundError`.
+      #
+      # @param [String] id the ID of the query to be searched for
+      # @return [CachedQuery]
+      # @raise [Readyset::Query::NotFoundError] raised if a cached query with the given ID cannot be
+      # found
+      def self.find(id)
+        super('SHOW CACHES WHERE query_id = ?', id)
+      end
+
+      # Constructs a new `CachedQuery` from the given attributes.
+      #
+      # @param [Hash] attributes the attributes from which the `CachedQuery` should be
+      # constructed
+      # @return [CachedQuery]
+      def initialize(id: nil, text:, name: nil, always: nil, count: nil)
+        @id = id
+        @text = text
+        @name = name
+        @always = always
+        @count = count
+      end
+
+      # Checks two queries for equality by comparing all of their attributes.
+      #
+      # @param [CachedQuery] the query against which `self` should be compared
+      # @return [Boolean]
+      def ==(other)
+        id == other.id &&
+          text == other.text &&
+          name == other.name &&
+          always == other.always
+      end
+
+      # Returns false if the cached query supports falling back to the upstream database and true
+      # otherwise.
+      #
+      # @return [Boolean]
+      def always?
+        always
+      end
+
+      # Drops the cache associated with this query.
+      #
+      # @return [void]
+      def drop!
+        Readyset.drop_cache!(name)
+        ProxiedQuery.find(id)
+      end
+
+      private
+
+      # Constructs a new `CachedQuery` from the given attributes. The attributes accepted
+      # by this method of this hash should correspond to the columns in the results returned by the
+      # `SHOW CACHES` ReadySet SQL extension.
+      #
+      # @param [Hash] attributes the attributes from which the `CachedQuery` should be
+      # constructed
+      # @return [CachedQuery]
+      def self.from_readyset_result(**attributes)
+        new(
+          id: attributes[:'query id'],
+          text: attributes[:'query text'],
+          name: attributes[:'cache name'],
+          always: attributes[:'fallback behavior'] != 'fallback allowed',
+          count: attributes[:count].to_i,
+        )
+      end
+      private_class_method :from_readyset_result
+    end
+  end
+end

data/lib/readyset/query/proxied_query.rb

@@ -0,0 +1,116 @@
+require 'active_model'
+
+require 'readyset/query'
+require 'readyset/query/queryable'
+
+module Readyset
+  module Query
+    # Represents an uncached query that has been proxied by ReadySet.
+    class ProxiedQuery
+      include ActiveModel::AttributeMethods
+      extend Queryable
+
+      # An error raised when a `ProxiedQuery` is expected to be supported but isn't.
+      class UnsupportedError < Query::BaseError
+        def to_s
+          "Query #{id} is unsupported"
+        end
+      end
+
+      attr_reader :id, :text, :supported, :count
+
+      # Returns all of the queries proxied by ReadySet that are not currently cached. This list is
+      # retrieved by invoking the `SHOW PROXIED QUERIES` SQL extension on ReadySet.
+      #
+      # @return [Array<ProxiedQuery>]
+      def self.all
+        super('SHOW PROXIED QUERIES')
+      end
+
+      # Creates a cache for every proxied query that is not already cached.
+      #
+      # @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 [Array<CachedQuery>] an array of the newly-created caches
+      def self.cache_all_supported!(always: false)
+        all.
+          select { |query| query.supported == :yes }.
+          map { |query| query.cache!(always: always) }
+      end
+
+      # Clears the list of proxied queries on ReadySet.
+      def self.drop_all!
+        Readyset.raw_query('DROP ALL PROXIED QUERIES')
+      end
+
+      # Returns the proxied query with the given query ID. The query is searched for by directly
+      # querying ReadySet. If a proxied query with the given ID doesn't exist, this method raises a
+      # `Readyset::Query::NotFoundError`.
+      #
+      # @param [String] id the ID of the query to be searched for
+      # @return [ProxiedQuery]
+      # @raise [Readyset::Query::NotFoundError] raised if a proxied query with the given
+      # ID cannot be found
+      def self.find(id)
+        super('SHOW PROXIED QUERIES WHERE query_id = ?', id)
+      end
+
+      # Constructs a new `ProxiedQuery` from the given attributes.
+      #
+      # @param [Hash] attributes the attributes from which the `ProxiedQuery` should be
+      # constructed
+      # @return [ProxiedQuery]
+      def initialize(id:, text:, supported:, count:)
+        @id = id
+        @text = text
+        @supported = supported
+        @count = count
+      end
+
+      # Checks two proxied queries for equality by comparing all of their attributes.
+      #
+      # @param [ProxiedQuery] the query against which `self` should be compared
+      # @return [Boolean]
+      def ==(other)
+        id == other.id &&
+          text == other.text &&
+          supported == other.supported
+      end
+
+      # Creates a cache on ReadySet for this query.
+      #
+      # @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 [CachedQuery] the newly-cached query
+      # @raise [ProxiedQuery::UnsupportedError] raised if this method is invoked on an
+      # unsupported query
+      def cache!(name: nil, always: false)
+        if supported == :unsupported
+          raise UnsupportedError, id
+        else
+          Readyset.create_cache!(id: id, name: name, always: always)
+          CachedQuery.find(id)
+        end
+      end
+
+      private
+
+      # Constructs a new `ProxiedQuery` from the given attributes. The attributes accepted
+      # by this method of this hash should correspond to the columns in the results returned by the
+      # `SHOW PROXIED QUERIES` ReadySet SQL extension.
+      #
+      # @param [Hash] attributes the attributes from which the `ProxiedQuery` should be constructed
+      # @return [ProxiedQuery]
+      def self.from_readyset_result(**attributes)
+        new(
+          id: attributes[:'query id'],
+          text: attributes[:'proxied query'],
+          supported: attributes[:'readyset supported'].to_sym,
+          count: attributes[:count].to_i,
+        )
+      end
+      private_class_method :from_readyset_result
+    end
+  end
+end

data/lib/readyset/query/queryable.rb

@@ -0,0 +1,23 @@
+module Readyset
+  module Query
+    module Queryable
+      private
+
+      def all(query) # :nodoc:
+        Readyset.raw_query(query).map do |result|
+          from_readyset_result(**result.symbolize_keys)
+        end
+      end
+
+      def find(query, id) # :nodoc:
+        result = Readyset.raw_query_sanitize(query, id).first
+
+        if result.nil?
+          raise NotFoundError, id
+        else
+          from_readyset_result(**result.to_h.symbolize_keys)
+        end
+      end
+    end
+  end
+end

data/lib/readyset/query.rb

@@ -0,0 +1,23 @@
+# lib/readyset/query.rb
+
+require 'active_model'
+
+module Readyset
+  module Query
+    class BaseError < StandardError
+      attr_reader :id
+
+      def initialize(id)
+        @id = id
+      end
+    end
+
+    # An error raised when a query with the given ID can't be found on the ReadySet
+    # instance.
+    class NotFoundError < BaseError
+      def to_s
+        "Query not found for ID #{id}"
+      end
+    end
+  end
+end

data/lib/readyset/railtie.rb

@@ -0,0 +1,68 @@
+# lib/readyset/railtie.rb
+
+require 'active_record/readyset_connection_handling'
+
+module Readyset
+  class Railtie < Rails::Railtie
+    initializer 'readyset.action_controller' do
+      ActiveSupport.on_load(:action_controller) do
+        prepend Readyset::ControllerExtension
+      end
+    end
+
+    initializer 'readyset.active_record' do |app|
+      ActiveSupport.on_load(:active_record) do
+        ActiveRecord::Base.prepend(Readyset::ModelExtension)
+        ActiveRecord::Base.extend(ActiveRecord::ReadysetConnectionHandling)
+
+        ActiveRecord::Relation.prepend(Readyset::RelationExtension)
+      end
+    end
+
+    # This Railtie sets up the ReadySet connection pools, which prevents users from needing to
+    # add a call to `ActiveRecord::Base.connects_to` in their ApplicationRecord class.
+    initializer 'readyset.connection_pools' do |app|
+      ActiveSupport.on_load(:after_initialize) do
+        shard = Readyset.config.shard
+
+        ActiveRecord::Base.connected_to(role: ActiveRecord.reading_role, shard: shard) do
+          ActiveRecord::Base.establish_connection(:readyset)
+        end
+
+        ActiveRecord::Base.connected_to(role: ActiveRecord.writing_role, shard: shard) do
+          ActiveRecord::Base.establish_connection(:readyset)
+        end
+      end
+    end
+
+    rake_tasks do
+      Dir[File.join(File.dirname(__FILE__), '../tasks/*.rake')].each { |f| load f }
+    end
+
+    initializer 'readyset.query_annotator' do |app|
+      setup_query_annotator
+    end
+
+    def setup_query_annotator
+      config.after_initialize do
+        if Rails.env.development? || Rails.env.test?
+          if Rails.configuration.active_record.query_log_tags_enabled
+            Rails.configuration.active_record.query_log_tags ||= []
+            Rails.configuration.active_record.query_log_tags << {
+              destination: ->(context) do
+                ActiveRecord::Base.connection_db_config.name
+              end,
+            }
+          else
+            Rails.logger.warn 'Query log tags are currently disabled.' \
+              'The ReadySet gem uses these tags to display information' \
+              'in the logs about whether a query was routed to ReadySet.' \
+              'It is highly recommended that you enable query log tags by setting' \
+              '`Rails.configuration.active_record.query_log_tags_enabled` to true to' \
+              'verify that queries are being routed to ReadySet as expected.'
+          end
+        end
+      end
+    end
+  end
+end

data/lib/readyset/relation_extension.rb

@@ -0,0 +1,29 @@
+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 issued to do the eager loading will not have caches created. Those queries must
+      # have their caches created separately.
+      #
+      # @param always [Boolean] whether the queries to this cache should always be served by
+      #                         ReadySet, preventing fallback to the uptream database
+      # @return [void]
+      def create_readyset_cache!(always: false)
+        Readyset.create_cache!(to_sql, always: always)
+      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