module PgConn
  # Schema methods
  class SessionMethods
    attr_reader :conn

    def initialize(conn)
      @conn = conn
    end

    # Returns a list of users connected to the given database. If database is
    # nil, it returns a list of database/username tuples for all connected users
    def list(database)
      if database
        conn.values "select usename from pg_stat_activity where datname = '#{database}'"
      else
        conn.tuples %(
          select datname, usename
          from pg_stat_activity 
          where datname is not null and usename is not null
        )
      end
    end

    # Terminate sessions in the database of the given users or of all users if
    # the users is nil. Note that 'terminate(database)' is a nop because the
    # absent users argument defaults to an empty list
    #
    # TODO: Make is possible to terminate a single session of a user with
    #       multiple sessions (is this ever relevant?)
    def terminate(database, *users)
      !database.nil? or raise ArgumentError
      users = Array(users).flatten
      case users
        when []; return
        when [nil]; users = list(database)
        else users = Array(users).flatten
      end
      pids = self.pids(database, users)
      return if pids.empty?
      pids_sql = pids.map { |pid| "(#{pid})" }.join(", ")
      conn.execute "select pg_terminate_backend(pid) from ( values #{pids_sql} ) as x(pid)"
    end

    def disable(database)
      !database.nil? or raise ArgumentError
      conn.execute "alter database #{database} allow_connections = false"
    end

    def enable(database)
      !database.nil? or raise ArgumentError
      conn.execute "alter database #{database} allow_connections = true"
    end

    # Run block without any connected users. Existing sessions are terminated
    def exclusive(database, &block) 
      !database.nil? or raise ArgumentError
      begin
        disable(database)
        users = list(database)
        terminate(database, users)
        yield
      ensure
        enable(database)
      end
    end

  private
    # Like #list but returns the PIDs of the users
    def pids(database, users)
      users ||= list(database)
      if !users.empty?
        users_sql = "(" + users.map { |user| "'#{user}'" }.join(", ") + ")"
        conn.values "select pid from pg_stat_activity where datname = '#{database}' and usename in #{users_sql}"
      else
        conn.values "select pid from pg_stat_activity where datname = '#{database}'"
      end
    end
  end
end