require "active_record/connection_adapters/mysql2_adapter"
require "activerecord-import/adapters/mysql2_adapter"

module ActiveRecord
  module ConnectionHandling
    # Provide an extensible connection adapter based on mysql2_connection implemented in
    #
    # In the configuration of shard.yml or database.yml, one may declare to use this custom adapter:
    #
    #   adapter: "proxy_mysql2"
    #
    # Note that we can't just call it "proxysql" because database tasks looks for the pattern /mysql/ in the name.
    def proxy_mysql2_connection(config)
      config = config.symbolize_keys

      config[:username] ||= "root"
      config[:flags] ||= 0

      if Mysql2::Client.const_defined?(:FOUND_ROWS)
        if config[:flags].kind_of?(Array)
          config[:flags].push("FOUND_ROWS".freeze)
        else
          config[:flags] |= Mysql2::Client::FOUND_ROWS
        end
      end

      client = Mysql2::Client.new(config)
      ConnectionAdapters::ProxyMysql2Adapter.new(client, logger, nil, config)
    rescue Mysql2::Error => error
      if error.message.include?("Unknown database")
        raise ActiveRecord::NoDatabaseError
      else
        raise
      end
    end
  end
end

# When acquiring a new connection in Rails5, it sets the `sql_mode` in the beginning of the query. ProxySQL ignores all the variables after setting the `sql_mode`.
# ```
# # rails5 - ProxySQL ignores the subsequent variables after sql_mode.
# SET @@SESSION.sql_mode = CONCAT(REPLACE(REPLACE(REPLACE(@@sql_mode, 'STRICT_TRANS_TABLES', ''), 'STRICT_ALL_TABLES', ''), 'TRADITIONAL', ''), ',NO_AUTO_VALUE_ON_ZERO'), @@SESSION.max_statement_time = 300000,
#     @@SESSION.sql_auto_is_null = 0, @@SESSION.wait_timeout = 2147483"
# ```
#
# As a result, the connection isn't set up properly in Rails5.
#
# Except the `sql_mode`, the rest of the variables are the same. In practice, `sql_mode = ''` in Rails4 falls back to the default of MySQL server. `sql_model = 'NO_AUTO_VALUE_ON_ZERO,NO_ENGINE_SUBSTITUTION'` in Rails5, practically, has the same effect.
#
# The adapter ensures that the session variables are set properly and keeps multiplexing enabled in ProxySQL.
class ActiveRecord::ConnectionAdapters::ProxyMysql2Adapter < ActiveRecord::ConnectionAdapters::Mysql2Adapter
  include ActiveRecord::Import::MysqlAdapter

  # Copied from https://github.com/rails/rails/blob/4-2-stable/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb#L866-L906
  # and https://github.com/rails/rails/blob/4-2-stable/activerecord/lib/active_record/connection_adapters/mysql2_adapter.rb#L255
  def configure_connection
    @connection.query_options.merge!(as: :array)

    variables = @config.fetch(:variables, {}).stringify_keys

    # By default, MySQL 'where id is null' selects the last inserted id.
    # Turn this off. http://dev.rubyonrails.org/ticket/6778
    variables["sql_auto_is_null"] = 0

    # Increase timeout so the server doesn't disconnect us.
    wait_timeout = self.class.type_cast_config_to_integer(@config[:wait_timeout])
    wait_timeout = 2147483 unless wait_timeout.is_a?(Integer)
    variables["wait_timeout"] = wait_timeout

    # Make MySQL reject illegal values rather than truncating or blanking them, see
    # http://dev.mysql.com/doc/refman/5.0/en/server-sql-mode.html#sqlmode_strict_all_tables
    # If the user has provided another value for sql_mode, don't replace it.
    unless variables.has_key?("sql_mode")
      variables["sql_mode"] = strict_mode? ? "STRICT_ALL_TABLES" : ""
    end

    # NAMES does not have an equals sign, see
    # http://dev.mysql.com/doc/refman/5.0/en/set-statement.html#id944430
    # (trailing comma because variable_assignments will always have content)
    if @config[:encoding]
      encoding = "NAMES #{@config[:encoding]}"
      encoding << " COLLATE #{@config[:collation]}" if @config[:collation]
      encoding << ", "
    end

    # Gather up all of the SET variables...
    variable_assignments = variables.map do |k, v|
      if v == ":default" || v == :default
        "@@SESSION.#{k} = DEFAULT" # Sets the value to the global or compile default
      elsif !v.nil?
        "@@SESSION.#{k} = #{quote(v)}"
      end
      # or else nil; compact to clear nils out
    end.compact.join(", ")

    # The ProxySQL config at /etc/proxysql.cnf keeps multiplexing enabled on queries tagged with /* keep_multiplexing_enabled */
    # Otherwise multiplexing is disabled because it sets the session variables, and the db connections are not shared between processes.
    # https://github.com/sysown/proxysql/wiki/Multiplexing#ad-hoc-enabledisable-of-multiplexing
    @connection.query("SET #{encoding} #{variable_assignments} /* keep_multiplexing_enabled */")
  end
end