# rubocop:disable GitlabSecurity/PublicSend

module Gitlab
  module Database
    module LoadBalancing
      # Redirecting of ActiveRecord connections.
      #
      # The ConnectionProxy class redirects ActiveRecord connection requests to
      # the right load balancer pool, depending on the type of query.
      class ConnectionProxy
        attr_reader :load_balancer

        # These methods perform writes after which we need to stick to the
        # primary.
        STICKY_WRITES = %i(
          delete
          delete_all
          insert
          transaction
          update
          update_all
        ).freeze

        # hosts - The hosts to use for load balancing.
        def initialize(hosts = [])
          @load_balancer = LoadBalancer.new(hosts)
        end

        def select(*args)
          read_using_load_balancer(:select, args)
        end

        def select_all(arel, name = nil, binds = [])
          if arel.respond_to?(:locked) && arel.locked
            # SELECT ... FOR UPDATE queries should be sent to the primary.
            write_using_load_balancer(:select_all, [arel, name, binds],
                                      sticky: true)
          else
            read_using_load_balancer(:select_all, [arel, name, binds])
          end
        end

        STICKY_WRITES.each do |name|
          define_method(name) do |*args, &block|
            write_using_load_balancer(name, args, sticky: true, &block)
          end
        end

        # Delegates all unknown messages to a read-write connection.
        def method_missing(name, *args, &block)
          write_using_load_balancer(name, args, &block)
        end

        # Performs a read using the load balancer.
        #
        # name - The name of the method to call on a connection object.
        def read_using_load_balancer(name, args, &block)
          method = Session.current.use_primary? ? :read_write : :read

          @load_balancer.send(method) do |connection|
            connection.send(name, *args, &block)
          end
        end

        # Performs a write using the load balancer.
        #
        # name - The name of the method to call on a connection object.
        # sticky - If set to true the session will stick to the master after
        #          the write.
        def write_using_load_balancer(name, args, sticky: false, &block)
          result = @load_balancer.read_write do |connection|
            # Sticking has to be enabled before calling the method. Not doing so
            # could lead to methods called in a block still being performed on a
            # secondary instead of on a primary (when necessary).
            Session.current.write! if sticky

            connection.send(name, *args, &block)
          end

          result
        end
      end
    end
  end
end