repquota_gpfs.py

#!/usr/bin/env python

import os
import sys
import argparse
import subprocess
import time
import datetime
import pytz
import filelock
from getpass import getuser
import re
from bitmath import *
import shutil
import json
import logging

import osc.local_ldap
from osc.logs import setup_logging
from osc.gpfs import GPFS

DEFAULT_QUOTA_PATH = '/users/reporting/storage/quota'
LOCK_DIR = '/var/run/storage_reporting'

logger = logging.getLogger()

def fs_default(arg):
    filename=os.path.basename(__file__)
    m = re.search('^repquota_gpfs\.(.*)\.py$', filename)
    if m:
        fs = m.group(1)
        if arg == 'fs':
            return fs
        elif arg == 'mountpoint':
            return os.path.join('/fs', fs)
    return None

def get_xdmod_data():
    ldap_uris = osc.local_ldap.LocalLdap.ldap_server_uris()
    ldap_uri = ' '.join(ldap_uris)
    base_dn = osc.local_ldap.LocalLdap.ldap_base_dn()
    ldap = osc.local_ldap.LocalLdap(ldap_uri, False)
    ldap_groups = ldap.search_s(base_dn, osc.local_ldap.SCOPE_SUBTREE, '(objectClass=posixGroup)')
    ldap_group_gids = {}
    for g in ldap_groups:
        ldap_group_gids[g['gidNumber']] = g['cn']
    ldap_users = ldap.search_s(base_dn, osc.local_ldap.SCOPE_SUBTREE, '(objectClass=oscUser)')
    ldap_user_gids = {}
    for u in ldap_users:
        ldap_user_gids[u['uid']] = u['gidNumber']
    return ldap_group_gids, ldap_user_gids

