Use composite_primary_key gem with spanner

test.rb

require "bundler/inline"

gemfile do
  source "https://rubygems.org"

  git_source(:github) { |repo| "https://github.com/#{repo}.git" }

  gem 'activerecord', '6.1.5'
  gem 'activerecord-spanner-adapter'
  gem 'composite_primary_keys'
end

require "active_record"
require "minitest/autorun"
require "activerecord-spanner-adapter"
require "google/cloud/spanner"
require "composite_primary_keys"

spanner = Google::Cloud::Spanner.new project: "test-project", emulator_host: "localhost:9010"
job = spanner.create_instance "test-instance",
                              name: "Test Instance",
                              config: "emulator-config",
                              nodes: 1
job.wait_until_done!

instance = spanner.instance "test-instance"
job = instance.create_database "testdb", statements: []
job.wait_until_done!

ActiveRecord::Base.establish_connection(project: "test-project", adapter: "spanner", instance: "test-instance", database: "testdb", emulator_host: "localhost:9010")

ActiveRecord::Schema.define do
  create_table :singers, force: true, id: false do |t|
    t.integer :singer_id, null: false, primary_key: true
    t.column :first_name, :string, limit: 200
    t.string :last_name
  end

  create_table :albums, force: true, id: false do |t|
    t.integer :album_id, null: false, primary_key: true
    t.interleave_in :singers, :cascade
    t.parent_key :singer_id
    t.string :title
  end
end

module PatchSpannerAdapter
  # Fix to support single key only
  def prefetch_primary_key? table_name = nil
    primary_keys(table_name).size == 1
  end

  # Fix primary_keys to includes parent key
  def primary_keys(table_name)
    primary_and_parent_keys(table_name)
  end
end
ActiveRecord::ConnectionAdapters::SpannerAdapter.prepend PatchSpannerAdapter

class Singer < ActiveRecord::Base
  has_many :albums
end

class Album < ActiveRecord::Base
  belongs_to :singer

  attribute :album_id, default: -> { next_sequence_value }
end

class QueryTracer
  attr_reader :store

  def initialize
    @store = []
  end

  def call(name, started, finished, unique_id, payload)
    @store << payload[:sql]
  end
end

class CompositePrimaryKeyTest < Minitest::Test
  def setup
    @tracer = QueryTracer.new

    ActiveSupport::Notifications.subscribe('sql.active_record', @tracer)
  end

  def test_query
    assert_equal Singer.primary_key, "singer_id"
    assert_equal Album.primary_key, ["singer_id", "album_id"]

    singer = Singer.new(first_name: "foo", last_name: "bar")
    album = singer.albums.build(title: "baz")

    assert singer.save
    assert album.persisted?
    assert_equal album.id, [album.singer_id, album.album_id]

    assert_equal @tracer.store[0], "BEGIN"
    assert_equal @tracer.store[1], "INSERT INTO `singers` (`first_name`, `last_name`, `singer_id`) VALUES (@p1, @p2, @p3)"
    assert_equal @tracer.store[2], "INSERT INTO `albums` (`album_id`, `singer_id`, `title`) VALUES (@p1, @p2, @p3)"
    assert_equal @tracer.store[3], "COMMIT"

    album = Album.find(album.id)
    assert album.update!(title: "baz1")

    assert_equal @tracer.store[4], "SELECT `albums`.* FROM `albums` WHERE `albums`.`singer_id` = #{album.singer_id} AND `albums`.`album_id` = #{album.album_id} LIMIT @p1"
    assert_equal @tracer.store[5], "BEGIN"
    assert_equal @tracer.store[6], "UPDATE `albums` SET `title` = @p1 WHERE `albums`.`singer_id` = @p2 AND `albums`.`album_id` = @p3"
    assert_equal @tracer.store[7], "COMMIT"
  end
end