Ridgepole 含め Active Record がいい感じに Aurora DSQL で使える Connection Adapter

dsql_adapter.rb

require "active_record"
require "active_record/connection_adapters/postgresql_adapter"
require "socket"
require "aws-sdk-dsql"

module ActiveRecord::ConnectionAdapters::DSQL
  class SchemaCreation < ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaCreation
    private def visit_CreateIndexDefinition(o)
      super(o).gsub(/CREATE( UNIQUE)? INDEX/) { "CREATE#{$1} INDEX ASYNC" }
    end
  end
end

class ActiveRecord::ConnectionAdapters::DSQLAdapter < ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
  ADAPTER_NAME = "DSQL"

  def initialize(config)
    super(self.class.transform_config(config))
  end

  def schema_creation
    ActiveRecord::ConnectionAdapters::DSQL::SchemaCreation.new(self)
  end

  def primary_keys(table_name)
    query_values(<<~SQL, "SCHEMA")
      select
          kcu.column_name
      from
          information_schema.table_constraints tco
          join information_schema.key_column_usage kcu
              on kcu.constraint_name = tco.constraint_name
              and kcu.constraint_schema = tco.constraint_schema
              and kcu.constraint_name = tco.constraint_name
      where
          tco.constraint_type = 'PRIMARY KEY'
          and kcu.table_name = #{quote(table_name)}
      order by 1
    SQL
  end

  def default_index_type?(index)
    index.using == :btree_index
  end

  class << self
    def dbconsole(config, *args)
      new_config =
        case config
        when ActiveRecord::DatabaseConfigurations::HashConfig
          config_hash = transform_config(config.configuration_hash)
          ActiveRecord::DatabaseConfigurations::HashConfig.new(config.env_name, config.name, config_hash)
        end
      super(new_config, *args)
    end

    def transform_config(config)
      c = config.dup

      unless c[:password]
        password = generate_auth_token(
          endpoint: c[:host],
          region: c[:region],
          assume_role_arn: c.delete(:assume_role_arn),
          is_admin: c[:username] == "admin",
        )
      end

      {
        database: "postgres",
        sslmode: "require",
        advisory_locks: false,
        prepared_statements: false,
        password:,
      }.merge(c)
    end

    private def generate_auth_token(endpoint:, region:, assume_role_arn:, is_admin:)
      credentials =
        if assume_role_arn
          Aws::AssumeRoleCredentials.new(
            role_arn: assume_role_arn,
            role_session_name: "#{ENV['USER']}@#{Socket.gethostname}-#{Time.now.to_i}",
          )
        else
          Aws::CredentialProviderChain.new.resolve
        end
      token_generator = Aws::DSQL::AuthTokenGenerator.new(credentials:)
      region ||= endpoint&.then { _1.split(".")[-3] } || "us-east-1"
      generate_options = { endpoint:, region: }
      if is_admin
        token_generator.generate_db_connect_admin_auth_token(generate_options)
      else
        token_generator.generate_db_connect_auth_token(generate_options)
      end
    end
  end

  # Aurora DSQL does not support setting standard_conforming_strings in the connection parameters
  def set_standard_conforming_strings; end

  # Aurora DSQL does not support setting min_messages in the connection parameters
  def client_min_messages=(level); end

  # Aurora DSQL does not support running multiple DDL or DDL + DML statements in the same transaction
  def supports_ddl_transactions? = false
end

ActiveRecord::ConnectionAdapters.register "dsql", "ActiveRecord::ConnectionAdapters::DSQLAdapter"