imip-agent

tools/update_quotas.py

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