def main():
    usage_examples = """

Usage Examples:

    Get fileset quotas:
        repquota_gpfs.project.py -t fileset
        %(prog)s -t fileset -f project -m /fs/project

    Get user quotas:
        repquota_gpfs.project.py -t user
        %(prog)s -t user -f project -m /fs/project

    Get fileset and user quotas
        repquota_gpfs.project.py -t fileset user

    Get group quotas:
        repquota_gpfs.project.py -t group
        %(prog)s -t group -f project -m /fs/project

    """
    default_fs = fs_default(arg='fs')
    default_quotafile = "gpfs.%s_quota.txt" % default_fs
    default_histdir = os.path.join(DEFAULT_QUOTA_PATH, 'historical', "gpfs.%s" % default_fs)
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter, epilog=usage_examples)
    parser.add_argument('-t', dest='types', help='type of quotas to query', nargs='+', choices=['user', 'group', 'fileset'], default=['fileset'])
    parser.add_argument('-f', dest='fs', help='file system (default: %(default)s)', default=default_fs)
    parser.add_argument('-m', dest='mountpoint', help='file system mountpoint (default: %(default)s)', default=fs_default(arg='mountpoint'))
    parser.add_argument('--histdir', help='historical directory (default: %(default)s)', default=default_histdir)
    parser.add_argument('--quotadir', help='quota output directory (default: %(default)s)', default=DEFAULT_QUOTA_PATH)
    parser.add_argument('--quotafile', help='quota file name (default: %(default)s)', default=default_quotafile)
    parser.add_argument('--cron', help='write output to files instead of stdout', action='store_true', default=False)
    parser.add_argument('--xdmod', help='run xdmod collection', action='store_true', default=False)
    args = parser.parse_args()

    setup_logging(logger=logger)

    gpfs = GPFS()
    output = []
    quota_data = {}
    quotas = {}
    now = datetime.datetime.now()
    now_str = now.strftime('%Y-%m-%dT%H:%M:%S')
    est = pytz.timezone('America/New_York')
    now_est = est.localize(now)
    utc = now_est.astimezone(pytz.utc)
    utc_str = utc.strftime('%Y-%m-%dT%H:%M:%SZ')
    date_str = "As of %s.000000" % now_str

    # Run with --xdmod or at 1AM
    if args.xdmod or (now.hour == 1 and now.minute >= 0 and now.minute <= 4):
        xdmod = True
    else:
        xdmod = False
    if xdmod:
        xdmod_quotadir = os.path.join(args.quotadir, 'xdmod', args.fs)
        xdmod_hist_quotadir = os.path.join(args.quotadir, 'xdmod', 'historical', args.fs)
        if not os.path.isdir(xdmod_quotadir):
            os.mkdir(xdmod_quotadir)
        if not os.path.isdir(xdmod_hist_quotadir):
            os.makedirs(xdmod_hist_quotadir)

        ldap_group_gids, ldap_user_gids = get_xdmod_data()

    # BEGIN: data collection
    for _type in args.types:
        if _type == 'user':
            type_str = 'userid'
        elif _type == 'group':
            type_str = 'group'
        elif _type == 'fileset':
            type_str = 'project/group'

        results = gpfs.mmrepquota(fs=args.fs, t=_type)
        if results is None:
            logger.error("No mmrepquota data returned")
            sys.exit(1)

        for data in results:
            name = data['name']
            # Empty line from mmrepquota
            if not name:
                continue
            filesetname = data.get('filesetname', None)
            if type_str == 'userid':
                if filesetname and filesetname != 'root':
                    path = os.path.join(args.mountpoint, str(filesetname))
                else:
                    path = args.mountpoint
            else:
                path = args.mountpoint
            block_used = KiB(data['blockUsage'])
            block_used_str = block_used.to_GiB().format("{value:.0f} {unit}")
            block_quota = KiB(data['blockQuota'])
            block_quota_str = block_quota.to_GiB().format("{value:.0f} {unit}")
            output.append("%s %s %s on %s used %s of quota %s and %s files of quota %s files" % (
                date_str, type_str, name, path, block_used_str, block_quota_str, data['filesUsage'], data['filesQuota']
            ))    
            data = {
                'name': name,
                'path': path,
                'block_usage': block_used,
                'block_limit': block_quota,
                'file_usage': data['filesUsage'],
                'file_limit': data['filesQuota'],
                'parent': filesetname,
                'type': _type,
            }
            if name not in quota_data:
                quota_data[name] = []
            quota_data[name].append(data)
    # END: data collection

    quotas['version'] = 1
    quotas['timestamp'] = int(time.mktime(now.timetuple()))
    quotas['quotas'] = []
    quotas['quotas_other'] = []
    xdmod_quotas = []
    for name, datas in quota_data.iteritems():
        for d in datas:
            data = {}
            data['path'] = d['path']
            if d['type'] == 'user':
                key = 'quotas'
                parent = quota_data[d['parent']][0]
                data['user'] = name
                data['type'] = 'fileset'
                data['block_usage'] = int(d['block_usage'])
                data['total_block_usage'] = int(parent['block_usage'])
                data['block_limit'] = int(parent['block_limit'])
                data['file_usage'] = d['file_usage']
                data['total_file_usage'] = parent['file_usage']
                data['file_limit'] = parent['file_limit']
                if xdmod:
                    gid = ldap_user_gids.get(name, 65534)
                    primary_group = ldap_group_gids.get(gid, None)
                    if str(name) == 'root':
                        primary_group = 'root'
                    if not primary_group:
                        primary_group = 'nfsnobody'
                    xdmod_quota = {
                        'resource': args.fs,
                        'mountpoint': d['path'],
                        'user': str(name),
                        'pi': primary_group,
                        'dt': utc_str,
                        'soft_threshold': int(parent['block_limit'].bytes),
                        'hard_threshold': int(parent['block_limit'].bytes),
                        'file_count': d['file_usage'],
                        'logical_usage': int(d['block_usage'].bytes),
                        'physical_usage': int(d['block_usage'].bytes),
                    }
                    if not str(name).isdigit():
                        xdmod_quotas.append(xdmod_quota)
            else:
                key = 'quotas_other'
                data['group'] = name
                data['total_block_usage'] = int(d['block_usage'])
                data['block_limit'] = int(d['block_limit'])
                data['total_file_usage'] = d['file_usage']
                data['file_limit'] = d['file_limit']
            quotas[key].append(data)

    if args.cron:
        hist_latest_file = os.path.join(args.histdir, 'latest')
        hist_date_file = os.path.join(args.histdir, now_str)
        latest_file = os.path.join(args.quotadir, args.quotafile)
        with open(hist_latest_file, 'w') as f:
            for o in output:
                f.write("%s\n" % o)
        shutil.copy2(hist_latest_file, latest_file)
        # Generate JSON
        latest_file_path, extension = os.path.splitext(latest_file)
        latest_json_file = "%s.json" % latest_file_path
        latest_json_file_new = "%s.new" % latest_json_file
        latest_json_file_old = "%s.old" % latest_json_file
        with open(latest_json_file_new, 'w') as f:
            json.dump(quotas, f, sort_keys=True, indent=4)
        # Move .json to .json.old then .json.new to .json
        if os.path.isfile(latest_json_file):
            shutil.move(latest_json_file, latest_json_file_old)
        shutil.move(latest_json_file_new, latest_json_file)
        # Copy historical files
        json_hist_file = os.path.join(args.histdir, "%s.json" % now_str)
        json_hist_latest = os.path.join(args.histdir, 'latest.json')
        shutil.copy2(latest_json_file, json_hist_latest)
        # Only at 4AM and 4PM
        if now.hour in [4, 16]:
            if now.minute >= 0 and now.minute <= 4:
                shutil.copy2(hist_latest_file, hist_date_file)
                shutil.copy2(latest_json_file, json_hist_file)
        if xdmod:
            # Write XDMoD file
            xdmod_file = os.path.join(xdmod_quotadir, '%s.json' % args.fs)
            xdmod_hist_file = os.path.join(xdmod_hist_quotadir, "%s.json" % now_str)
            with open(xdmod_file, 'w') as f:
                json.dump(xdmod_quotas, f, sort_keys=True, indent=4)
            shutil.copy2(xdmod_file, xdmod_hist_file)
    else:
        for o in output:
            print o
    
if __name__ == '__main__':
    lockfile = os.path.join(LOCK_DIR, "%s-%s.lock" % (getuser(), os.path.splitext(os.path.basename(__file__))[0]))
    flock = filelock.FileLock(lockfile, 10)
    try:
        flock.acquire()
    except filelock.Timeout:
        sys.stderr.write("Failed to acquire lock on %s\n" % lockfile)
        sys.exit(1)
    try:
        main()
    finally:
        flock.release()