imip-agent

Annotated tools/update_quotas.py

1465:0c623c86704f
2020-08-03 Paul Boddie Changed diagram font to sans-serif.
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@1215 6
Copyright (C) 2014, 2015, 2016, 2017 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@1140 22
from os.path import abspath, 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@1140 30
    parent = abspath(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@1215 35
from imiptools.config import settings
paul@1049 36
from imiptools.dates import get_datetime, get_default_timezone, get_time, \
paul@1049 37
                            to_utc_datetime
paul@1088 38
from imiptools.stores import get_journal
paul@1049 39
paul@1049 40
def remove_expired_entries(entries, expiry):
paul@1049 41
paul@1049 42
    "Remove from 'entries' events that end at or before 'expiry'."
paul@1049 43
paul@1049 44
    removed = []
paul@1059 45
    i = 0
paul@1049 46
paul@1049 47
    while i < len(entries):
paul@1059 48
        period = entries[i]
paul@1049 49
paul@1059 50
        if period.get_end_point() <= expiry:
paul@1059 51
            removed.append(period)
paul@1049 52
            del entries[i]
paul@1049 53
        else:
paul@1049 54
            i += 1
paul@1049 55
paul@1049 56
    return removed
paul@1049 57
paul@1049 58
def update_entries(journal, quota, expiry, store, verbose):
paul@1049 59
paul@1049 60
    """
paul@1049 61
    Using the given 'journal' process quota records for the given 'quota', with
paul@1049 62
    the given 'expiry' time used to expire events ending before or at this time,
paul@1049 63
    with None meaning the current time.
paul@1049 64
paul@1049 65
    If 'store' is set, the stored details will be updated; otherwise, the
paul@1049 66
    details will be written to standard output.
paul@1049 67
paul@1049 68
    If 'verbose' is set, messages will be written to standard error.
paul@1049 69
    """
paul@1049 70
paul@1049 71
    if not store:
paul@1049 72
        stdout = getwriter("utf-8")(sys.stdout)
paul@1049 73
    if verbose:
paul@1049 74
        stderr = getwriter("utf-8")(sys.stderr)
paul@1049 75
paul@1049 76
    if not expiry:
paul@1049 77
        expiry = get_time()
paul@1049 78
paul@1049 79
    journal.acquire_lock(quota)
paul@1049 80
paul@1049 81
    try:
paul@1049 82
        for user in journal.get_quota_users(quota):
paul@1049 83
            if verbose:
paul@1049 84
                print >>stderr, user
paul@1049 85
paul@1049 86
            entries = journal.get_entries(quota, user)
paul@1049 87
            removed = remove_expired_entries(entries, expiry)
paul@1049 88
paul@1049 89
            if verbose:
paul@1059 90
                for period in removed:
paul@1059 91
                    print >>stderr, "\t".join(("Removed",) + period.as_tuple(strings_only=True))
paul@1049 92
paul@1049 93
            # Store the processed entries.
paul@1049 94
paul@1049 95
            if store:
paul@1049 96
                journal.set_entries(quota, user, entries)
paul@1049 97
paul@1049 98
            # Alternatively, just write the entries to standard output.
paul@1049 99
paul@1049 100
            else:
paul@1059 101
                for period in entries:
paul@1059 102
                    print >>stdout, "\t".join(period.as_tuple(strings_only=True))
paul@1049 103
    finally:
paul@1049 104
        journal.release_lock(quota)
paul@1049 105
paul@1049 106
# Main program.
paul@1049 107
paul@1049 108
if __name__ == "__main__":
paul@1049 109
paul@1049 110
    # Interpret the command line arguments.
paul@1049 111
paul@1049 112
    quotas = []
paul@1049 113
    args = []
paul@1088 114
    store_type = []
paul@1049 115
    journal_dir = []
paul@1049 116
    expiry = []
paul@1049 117
    ignored = []
paul@1049 118
paul@1049 119
    # Collect quota details first, switching to other arguments when encountering
paul@1049 120
    # switches.
paul@1049 121
paul@1049 122
    l = quotas
paul@1049 123
paul@1049 124
    for arg in sys.argv[1:]:
paul@1049 125
        if arg in ("-s", "-v"):
paul@1049 126
            args.append(arg)
paul@1049 127
            l = ignored
paul@1088 128
        elif arg == "-T":
paul@1088 129
            l = store_type
paul@1049 130
        elif arg == "-j":
paul@1049 131
            l = journal_dir
paul@1049 132
        elif arg == "-e":
paul@1049 133
            l = expiry
paul@1049 134
        else:
paul@1049 135
            l.append(arg)
paul@1049 136
paul@1049 137
    try:
paul@1049 138
        quota = quotas[0]
paul@1049 139
    except IndexError:
paul@1049 140
        print >>sys.stderr, """\
paul@1049 141
Usage: %s <quota> <options>
paul@1049 142
paul@1049 143
Need a quota along with the -s option if updating the journal.
paul@1049 144
Specify -v for additional messages on standard error.
paul@1049 145
paul@1049 146
General options:
paul@1049 147
paul@1088 148
-e  Indicates an expiry time for events (default is now)
paul@1088 149
-j  Indicates the journal directory location
paul@1088 150
-T  Indicates the store type (the configured value if omitted)
paul@1049 151
""" % split(sys.argv[0])[1]
paul@1049 152
        sys.exit(1)
paul@1049 153
paul@1049 154
    # Define any other options.
paul@1049 155
paul@1049 156
    store = "-s" in args
paul@1049 157
    verbose = "-v" in args
paul@1049 158
paul@1049 159
    # Override defaults if indicated.
paul@1049 160
paul@1088 161
    getvalue = lambda value, default=None: value and value[0] or default
paul@1088 162
paul@1215 163
    store_type = getvalue(store_type, settings["STORE_TYPE"])
paul@1088 164
    journal_dir = getvalue(journal_dir)
paul@1088 165
    expiry = getvalue(expiry)
paul@1049 166
paul@1049 167
    if expiry:
paul@1049 168
        expiry = to_utc_datetime(get_datetime(expiry), get_default_timezone())
paul@1049 169
        if not expiry:
paul@1049 170
            print >>sys.stderr, "Expiry time must be a valid datetime."
paul@1049 171
            sys.exit(1)
paul@1049 172
paul@1049 173
    # Obtain store-related objects.
paul@1049 174
paul@1088 175
    journal = get_journal(store_type, journal_dir)
paul@1338 176
    if not journal:
paul@1338 177
        print >>sys.stderr, "Journal not present."
paul@1338 178
        sys.exit(1)
paul@1049 179
paul@1049 180
    # Obtain a list of users for processing.
paul@1049 181
paul@1049 182
    if quota in ("*", "all"):
paul@1049 183
        quotas = journal.get_quotas()
paul@1049 184
paul@1049 185
    # Process the given users.
paul@1049 186
paul@1049 187
    if verbose:
paul@1049 188
        stderr = getwriter("utf-8")(sys.stderr)
paul@1049 189
paul@1049 190
    for quota in quotas:
paul@1049 191
        if verbose:
paul@1049 192
            print >>stderr, quota
paul@1049 193
        update_entries(journal, quota, expiry, store, verbose)
paul@1049 194
paul@1049 195
# vim: tabstop=4 expandtab shiftwidth=4