From bdb78a368eb62a8496ee93573c7a9b9ff5bd38b2 Mon Sep 17 00:00:00 2001 From: Ulysse Buonomo Date: Tue, 15 Aug 2023 09:56:51 -0500 Subject: [PATCH] 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 + .../cockroachdb/arel_tosql.rb | 8 +++ .../cockroachdb_adapter.rb | 1 + .../relation/query_methods_ext.rb | 31 +++++++++- lib/arel/nodes/join_source_ext.rb | 28 +++++++++ test/cases/arel/nodes_test.rb | 36 ++++++++++++ test/cases/relation/aost_test.rb | 57 +++++++++++++++++++ test/excludes/Arel/Nodes/TestNodes.rb | 1 + 8 files changed, 163 insertions(+), 1 deletion(-) create mode 100644 lib/arel/nodes/join_source_ext.rb create mode 100644 test/cases/arel/nodes_test.rb create mode 100644 test/cases/relation/aost_test.rb create mode 100644 test/excludes/Arel/Nodes/TestNodes.rb diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f81e116..78a4afef 100644 --- a/CHANGELOG.md +++ b/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/)). diff --git a/lib/active_record/connection_adapters/cockroachdb/arel_tosql.rb b/lib/active_record/connection_adapters/cockroachdb/arel_tosql.rb index 5c2529ea..c0ef7bf3 100644 --- a/lib/active_record/connection_adapters/cockroachdb/arel_tosql.rb +++ b/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 diff --git a/lib/active_record/connection_adapters/cockroachdb_adapter.rb b/lib/active_record/connection_adapters/cockroachdb_adapter.rb index 5c99302f..c46d4943 100644 --- a/lib/active_record/connection_adapters/cockroachdb_adapter.rb +++ b/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" diff --git a/lib/active_record/relation/query_methods_ext.rb b/lib/active_record/relation/query_methods_ext.rb index 42d3eb55..65c63698 100644 --- a/lib/active_record/relation/query_methods_ext.rb +++ b/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 diff --git a/lib/arel/nodes/join_source_ext.rb b/lib/arel/nodes/join_source_ext.rb new file mode 100644 index 00000000..aa95f57f --- /dev/null +++ b/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 diff --git a/test/cases/arel/nodes_test.rb b/test/cases/arel/nodes_test.rb new file mode 100644 index 00000000..4060f45b --- /dev/null +++ b/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 diff --git a/test/cases/relation/aost_test.rb b/test/cases/relation/aost_test.rb new file mode 100644 index 00000000..31b5fa28 --- /dev/null +++ b/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 diff --git a/test/excludes/Arel/Nodes/TestNodes.rb b/test/excludes/Arel/Nodes/TestNodes.rb new file mode 100644 index 00000000..0c05d8d3 --- /dev/null +++ b/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"