imip-agent

Annotated tools/update_quotas.py

1065:2175787d1896
2016-03-04 Paul Boddie Merged changes from the default branch. freebusy-collections
paul@1049 1
#!/usr/bin/env python
paul@1049 2
paul@1049 3
"""
paul@1049 4
Remove expired events from quota journals.
paul@1049 5
paul@1049 6
Copyright (C) 2014, 2015, 2016 Paul Boddie <paul@boddie.org.uk>
paul@1049 7
paul@1049 8
This program is free software; you can redistribute it and/or modify it under
paul@1049 9
the terms of the GNU General Public License as published by the Free Software
paul@1049 10
Foundation; either version 3 of the License, or (at your option) any later
paul@1049 11
version.
paul@1049 12
paul@1049 13
This program is distributed in the hope that it will be useful, but WITHOUT
paul@1049 14
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
paul@1049 15
FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
paul@1049 16
details.
paul@1049 17
paul@1049 18
You should have received a copy of the GNU General Public License along with
paul@1049 19
this program.  If not, see <http://www.gnu.org/licenses/>.
paul@1049 20
"""
paul@1049 21
paul@1049 22
from os.path import split
paul@1049 23
import sys
paul@1049 24
paul@1049 25
# Find the modules.
paul@1049 26
paul@1049 27
try:
paul@1049 28
    import imiptools
paul@1049 29
except ImportError:
paul@1049 30
    parent = split(split(__file__)[0])[0]
paul@1049 31
    if split(parent)[1] == "imip-agent":
paul@1049 32
        sys.path.append(parent)
paul@1049 33
paul@1049 34
from codecs import getwriter
paul@1049 35
from imiptools.dates import get_datetime, get_default_timezone, get_time, \
paul@1049 36
                            to_utc_datetime
paul@1049 37
from imip_store import FileJournal
paul@1049 38
paul@1049 39
def remove_expired_entries(entries, expiry):
paul@1049 40
paul@1049 41
    "Remove from 'entries' events that end at or before 'expiry'."
paul@1049 42
paul@1049 43
    removed = []
paul@1059 44
    i = 0
paul@1049 45
paul@1049 46
    while i < len(entries):
paul@1059 47
        period = entries[i]
paul@1049 48
paul@1059 49
        if period.get_end_point() <= expiry:
paul@1059 50
            removed.append(period)
paul@1049 51
            del entries[i]
paul@1049 52
        else:
paul@1049 53
            i += 1
paul@1049 54
paul@1049 55
    return removed
paul@1049 56
paul@1049 57
def update_entries(journal, quota, expiry, store, verbose):
paul@1049 58
paul@1049 59
    """
paul@1049 60
    Using the given 'journal' process quota records for the given 'quota', with
paul@1049 61
    the given 'expiry' time used to expire events ending before or at this time,
paul@1049 62
    with None meaning the current time.
paul@1049 63
paul@1049 64
    If 'store' is set, the stored details will be updated; otherwise, the
paul@1049 65
    details will be written to standard output.
paul@1049 66
paul@1049 67
    If 'verbose' is set, messages will be written to standard error.
paul@1049 68
    """
paul@1049 69
paul@1049 70
    if not store:
paul@1049 71
        stdout = getwriter("utf-8")(sys.stdout)
paul@1049 72
    if verbose:
paul@1049 73
        stderr = getwriter("utf-8")(sys.stderr)
paul@1049 74
paul@1049 75
    if not expiry:
paul@1049 76
        expiry = get_time()
paul@1049 77
paul@1049 78
    journal.acquire_lock(quota)
paul@1049 79
paul@1049 80
    try:
paul@1049 81
        for user in journal.get_quota_users(quota):
paul@1049 82
            if verbose:
paul@1049 83
                print >>stderr, user
paul@1049 84
paul@1049 85
            entries = journal.get_entries(quota, user)
paul@1049 86
            removed = remove_expired_entries(entries, expiry)
paul@1049 87
paul@1049 88
            if verbose:
paul@1059 89
                for period in removed:
