Skip to content

Commit

Permalink
feat: Add AOST queries
Browse files Browse the repository at this point in the history
Add _As Of System Time_ support for models.

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

Fixes #281
  • Loading branch information
BuonOmo authored and rafiss committed Sep 19, 2023
1 parent bbb320a commit 1b17ecd
Show file tree
Hide file tree
Showing 8 changed files with 167 additions and 1 deletion.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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/)).
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
31 changes: 30 additions & 1 deletion lib/active_record/relation/query_methods_ext.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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?
Expand Down Expand Up @@ -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
28 changes: 28 additions & 0 deletions lib/arel/nodes/join_source_ext.rb
Original file line number Diff line number Diff line change
@@ -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
40 changes: 40 additions & 0 deletions test/cases/arel/nodes_test.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
# frozen_string_literal: true
# This file may be remove after next
# rails release.
# Whenever https://github.com/rails/rails/commit/8bd463c
# is part of a rails version.

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
57 changes: 57 additions & 0 deletions test/cases/relation/aost_test.rb
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions test/excludes/Arel/Nodes/TestNodes.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
exclude :test_every_arel_nodes_have_hash_eql_eqeq_from_same_class, "Overwitten, see https://github.com/rails/rails/issues/49274"

0 comments on commit 1b17ecd

Please sign in to comment.