quota.py

"""mpsd-quota command line tool.

Provides the current use of storage for a user on /home and /scratch .

Options are available to

- display this for another user
- display in units of bytes (for debugging)
- show the version of the mpsd_hpc_tools package
- show some additional output

"""


import argparse
import os
import pathlib
import subprocess
import sys
import textwrap
from typing import Tuple
from mpsd_hpc_tools.lib import humanise_size
from mpsd_hpc_tools import __version__


def get_ceph_attribute(path: str, attribute: str) -> int:
    """Return output of 'getfattr' command.

    Parameters
    ----------
    path : str
        path to ceph managed volume, for example '/scratch/fangohr'
    attribute : str
        getfattrs attribute, for example 'ceph.dir.rbytes' or
        'ceph.quota.max_bytes'

    Returns
    -------
    value : int
        value returned by getfattrs command (only integers support)
    """
    # getfattr command should return a single integer number
    cmd = ["getfattr", "-n", attribute, "--only-values", "--absolute-names", path]
    output = subprocess.check_output(cmd).decode()
    return int(output)


def df_ceph(path: pathlib.Path) -> Tuple[int, int, int]:
    """
    Given a path (such as '/scratch/fangohr') return the number of bytes
    (size, used, available).

    Assumes that path_obj is provided by ceph.

    Example:

    >>> df_ceph("/scratch/fangohr")
    (25000000000000, 1780686359, 24998219313641)
    """
    try:
        used = get_ceph_attribute(path, "ceph.dir.rbytes")
        size = get_ceph_attribute(path, "ceph.quota.max_bytes")
    except Exception:
        used, size, avail = None, None, None
    else:
        avail = size - used
    return size, used, avail


def df_mountpoint(path: pathlib.Path) -> Tuple[int, int, int]:
    """
    Given a path (such as '/home/fangohr') return the number of bytes
    (size, used, available).

    Assumes that path_obj is a distinct mountpoint.

    Example:

    >>> df("/home/fangohr")
    (25000000000000, 1780686359, 24998219313641)
    """
    size, used, avail = None, None, None
    cmd = ["df", "--block-size=1", path]
    try:
        output = subprocess.check_output(cmd).decode()
    except Exception:
        return None, None, None
    for line in output.splitlines():
        df_fields = line.split()
        if df_fields[0] == "Filesystem":
            continue
        (size, used, avail) = [int(df_field) for df_field in df_fields[1:4]]
    return size, used, avail


def print_quota_line(location, size, used, avail, rel=None, raw=False):
    format_string = "{:20} {:>14} {:>14} {:>14} {:>14.4}%\n"
    if size is None:
        return
    if rel is None:
        rel = 100.0 * used / size
    if not raw:
        size, used, avail = map(humanise_size, [size, used, avail])
    sys.stdout.write(format_string.format(str(location), used, avail, size, rel))


def main(argv=None):
    """Main function for mpsd-quota command.

    Parameters
    ----------

    argv : List[str] or None
        Can provide list of arguments (reflecting sys.argv) for testing purposes
        The default value (None) will result in `sys.argv` being used.
    """
    parser = argparse.ArgumentParser(
        prog="mpsd-quota",
        description="Shows current usage of /home and /scratch",
        epilog=textwrap.dedent(
            """
            Output of mpsd-quota is in multiples of 1000 (i.e. decimal
            prefixes) and not 1024 (i.e. binary prefixes). If you want to
            compare output from mpsd-quota with output from "df -h", then
            use "df -h --si" to enforce use of decimal prefixes by df.
            """
        ),
    )

    # if executed on CI, "USER" may not be defined in environment. Then use
    # "UNKNOWN":
    userdefault = os.environ.get("USER", "UNKNOWN")

    parser.add_argument(
        type=str,
        default=userdefault,
        dest="user",
        nargs="?",
        help=f'quota for which user (by default "{userdefault}")',
    )
    parser.add_argument(
        "--version", "-V", help="display version and exit", action="store_true"
    )
    parser.add_argument(
        "--bytes",
        help="display storage in bytes (default: human readable)",
        action="store_true",
    )

    # if argv==None, then parse_args will use sys.argv (normal scenario) if we
    # are testing `main(argv)` we can parse custom arguments into `argv` (for
    # testing)
    args = parser.parse_args(argv)

    if args.version:
        print(__version__)
        sys.exit(0)

    # by default the user executing the command, or user provided on command line
    user = args.user

    # check user exists
    # print header
    print_quota_line("#location", "quota", "used", "avail", "use", raw=True)
    # print home quota
    homedir = pathlib.Path("/home") / user
    print_quota_line(homedir, *df_mountpoint(homedir), raw=args.bytes)

    # print scratch quota
    scratchdir = pathlib.Path("/scratch") / user
    print_quota_line(scratchdir, *df_ceph(scratchdir), raw=args.bytes)


if __name__ == "__main__":
    main()