paul@1059 90
                    print >>stderr, "\t".join(("Removed",) + period.as_tuple(strings_only=True))
paul@1049 91
paul@1049 92
            # Store the processed entries.
paul@1049 93
paul@1049 94
            if store:
paul@1049 95
                journal.set_entries(quota, user, entries)
paul@1049 96
paul@1049 97
            # Alternatively, just write the entries to standard output.
paul@1049 98
paul@1049 99
            else:
paul@1059 100
                for period in entries:
paul@1059 101
                    print >>stdout, "\t".join(period.as_tuple(strings_only=True))
paul@1049 102
    finally:
paul@1049 103
        journal.release_lock(quota)
paul@1049 104
paul@1049 105
# Main program.
paul@1049 106
paul@1049 107
if __name__ == "__main__":
paul@1049 108
paul@1049 109
    # Interpret the command line arguments.
paul@1049 110
paul@1049 111
    quotas = []
paul@1049 112
    args = []
paul@1049 113
    journal_dir = []
paul@1049 114
    expiry = []
paul@1049 115
    ignored = []
paul@1049 116
paul@1049 117
    # Collect quota details first, switching to other arguments when encountering
paul@1049 118
    # switches.
paul@1049 119
paul@1049 120
    l = quotas
paul@1049 121
paul@1049 122
    for arg in sys.argv[1:]:
paul@1049 123
        if arg in ("-s", "-v"):
paul@1049 124
            args.append(arg)
paul@1049 125
            l = ignored
paul@1049 126
        elif arg == "-j":
paul@1049 127
            l = journal_dir
paul@1049 128
        elif arg == "-e":
paul@1049 129
            l = expiry
paul@1049 130
        else:
paul@1049 131
            l.append(arg)
paul@1049 132
paul@1049 133
    try:
paul@1049 134
        quota = quotas[0]
paul@1049 135
    except IndexError:
paul@1049 136
        print >>sys.stderr, """\
paul@1049 137
Usage: %s <quota> <options>
paul@1049 138
paul@1049 139
Need a quota along with the -s option if updating the journal.
paul@1049 140
Specify -v for additional messages on standard error.
paul@1049 141
paul@1049 142
General options:
paul@1049 143
paul@1049 144
-e  indicate an expiry time for events (default is now)
paul@1049 145
-j  indicate the journal directory location
paul@1049 146
""" % split(sys.argv[0])[1]
paul@1049 147
        sys.exit(1)
paul@1049 148
paul@1049 149
    # Define any other options.
paul@1049 150
paul@1049 151
    store = "-s" in args
paul@1049 152
    verbose = "-v" in args
paul@1049 153
paul@1049 154
    # Override defaults if indicated.
paul@1049 155
paul@1049 156
    journal_dir = journal_dir and journal_dir[0] or None
paul@1049 157
    expiry = expiry and expiry[0] or None
paul@1049 158
paul@1049 159
    if expiry:
paul@1049 160
        expiry = to_utc_datetime(get_datetime(expiry), get_default_timezone())
paul@1049 161
        if not expiry:
paul@1049 162
            print >>sys.stderr, "Expiry time must be a valid datetime."
paul@1049 163
            sys.exit(1)
paul@1049 164
paul@1049 165
    # Obtain store-related objects.
paul@1049 166
paul@1049 167
    journal = FileJournal(journal_dir)
paul@1049 168
paul@1049 169
    # Obtain a list of users for processing.
paul@1049 170
paul@1049 171
    if quota in ("*", "all"):
paul@1049 172
        quotas = journal.get_quotas()
paul@1049 173
paul@1049 174
    # Process the given users.
paul@1049 175
paul@1049 176
    if verbose:
paul@1049 177
        stderr = getwriter("utf-8")(sys.stderr)
paul@1049 178
paul@1049 179
    for quota in quotas:
paul@1049 180
        if verbose:
paul@1049 181
            print >>stderr, quota
paul@1049 182
        update_entries(journal, quota, expiry, store, verbose)
paul@1049 183
paul@1049 184
# vim: tabstop=4 expandtab shiftwidth=4