depot_verify_chunks.py #6

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# ==============================================================================
# Copyright and license info is available in the LICENSE file included with
# the Server Deployment Package (SDP), and also available online:
# https://swarm.workshop.perforce.com/projects/perforce-software-sdp/view/main/LICENSE
# ------------------------------------------------------------------------------

"""
NAME:
    depot_verify_chunks.py

DESCRIPTION:
    This script prints out a list of directories to verify that don't exceed a specific
    size limit, e.g. 100M or 1G.

    It is intended as an optional input to p4verify.py script or to be run using GNU parallel.

    It walks down the tree of specified paths using the p4 dirs/p4 sizes -az commands to identify
    appropriate sub-trees.

    The output is a list of depot paths:
        //depot/...
        //big_depot/subdir1/...
        //big_depot/subdir2/...
        //big_depot/subdir3/...
        //big_depot/*

    You can run:

        depot_verify_chunks.py -m 1G //some/path/... //another_depot/...

    With parallel (see https://www.gnu.org/software/parallel/):

        depot_verify_chunks.py -v ERROR -m 1G //some/path/... | parallel p4 verify -qzt {} >> verify.out
        nohup depot_verify_chunks.py -v ERROR -m 1G //some/path/... | parallel p4 verify -qzt {} >> verify.out 2>&1 &

"""

# Python 2.7/3.3 compatibility.
from __future__ import print_function

import sys
import os
import textwrap
import argparse
import logging
import P4
import re
from collections import OrderedDict

script_name = os.path.basename(os.path.splitext(__file__)[0])
LOGDIR = os.getenv('LOGS', '/p4/1/logs')

DEFAULT_LOG_FILE = "log-%s.log" % script_name
if os.path.exists(LOGDIR):
    DEFAULT_LOG_FILE = os.path.join(LOGDIR, "%s.log" % script_name)
DEFAULT_VERBOSITY = 'DEBUG'
LOGGER_NAME = 'verify_dir_list'

def parse_human_fmt(opt):
    "Parses 1, 1K, 2M, 3.1G etc and returns int"
    units = "kmgt"
    opt = opt.lower()
    m = re.match(r"^([0-9]*\.{0,1}[0-9]*)([kmgt]*)$", opt)
    if not m:
        raise TypeError("Argument must be a human readable size number, e.g. 1, 1K, 2M, 3.1G")
    num = float(m.group(1))
    unit = m.group(2)
    if unit and unit in units:
        multiplier = 1.0
        for u in ['k','m','g','t']:
            multiplier *= 1024.0
            if u == unit:
                break
        num *= multiplier
    return int(num)

class depot_verify_chunks(object):
    """See module doc string for details"""

    def __init__(self, *args, **kwargs):
        self.parse_args(__doc__, args)

    def parse_args(self, doc, args):
        """Common parsing and setting up of args"""
        desc = textwrap.dedent(doc)
        parser = argparse.ArgumentParser(
            formatter_class=argparse.RawDescriptionHelpFormatter,
            description=desc,
            epilog="Copyright (c) 2008-2018 Perforce Software, Inc."
        )
        self.add_parse_args(parser)
        self.options = parser.parse_args(args=args)
        self.init_logger()
        self.logger.debug("Command Line Options: %s\n" % self.options)
        self.options.max_size = parse_human_fmt(self.options.max_size)  # May raise error

    def add_parse_args(self, parser, default_log_file=None, default_verbosity=None):
        """Default trigger arguments - common to all triggers
        :param default_verbosity:
        :param default_log_file:
        :param parser:
        """
        if not default_log_file:
            default_log_file = DEFAULT_LOG_FILE
        if not default_verbosity:
            default_verbosity = DEFAULT_VERBOSITY
        parser.add_argument('-p', '--port', default=None,
                            help="Perforce server port. Default: $P4PORT")
        parser.add_argument('-u', '--user', default=None, help="Perforce user. Default: $P4USER")
        parser.add_argument('-L', '--log', default=default_log_file, help="Default: " + default_log_file)
        parser.add_argument('-m', '--max-size', help="Max size of each depot chunk. Can be specified in K/M/G/T, e.g. 1G")
        parser.add_argument('path', nargs='*', help="Perforce depot path(s)")
        parser.add_argument('-v', '--verbosity',
                            nargs='?',
                            const="INFO",
                            default=default_verbosity,
                            choices=('DEBUG', 'WARNING', 'INFO', 'ERROR', 'FATAL'),
                            help="Output verbosity level. Default is: " + default_verbosity)

    def init_logger(self, logger_name=None):
        if not logger_name:
            logger_name = LOGGER_NAME
        self.logger = logging.getLogger(logger_name)
        self.logger.setLevel(self.options.verbosity)
        logformat = '%(levelname)s %(asctime)s %(filename)s %(lineno)d: %(message)s'
        logging.basicConfig(format=logformat, filename=self.options.log, level=self.options.verbosity)
        formatter = logging.Formatter('%(message)s')
        ch = logging.StreamHandler(sys.stderr)
        ch.setLevel(logging.INFO)
        ch.setFormatter(formatter)
        self.logger.addHandler(ch)

    def get_chunks(self, results, startdir, max_size):
        # Use a stack for iterative processing to avoid recursion depth limits
        if startdir.endswith("/..."):
            startdir = startdir[:-4]

        stack = [startdir]
        visited = set()

        while stack:
            current_dir = stack.pop()

            # Remove trailing /... if present
            if current_dir.endswith("/..."):
                current_dir = current_dir[:-4]

            # Guard against cycles: skip already-visited directories.
            # This prevents infinite loops when p4 dirs returns a path that
            # resolves back to the current directory (e.g. a result with a
            # trailing /... that strips to the parent path).
            if current_dir in visited:
                continue
            visited.add(current_dir)

            # Check size of current directory
            sizes = self.p4.run_sizes("-az", "%s/..." % current_dir)
            if int(sizes[0]['fileSize']) <= max_size:
                # Directory is small enough, add it to results
                results["%s/..." % current_dir] = 1
                continue

            # Directory is too large, need to subdivide
            subdirs = self.p4.run_dirs('-D', "%s/*" % current_dir)

            # Only queue subdirectories not already visited
            new_subdirs = [d['dir'] for d in subdirs if d['dir'] not in visited]

            if new_subdirs:
                for subdir in new_subdirs:
                    stack.append(subdir)
                # Mark this directory as subdivided
                results["%s/*" % current_dir] = 1
            else:
                # No unvisited subdirs to split into; include directory as-is.
                # This handles the case where p4 dirs returns no children or
                # only paths that cycle back to already-visited directories.
                self.logger.warning("Cannot subdivide %s further; including as-is" % current_dir)
                results["%s/..." % current_dir] = 1

    def run(self):
        """Runs script"""
        self.p4 = P4.P4()
        self.p4.logger = self.logger
        if self.options.port:
            self.p4.port = self.options.port
        if self.options.user:
            self.p4.user = self.options.user
        self.p4.connect()
        results = OrderedDict()
        for path in self.options.path:
            self.get_chunks(results, path, int(self.options.max_size))
        print("\n".join([k for k in results.keys()]))
        return [k for k in results.keys()]

if __name__ == '__main__':
    """ Main Program"""
    obj = depot_verify_chunks(*sys.argv[1:])
    obj.run()