feat: translate annotate to tags (#267)

* feat: translate annotate to tags

Adds support for adding request and transaction tags through the
annotate method.

Fixes #125

* fix: test and format errors

* chore: address review comments

examples/snippets/tags/README.md

@@ -0,0 +1,29 @@
+# Sample - Request tags and Transaction tags
+
+Queries can be annotated with request tags and transaction tags. These can be used to give
+you more insights in your queries and transactions.
+See https://cloud.google.com/spanner/docs/introspection/troubleshooting-with-tags for more
+information about request and transaction tags in Cloud Spanner.
+
+You can use the [`annotate`](https://api.rubyonrails.org/classes/ActiveRecord/QueryMethods.html#method-i-annotate)
+method in Ruby ActiveRecord to add a request and/or transaction tag to a __query__. You must
+prefix the annotation with either `request_tag:` or `transaction_tag:` to instruct the
+Cloud Spanner ActiveRecord provider to recognize the annotation as a tag.
+
+__NOTE:__ Ruby ActiveRecord does not add comments for `INSERT`, `UPDATE` and `DELETE` statements
+when you add annotations to a model. This means that these statements will not be tagged.
+
+Example:
+
+```ruby
+
+```
+
+The sample will automatically start a Spanner Emulator in a docker container and execute the sample
+against that emulator. The emulator will automatically be stopped when the application finishes.
+
+Run the application with the command
+
+```bash
+bundle exec rake run
+```

examples/snippets/tags/Rakefile

@@ -0,0 +1,13 @@
+# Copyright 2023 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+require_relative "../config/environment"
+require "sinatra/activerecord/rake"
+
+desc "Sample showing how to tag queries on Cloud Spanner with ActiveRecord."
+task :run do
+  Dir.chdir("..") { sh "bundle exec rake run[tags]" }
+end

examples/snippets/tags/application.rb

@@ -0,0 +1,32 @@
+# Copyright 2023 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+require "io/console"
+require_relative "../config/environment"
+require_relative "models/singer"
+require_relative "models/album"
+
+class Application
+  def self.run
+    puts ""
+    puts "Query all Albums and include a request tag"
+    albums = Album.annotate("request_tag: query-all-albums").all
+    puts "Queried #{albums.length} albums using a request tag"
+
+    puts ""
+    puts "Query all Albums in a transaction and include a request tag and a transaction tag"
+    Album.transaction do
+      albums = Album.annotate("request_tag: query-all-albums", "transaction_tag: sample-transaction").all
+      puts "Queried #{albums.length} albums using a request and a transaction tag"
+    end
+
+    puts ""
+    puts "Press any key to end the application"
+    STDIN.getch
+  end
+end
+
+Application.run

examples/snippets/tags/config/database.yml

@@ -0,0 +1,8 @@
+development:
+  adapter: spanner
+  emulator_host: localhost:9010
+  project: test-project
+  instance: test-instance
+  database: testdb
+  pool: 5
+  timeout: 5000

examples/snippets/tags/db/migrate/01_create_tables.rb

@@ -0,0 +1,21 @@
+# Copyright 2023 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+class CreateTables < ActiveRecord::Migration[6.0]
+  def change
+    connection.ddl_batch do
+      create_table :singers do |t|
+        t.string :first_name
+        t.string :last_name
+      end
+
+      create_table :albums do |t|
+        t.string :title
+        t.references :singer, index: false, foreign_key: true
+      end
+    end
+  end
+end

examples/snippets/tags/db/schema.rb

@@ -0,0 +1,31 @@
+# This file is auto-generated from the current state of the database. Instead
+# of editing this file, please use the migrations feature of Active Record to
+# incrementally modify your database, and then regenerate this schema definition.
+#
+# This file is the source Rails uses to define your schema when running `bin/rails
+# db:schema:load`. When creating a new database, `bin/rails db:schema:load` tends to
+# be faster and is potentially less error prone than running all of your
+# migrations from scratch. Old migrations may fail to apply correctly if those
+# migrations use external dependencies or application code.
+#
+# It's strongly recommended that you check this file into your version control system.
+
+ActiveRecord::Schema.define(version: 1) do
+  connection.start_batch_ddl
+
+  create_table "albums", id: { limit: 8 }, force: :cascade do |t|
+    t.string "title"
+    t.integer "singer_id", limit: 8
+  end
+
+  create_table "singers", id: { limit: 8 }, force: :cascade do |t|
+    t.string "first_name"
+    t.string "last_name"
+  end
+
+  add_foreign_key "albums", "singers"
+  connection.run_batch
+rescue
+  abort_batch
+  raise
+end

examples/snippets/tags/db/seeds.rb

@@ -0,0 +1,24 @@
+# Copyright 2023 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+require_relative "../../config/environment.rb"
+require_relative "../models/singer"
+require_relative "../models/album"
+
+first_names = ["Pete", "Alice", "John", "Ethel", "Trudy", "Naomi", "Wendy", "Ruben", "Thomas", "Elly"]
+last_names = ["Wendelson", "Allison", "Peterson", "Johnson", "Henderson", "Ericsson", "Aronson", "Tennet", "Courtou"]
+
+adjectives = ["daily", "happy", "blue", "generous", "cooked", "bad", "open"]
+nouns = ["windows", "potatoes", "bank", "street", "tree", "glass", "bottle"]
+
+5.times do
+  Singer.create first_name: first_names.sample, last_name: last_names.sample
+end
+
+20.times do
+  singer_id = Singer.all.sample.id
+  Album.create title: "#{adjectives.sample} #{nouns.sample}", singer_id: singer_id
+end

examples/snippets/tags/models/album.rb

@@ -0,0 +1,9 @@
+# Copyright 2023 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+class Album < ActiveRecord::Base
+  belongs_to :singer
+end

examples/snippets/tags/models/singer.rb

@@ -0,0 +1,9 @@
+# Copyright 2023 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+class Singer < ActiveRecord::Base
+  has_many :albums
+end

lib/active_record/connection_adapters/spanner/database_statements.rb

@@ -34,6 +34,10 @@ def execute sql, name = nil, binds = []
               selector = Google::Cloud::Spanner::Session.single_use_transaction staleness_hint.value
               binds.delete staleness_hint
             end
+            request_options = binds.find { |b| b.is_a? Google::Cloud::Spanner::V1::RequestOptions }
+            if request_options
+              binds.delete request_options
+            end
 
             log_args = [sql, name]
             log_args.concat [binds, type_casted_binds(binds)] if log_statement_binds
@@ -43,10 +47,11 @@ def execute sql, name = nil, binds = []
               ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
                 if transaction_required
                   transaction do
-                    @connection.execute_query sql, params: params, types: types
+                    @connection.execute_query sql, params: params, types: types, request_options: request_options
                   end
                 else
-                  @connection.execute_query sql, params: params, types: types, single_use_selector: selector
+                  @connection.execute_query sql, params: params, types: types, single_use_selector: selector,
+                                            request_options: request_options
                 end
               end
             end

lib/activerecord_spanner_adapter/connection.rb

@@ -200,7 +200,7 @@ def run_batch
 
     # DQL, DML Statements
 
-    def execute_query sql, params: nil, types: nil, single_use_selector: nil
+    def execute_query sql, params: nil, types: nil, single_use_selector: nil, request_options: nil
       if params
         converted_params, types = \
           Google::Cloud::Spanner::Convert.to_input_params_and_types(
@@ -214,15 +214,16 @@ def execute_query sql, params: nil, types: nil, single_use_selector: nil
       end
 
       selector = transaction_selector || single_use_selector
-      execute_sql_request sql, converted_params, types, selector
+      execute_sql_request sql, converted_params, types, selector, request_options
     end
 
-    def execute_sql_request sql, converted_params, types, selector
+    def execute_sql_request sql, converted_params, types, selector, request_options = nil
       res = session.execute_query \
         sql,
         params: converted_params,
         types: types,
         transaction: selector,
+        request_options: request_options,
         seqno: (current_transaction&.next_sequence_number)
       current_transaction.grpc_transaction = res.metadata.transaction \
           if current_transaction && res&.metadata&.transaction

lib/arel/visitors/spanner.rb

@@ -36,6 +36,7 @@ def compile node, collector = Arel::Collectors::SQLString.new
 
         if binds
           binds << collector.hints[:staleness] if collector.hints[:staleness]
+          binds << collector.hints[:request_options] if collector.hints[:request_options]
           [sql, binds]
         else
           sql
@@ -94,6 +95,26 @@ def visit_Arel_Nodes_OptimizerHints o, collector
         collector
       end
 
+      def visit_Arel_Nodes_Comment o, collector
+        o.values.each do |v|
+          if v.start_with?("request_tag:") || v.start_with?("transaction_tag:")
+            collector.hints[:request_options] ||= \
+              Google::Cloud::Spanner::V1::RequestOptions.new
+          end
+
+          if v.start_with? "request_tag:"
+            collector.hints[:request_options].request_tag = v.delete_prefix("request_tag:").strip
+            next
+          end
+          if v.start_with? "transaction_tag:"
+            collector.hints[:request_options].transaction_tag = v.delete_prefix("transaction_tag:").strip
+            next
+          end
+        end
+        # Also include the annotations as comments by calling the super implementation.
+        super
+      end
+
       def visit_Arel_Table o, collector
         return super unless collector.table_hints[o.name]
         if o.table_alias

test/activerecord_spanner_mock_server/spanner_active_record_with_mock_server_test.rb

@@ -831,6 +831,65 @@ def test_join_hint
       assert_equal sql, execute_sql_request.sql
     end
 
+    def test_query_annotate_request_tag
+      sql = "SELECT `singers`.* FROM `singers` /* request_tag: selecting all singers */"
+      @mock.put_statement_result sql, MockServerTests::create_random_singers_result(4)
+      Singer.annotate("request_tag: selecting all singers").all.each do |singer|
+        refute_nil singer.id, "singer.id should not be nil"
+      end
+      select_requests = @mock.requests.select { |req| req.is_a?(Google::Cloud::Spanner::V1::ExecuteSqlRequest) && req.sql == sql }
+      select_requests.each do |request|
+        assert request.request_options
+        assert_equal "selecting all singers", request.request_options.request_tag
+      end
+    end
+
+    def test_query_annotate_transaction_tag
+      sql = "SELECT `singers`.* FROM `singers` /* transaction_tag: selecting all singers */"
+      @mock.put_statement_result sql, MockServerTests::create_random_singers_result(4)
+      Singer.annotate("transaction_tag: selecting all singers").all.each do |singer|
+        refute_nil singer.id, "singer.id should not be nil"
+      end
+      select_requests = @mock.requests.select { |req| req.is_a?(Google::Cloud::Spanner::V1::ExecuteSqlRequest) && req.sql == sql }
+      select_requests.each do |request|
+        assert request.request_options
+        assert_equal "selecting all singers", request.request_options.transaction_tag
+      end
+    end
+
+    def test_query_annotate_request_and_transaction_tag
+      sql = "SELECT `singers`.* FROM `singers` /* transaction_tag: tx tag */ /* request_tag: req tag */"
+      @mock.put_statement_result sql, MockServerTests::create_random_singers_result(4)
+      Singer.annotate("transaction_tag: tx tag", "request_tag: req tag").all.each do |singer|
+        refute_nil singer.id, "singer.id should not be nil"
+      end
+      select_requests = @mock.requests.select { |req| req.is_a?(Google::Cloud::Spanner::V1::ExecuteSqlRequest) && req.sql == sql }
+      select_requests.each do |request|
+        assert request.request_options
+        assert_equal "req tag", request.request_options.request_tag
+        assert_equal "tx tag", request.request_options.transaction_tag
+      end
+    end
+
+    def test_query_annotate_request_and_transaction_tag_and_binds
+      sql = "SELECT `singers`.* FROM `singers` WHERE `singers`.`id` = @p1 /* transaction_tag: tx tag */ /* request_tag: req tag */ LIMIT @p2"
+      @mock.put_statement_result sql, MockServerTests::create_random_singers_result(1)
+      singer = Singer.annotate("transaction_tag: tx tag", "request_tag: req tag").find_by id: 1
+      assert singer
+
+      select_requests = @mock.requests.select { |req| req.is_a?(Google::Cloud::Spanner::V1::ExecuteSqlRequest) && req.sql == sql }
+      select_requests.each do |request|
+        assert request.request_options
+        assert_equal "req tag", request.request_options.request_tag
+        assert_equal "tx tag", request.request_options.transaction_tag
+        assert_equal 2, request.params.fields.length
+        assert_equal "1", request.params["p2"]
+        assert_equal "1", request.params["p1"]
+        assert_equal :INT64, request.param_types["p2"].code
+        assert_equal :INT64, request.param_types["p1"].code
+      end
+    end
+
     def test_insert_all
       values = [
         {id: 1, first_name: "Dave", last_name: "Allison"},

test/test_helper.rb

@@ -192,11 +192,12 @@ def release!
     end
 
     def execute_query sql, params: nil, types: nil, transaction: nil,
-                      partition_token: nil, seqno: nil
+                      partition_token: nil, request_options: nil, seqno: nil
       MockGoogleSpanner.last_executed_sqls OpenStruct.new(
         sql: sql, options: {
           params: params, types: types, transaction: transaction,
-          partition_token: partition_token, seqno: seqno
+          partition_token: partition_token, request_options: request_options,
+          seqno: seqno
         }
       )
       OpenStruct.new(rows: MockGoogleSpanner.mocked_result || [])