Generated schema lacks table locality data

[mmalek-sa]

When developing a multi-region solution, schema generated by the adapter doesn't include table locality settings. Also environment specific primary region data (e.g. create_enum "crdb_internal_region", ["us-east-1"]) leaks to the schema which is not desired.

To reproduce it:

# frozen_string_literal: true

require 'bundler/inline'

gemfile(true) do
  source 'https://rubygems.org'

  gem 'activerecord', '7.1.4.2'

  gem 'activerecord-cockroachdb-adapter'
end

require 'active_record'
require 'activerecord-cockroachdb-adapter'
require 'minitest/autorun'
require 'logger'

ActiveRecord::Base.establish_connection('cockroachdb://YOUR_USER:YOUR_PASSWORD@localhost:5432/bugs')
ActiveRecord::Base.logger = Logger.new(STDOUT)

class User < ActiveRecord::Base
end

class BugTest < ActiveSupport::TestCase
  def dump_schema
    stream = StringIO.new
    ActiveRecord::SchemaDumper.dump(
      ActiveRecord::Base.connection,
      stream
    )
    stream.string
  end

  def test_schema_migration_with_global_tables
    users_table_migration = Class.new(ActiveRecord::Migration::Current) do
      def up
        create_table :users, force: true do |t|
          t.string 'username'
        end

        exec_query("ALTER TABLE users SET LOCALITY GLOBAL", "Setting table locality")
      end

      def down
        say "users table locality migration is not reversible", true
      end
    end

    users_table_migration.migrate(:up)
    schema = dump_schema
    assert schema.include?("LOCALITY GLOBAL"), "Expected schema to include LOCALITY GLOBAL for users table, but it didn't"
    assert_not schema.include?("crdb_internal_region", "Schema should not include environment specific crdb_internal_region data")
  end

end

# generates this schema which lacks options: "LOCALITY GLOBAL" on the users table and includes environment specific
# primary region data crdb_internal_region which was set via ALTER DATABASE bugs SET PRIMARY REGION "us-east-1";
# 
#
# 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[7.1].define(version: 0) do
#   # Custom types defined in this database.
#   # Note that some types may not work with other database engines. Be careful if changing database.
#   create_enum "crdb_internal_region", ["us-east-1"]

#   create_table "users", id: :bigint, default: -> { "nextval('public.users_id_seq'::REGCLASS)" }, force: :cascade do |t|
#     t.string "username"
#   end
# end

Patching the table_options method seems to fix the table locality issue:

    # This method is called by the schema dump to get additional options to use when creating a table. The
    # current implementation in the CRDB adapter does not add the locality information to the tables, which
    # means they're all created with the default locality "REGIONAL BY TABLE". The monkey patch below uses the
    # table metadata to query the locality from the database and add it to the table options
    def table_options(table_name)
      options = {}
      table_metadata = tables_metadata[table_name]
      if table_metadata && table_metadata["locality"]
        case table_metadata["locality"].downcase
        when "global"
          options[:options] = "LOCALITY GLOBAL"
        when "regional by row"
          options[:options] = "LOCALITY REGIONAL BY ROW"
        end
      end
      options
    end