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

    # Return true if the given database accepts connections
    def enabled?(database)
      !database.nil? or raise ArgumentError
      conn.value "select datallowconn from pg_catalog.pg_database where datname = '#{database}'"
    end

    # Ensure connections to the given database are enabled
    def enable(database)
      !database.nil? or raise ArgumentError
      conn.execute "alter database #{database} allow_connections = true"
    end

    # Ensure connections to the given database are disabled
    def disable(database)
      !database.nil? or raise ArgumentError
      conn.execute "alter database #{database} allow_connections = false"
    end

    # TODO: Why not let a nil database argument have the current database as default?

    # Terminate sessions in the database of the given users or of all users if
    # 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)
      users = Array(users).flatten
      !database.nil? or raise ArgumentError
      enabled = self.enabled?(database)

      case users
        when [];
          return
        when [nil]
          self.disable(database) if enabled
          users = self.list(database)
        else
          users = Array(users).flatten
      end
      pids = self.pids(database, users)

      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
      self.enable(database) if self.enabled?(database) != enabled
    end

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

    # Return true if session triggers are enabled. Triggers are enabled by
    # default by Postgres
    def triggers?() conn.value "select current_setting('session_replication_role') <> 'replica'" end

    # Enable session triggers
    def enable_triggers()
      conn.execute "set session session_replication_role = DEFAULT"
    end

    # Disable session triggers
    def disable_triggers()
      conn.execute "set session session_replication_role = replica"
    end

    # Execute block with session triggers on or off
    def triggers(on_off, &block)
      begin
        active = triggers?
        if on_off && !active
          enable_triggers
        elsif !on_off && active
          disable_triggers
        end
        yield
      ensure
        case active
          when true; enable_triggers if !triggers?
          when false; disable_triggers if triggers?
        end
      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