feat: Add AOST queries

Add _As Of System Time_ support for models.

See https://www.cockroachlabs.com/docs/stable/as-of-system-time

Fixes #281

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Ongoing
 
+- Add support for [AOST](cockroachlabs.com/docs/stable/as-of-system-time) queries ([#284](https://github.com/cockroachdb/activerecord-cockroachdb-adapter/pull/284))
+
 ## 7.0.3 - 2023-08-23
 
 - Fix Multiple Database connections ([#283](https://github.com/cockroachdb/activerecord-cockroachdb-adapter/pull/)).

lib/active_record/connection_adapters/cockroachdb/arel_tosql.rb

@@ -22,6 +22,14 @@ module Arel # :nodoc:
   module Visitors # :nodoc:
     class CockroachDB < PostgreSQL  # :nodoc:
       include RGeo::ActiveRecord::SpatialToSql
+
+      def visit_Arel_Nodes_JoinSource(o, collector)
+        super
+        if o.aost
+          collector << " AS OF SYSTEM TIME '#{o.aost.iso8601}'"
+        end
+        collector
+      end
     end
   end
 end

lib/active_record/connection_adapters/cockroachdb_adapter.rb

@@ -1,5 +1,6 @@
 require "rgeo/active_record"
 
+require_relative "../../arel/nodes/join_source_ext"
 require "active_record/connection_adapters/postgresql_adapter"
 require "active_record/connection_adapters/cockroachdb/attribute_methods"
 require "active_record/connection_adapters/cockroachdb/column_methods"

lib/active_record/relation/query_methods_ext.rb

@@ -3,6 +3,23 @@
 module ActiveRecord
   class Relation
     module QueryMethodsExt
+      def aost!(time) # :nodoc:
+        unless time.nil? || time.is_a?(Time)
+          raise ArgumentError, "Unsupported argument type: #{time} (#{time.class})"
+        end
+
+        @aost = time
+        self
+      end
+
+      # Set system time for the current query. Using
+      # `.aost(nil)` resets.
+      #
+      # See cockroachlabs.com/docs/stable/as-of-system-time
+      def aost(time)
+        spawn.aost!(time)
+      end
+
       def from!(...) # :nodoc:
         @force_index = nil
         @index_hint = nil
@@ -59,8 +76,20 @@ def index_hint!(hint)
         self
       end
 
+      # TODO: reset or no reset?
+
+      def show_create
+        connection.execute("show create table #{connection.quote_table_name self.table_name}").first["create_statement"]
+      end
+
       private
 
+      def build_arel(...)
+        arel = super
+        arel.aost(@aost) if @aost.present?
+        arel
+      end
+
       def from_clause_is_a_table_name?
         # if empty, we are just dealing with the current table.
         return true if from_clause.empty?
@@ -94,5 +123,5 @@ def build_from_clause_with_hints
   # as ancestor. That is how active_record is doing is as well.
   #
   # @see https://github.com/rails/rails/blob/914130a9f/activerecord/lib/active_record/querying.rb#L23
-  Querying.delegate(:force_index, :index_hint, to: :all)
+  Querying.delegate(:force_index, :index_hint, :aost, to: :all)
 end

lib/arel/nodes/join_source_ext.rb

@@ -0,0 +1,28 @@
+module Arel
+  module Nodes
+    module JoinSourceExt
+      def initialize(...)
+        super
+        @aost = nil
+      end
+
+      def hash
+        [*super, aost].hash
+      end
+
+      def eql?(other)
+        super && aost == other.aost
+      end
+      alias_method :==, :eql?
+    end
+    JoinSource.attr_accessor :aost
+    JoinSource.prepend JoinSourceExt
+  end
+  module SelectManagerExt
+    def aost(time)
+      @ctx.source.aost = time
+      nil
+    end
+  end
+  SelectManager.prepend SelectManagerExt
+end

test/cases/arel/nodes_test.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require "cases/arel/helper"
+
+module Arel
+  module Nodes
+    class TestNodes < Arel::Test
+      def test_every_arel_nodes_have_hash_eql_eqeq_from_same_class
+        # #descendants code from activesupport
+        node_descendants = []
+        ObjectSpace.each_object(Arel::Nodes::Node.singleton_class) do |k|
+          next if k.respond_to?(:singleton_class?) && k.singleton_class?
+          node_descendants.unshift k unless k == self
+        end
+        node_descendants.delete(Arel::Nodes::Node)
+        node_descendants.delete(Arel::Nodes::NodeExpression)
+
+        default_hash_owner = Object.instance_method(:hash).owner
+
+        bad_node_descendants = node_descendants.reject do |subnode|
+          eqeq_owner = subnode.instance_method(:==).owner
+          eql_owner = subnode.instance_method(:eql?).owner
+          hash_owner = subnode.instance_method(:hash).owner
+
+          hash_owner != default_hash_owner &&
+              eqeq_owner == eql_owner &&
+              eqeq_owner == hash_owner
+        end
+
+        problem_msg = "Some subclasses of Arel::Nodes::Node do not have a" \
+            " #== or #eql? or #hash defined from the same class as the others"
+        assert_empty bad_node_descendants, problem_msg
+      end
+    end
+  end
+end

test/cases/relation/aost_test.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require "cases/helper_cockroachdb"
+require "models/post"
+require "models/comment"
+
+module CockroachDB
+  class AostTest < ActiveRecord::TestCase
+    def test_simple_aost
+      time = 2.days.ago
+      re_time = Regexp.quote(time.iso8601)
+      assert_match(/from "posts" as of system time '#{re_time}'/i, Post.aost(time).to_sql)
+      assert_match(/from "posts" as of system time '#{re_time}'/i, Post.where(name: "foo").aost(time).to_sql)
+    end
+
+    def test_reset_aost
+      time = 1.second.from_now
+      assert_match(/from "posts"\z/i, Post.aost(time).aost(nil).to_sql)
+    end
+
+    def test_aost_with_join
+      time = Time.now
+      assert_match(
+        /FROM "posts" INNER JOIN "comments" ON "comments"."post_id" = "posts"."id" AS OF SYSTEM TIME '#{Regexp.quote time.iso8601}'/,
+        Post.joins(:comments).aost(time).to_sql
+      )
+    end
+
+    def test_aost_with_subquery
+      time = 4.seconds.ago
+      assert_match(/from \(.*?\) subquery as of system time '#{Regexp.quote time.iso8601}'/i, Post.from(Post.where(name: "foo")).aost(time).to_sql)
+    end
+
+    def test_only_time_input
+      time = 1.second.ago
+      expected = "SELECT \"posts\".* FROM \"posts\" AS OF SYSTEM TIME '#{time.iso8601}'"
+      assert_equal expected, Post.aost(time).to_sql
+      assert_raises(ArgumentError) { Post.aost("no time") }
+      assert_raises(ArgumentError) { Post.aost(true) }
+    end
+  end
+
+  class AostNoTransactionTest < ActiveRecord::TestCase
+    # AOST per query is not compatible with transactions.
+    self.use_transactional_tests = false
+
+    def test_aost_with_multiple_queries
+      time = 1.second.ago
+      queries = capture_sql {
+        Post.aost(time).limit(2).find_each(batch_size: 1).to_a
+      }
+      queries.each do
+        assert_match /FROM \"posts\" AS OF SYSTEM TIME '#{Regexp.quote time.iso8601}'/, _1
+      end
+    end
+  end
+end

test/excludes/Arel/Nodes/TestNodes.rb

@@ -0,0 +1 @@
+exclude :test_every_arel_nodes_have_hash_eql_eqeq_from_same_class, "Overwitten, see https://github.com/rails/rails/issues/49274"