Vitess Rails patches

vitess_rails.rb

require 'active_record/connection_adapters/mysql2_adapter'

module VitessRails
  module Patches
    module PrependMasterKeyspace
      def table_name
        raise "A keyspace needs to be provided for internal rails metadata" unless @keyspace
        "#{@keyspace}.#{super}"
      end

      def keyspace=(value)
        @keyspace = value
      end
    end

    module SchemaCollationFix
      def schema_collation(column)
        if column.collation
          schema, unqualified_table_name = @connection.send(:extract_schema_qualified_name, table_name)
          @table_collation_cache ||= {}
          @table_collation_cache[table_name] ||=
            @connection.exec_query("SHOW TABLE STATUS FROM #{schema} LIKE #{@connection.quote(unqualified_table_name)}", "SCHEMA").first["Collation"]
          column.collation.inspect if column.collation != @table_collation_cache[table_name]
        end
      end
    end

    module IgnoreMetaTables
      def ignored?(table_name)
        _schema, name = @connection.send(:extract_schema_qualified_name, table_name)
        [ActiveRecord::Base.schema_migrations_table_name, ActiveRecord::Base.internal_metadata_table_name, ignore_tables].flatten.any? do |ignored|
          ignored === remove_prefix_and_suffix(name)
        end
      end
    end

    module DumpTableNamesWithSchemaName
      module AbstractPatches
        def tables
          res = query_values(data_source_sql(type: "BASE TABLE"), "SCHEMA")
          res = res.select { |x| x.start_with?("vt_") }
          # The wrangling below is necessary because vtgate returns very inconsistent
          # results depending on what conditions the query has. See PR that tried to
          # make Vitess "play nice" with this function:
          # https://github.com/vitessio/vitess/issues/6894
          # https://github.com/vitessio/vitess/pull/6932
          #
          # For now we hack around this, but we probably need to upstream this in Vitess
          res = res.map do |x|
            db_name, table_name = x.split(".")

            db_name_parts = db_name.split("_")
            db_name = db_name_parts[1..-2].join("_")

            "#{db_name}.#{table_name}"
          end.uniq

          res
        end
      end

      module MysqlPatches
        def data_source_sql(name = nil, type: nil)
          scope = quoted_scope(name, type: type)

          sql = +"SELECT CONCAT(table_schema, '.', table_name) FROM (SELECT * FROM information_schema.tables "
          if scope[:schema] == "database()"
            sql << ") _subquery"
          else
            sql << " WHERE table_schema = #{scope[:schema]}) _subquery"
          end
          if scope[:type] || scope[:name]
            conditions = []
            conditions << "_subquery.table_type = #{scope[:type]}" if scope[:type]
            conditions << "_subquery.table_name = #{scope[:name]}" if scope[:name]
            sql << " WHERE #{conditions.join(" AND ")}"
          end
          sql
        end
      end
    end
  end
end

ActiveSupport.on_load(:active_record) do
  ActiveRecord::SchemaMigration.singleton_class.prepend(VitessRails::Patches::PrependMasterKeyspace)
  ActiveRecord::InternalMetadata.singleton_class.prepend(VitessRails::Patches::PrependMasterKeyspace)
  ActiveRecord::SchemaDumper.prepend(VitessRails::Patches::IgnoreMetaTables)
  ActiveRecord::ConnectionAdapters::MySQL::SchemaDumper.prepend(VitessRails::Patches::SchemaCollationFix)
  ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(VitessRails::Patches::DumpTableNamesWithSchemaName::AbstractPatches)
  ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(VitessRails::Patches::DumpTableNamesWithSchemaName::MysqlPatches)
end