module Semian::Mysql2

def acquire_semian_resource(**)
  super
rescue ::Mysql2::Error => error
  if error.is_a?(PingFailure) || (!error.is_a?(::Mysql2::SemianError) && error.message.match?(CONNECTION_ERROR))
    semian_resource.mark_failed(error)
    error.semian_identifier = semian_identifier
  end
  raise
end

def connect(*args)
  acquire_semian_resource(adapter: :mysql, scope: :connection) { raw_connect(*args) }
end

def included(base)
  base.send(:alias_method, :raw_query, :query)
  base.send(:remove_method, :query)
  base.send(:alias_method, :raw_connect, :connect)
  base.send(:remove_method, :connect)
  base.send(:alias_method, :raw_ping, :ping)
  base.send(:remove_method, :ping)
end

def ping
  return false if closed?
  result = nil
  acquire_semian_resource(adapter: :mysql, scope: :ping) do
    result = raw_ping
    raise PingFailure, result.to_s unless result
  end
  result
rescue ResourceBusyError, CircuitOpenError, PingFailure
  false
end

def query(*args)
  if query_whitelisted?(*args)
    raw_query(*args)
  else
    acquire_semian_resource(adapter: :mysql, scope: :query) { raw_query(*args) }
  end
end

def query_whitelisted?(sql, *)
  QUERY_ALLOWLIST =~ sql
rescue ArgumentError
  # The above regexp match can fail if the input SQL string contains binary
  # data that is not recognized as a valid encoding, in which case we just
  # return false.
  return false unless sql.valid_encoding?
  raise
end

def raw_semian_options
  return query_options[:semian] if query_options.key?(:semian)
  query_options["semian"] if query_options.key?("semian")
end

def resource_exceptions
  EXCEPTIONS
end

def semian_identifier
  @semian_identifier ||= begin
    name = semian_options && semian_options[:name]
    unless name
      host = query_options[:host] || DEFAULT_HOST
      port = query_options[:port] || DEFAULT_PORT
      name = "#{host}:#{port}"
    end
    :"mysql_#{name}"
  end
end

def with_resource_timeout(temp_timeout)
  prev_read_timeout = @read_timeout
  begin
    # C-ext reads this directly, writer method will configure
    # properly on the client but based on my read--this is good enough
    # until we get https://github.com/brianmario/mysql2/pull/955 in
    @read_timeout = temp_timeout
    yield
  ensure
    @read_timeout = prev_read_timeout
  end
end