Switch from optparse to argparse

[invirt/packages/invirt-images.git] / invirt-images

diff --git a/invirt-images b/invirt-images
index 1566cc6..1e77767 100755 (executable)
--- a/invirt-images
+++ b/invirt-images
@@ -1,135 +1,90 @@
-#!/usr/bin/python
+#!/usr/bin/env python3

 
 

-from invirt import database

 import os
 import subprocess
 import random
 import string
 import tempfile
 import os
 import subprocess
 import random
 import string
 import tempfile

-import urllib
+import urllib.request

 import math
 import math

-import optparse as op
-
-class InvirtImageException(Exception):
-    pass
+import argparse

 
 

-# verbosity = 0 means no output from the actual commands
-# verbosity = 1 means only errors from the actual commands
-# verbosity = 2 means all output from the actual commands
-verbosity = 0
+from invirt import database

 
 

-def getOutput():
-    global verbosity
-    return {
-        'stdout': subprocess.PIPE if verbosity < 2 else None,
-        'stderr': subprocess.PIPE if verbosity < 1 else None
-        }

 
 def lvcreate(name, size):
 
 def lvcreate(name, size):

-    lvc = subprocess.Popen(['lvcreate', '-L', size, '-n', name, 'xenvg'],
-                           stderr=subprocess.PIPE,
-                           stdout=getOutput()['stdout'])
-    if not lvc.wait():
-        return 0
-    stderr = lvc.stderr.read()
-    if 'already exists in volume group' in stderr:
-        return 5
-    else:
-        if verbosity > 0:
-            print stderr
-        return 6
+    subprocess.run(['lvcreate', '-L', size, '-n', name, 'xenvg'],
+                   check_output=True, encoding='utf-8', check=True)

 
 def lvrename(dest, src):
 
 def lvrename(dest, src):

-    lvr = subprocess.Popen(['lvrename', 'xenvg', src, dest],
-                           stderr=subprocess.PIPE,
-                           stdout=getOutput()['stdout'])
-    ret = lvr.wait()
-    if not ret:
-        return 0
-    stderr = lvr.stderr.read()
-    if 'not found in volume group' in stderr:
-        return 0
-    else:
-        if verbosity > 0:
-            print stderr
-        return ret
+    subprocess.run(['lvchange', '-an', f'xenvg/{src}'],
+                   check_output=True, encoding='utf-8', check=True)
+
+    subprocess.run(['lvrename', f'xenvg/{src}', f'xenvg/{dest}'],
+                   check_output=True, encoding='utf-8', check=True)
+
+    subprocess.run(['lvchange', '-ay', f'xenvg/{dest}'],
+                   check_output=True, encoding='utf-8', check=True)

 
 def lv_random(func, pattern, *args):
     """
 
 def lv_random(func, pattern, *args):
     """

-    Run some LVM-related command, optionally with a random string in
-    the LV name.
-    
+    Run some LVM-related command with a random string in the LV name.
+

     func takes an LV name plus whatever's in *args and returns the
     return code of some LVM command, such as lvcreate or lvrename
     func takes an LV name plus whatever's in *args and returns the
     return code of some LVM command, such as lvcreate or lvrename

-    
-    pattern can contain at most one '%s' pattern, which will be
-    replaced by a 6-character random string.
-    
-    If pattern contains a '%s', the script will attempt to re-run
-    itself if the error code indicates that the destination already
-    exists
+
+    pattern must contain one '{}' pattern, which will be replaced
+    by a 6-character random string.
+
+    the script will attempt to re-run itself if the error code
+    indicates that the destination already exists

     """
     """

+

     # Keep trying until it works
     while True:
     # Keep trying until it works
     while True:

-        rand_string = ''.join(random.choice(string.ascii_letters) \
-                                  for i in xrange(6))
-        if '%s' in pattern:
-            name = pattern % rand_string
-        else:
-            name = pattern
-        ret = func(name, *args)
-        if ret == 0:
-            return name
-        # 5 is the return code if the destination already exists
-        elif '%s' not in pattern or ret != 5:
-            raise InvirtImageException, 'E: Error running %s with args %s' % (func.__name__, args)
+        letters = (random.choice(string.ascii_letters) for _ in range(6))
+        rand_string = ''.join(letters)

 
 

-def lvcreate_random(pattern, size):
-    """
-    Creates an LV, optionally with a random string in the name.
-    
-    Call with a string formatting pattern with a single '%s' to use as
-    a pattern for the name of the new LV.
-    """
-    return lv_random(lvcreate, pattern, size)
+        name = pattern.format(rand_string)

 
 

-def lvrename_random(src, pattern):
-    """
-    Rename an LV to a new name with a random string incorporated.
-    
-    Call with a string formatting pattern with a single '%s' to use as
-    a pattern for the name of the new LV
-    """
-    return lv_random(lvrename, pattern, src)
+        try:
+            func(name, *args)
+        except subprocess.CalledProcessError as e:
+            # 5 is the return code if the destination already exists
+            if e.returncode != 5:
+                raise
+        else:
+            return name

 
 def fetch_image(cdrom):
     """
     Download a cdrom from a URI, shelling out to rsync if appropriate
     and otherwise trying to use urllib
     """
 
 def fetch_image(cdrom):
     """
     Download a cdrom from a URI, shelling out to rsync if appropriate
     and otherwise trying to use urllib
     """

+

     full_uri = os.path.join(cdrom.mirror.uri_prefix, cdrom.uri_suffix)
     temp_file = tempfile.mkstemp()[1]
     full_uri = os.path.join(cdrom.mirror.uri_prefix, cdrom.uri_suffix)
     temp_file = tempfile.mkstemp()[1]

+    print(f'Fetching image {cdrom.cdrom_id} from {full_uri} to {temp_file}')

     try:
         if full_uri.startswith('rsync://'):
     try:
         if full_uri.startswith('rsync://'):

-            if subprocess.call(['rsync', '--no-motd', '-tLP', full_uri, temp_file],
-                               **getOutput()):
-                raise InvirtImageException, "E: Unable to download '%s'" % full_uri
+            subprocess.run(['rsync', '--no-motd', '-tLP', full_uri, temp_file],
+                           check_output=True, encoding='utf-8', check=True)

         else:
             # I'm not going to look for errors here, because I bet it'll
             # throw its own exceptions
         else:
             # I'm not going to look for errors here, because I bet it'll
             # throw its own exceptions

-            urllib.urlretrieve(full_uri, temp_file)
+            urllib.request.urlretrieve(full_uri, temp_file)

         return temp_file
     except:
         return temp_file
     except:

-        os.unlink(temp_file)
+        os.remove(temp_file)

         raise
 
 def copy_file(src, dest):
     """
     Copy a file from one location to another using dd
     """
         raise
 
 def copy_file(src, dest):
     """
     Copy a file from one location to another using dd
     """

-    if subprocess.call(['dd', 'if=%s' % src, 'of=%s' % dest, 'bs=1M'],
-                       **getOutput()):
-        raise InvirtImageException, 'E: Unable to transfer %s into %s' % (src, dest)
+
+    subprocess.run(['dd', f'if={src}', f'of={dest}', 'bs=1M'],
+                   check_output=True, encoding='utf-8', check=True)

 
 def load_image(cdrom):
     """
 
 def load_image(cdrom):
     """

@@ -137,121 +92,113 @@ def load_image(cdrom):
     transferring it into an LV, moving the old LV out of the way and
     the new LV into place
     """
     transferring it into an LV, moving the old LV out of the way and
     the new LV into place
     """

+

     if cdrom.mirror_id is None:
         return
     if cdrom.mirror_id is None:
         return

+

     temp_file = fetch_image(cdrom)
     temp_file = fetch_image(cdrom)

+
+    st_size = os.stat(temp_file).st_size
+
+    assert st_size > 0, 'CD-ROM image size is 0'
+
+    megabytes = math.ceil((float(st_size) / (1024 * 1024)))
+    cdrom_size = f'{megabytes}M'
+

     try:
     try:

-        cdrom_size = '%sM' % math.ceil((float(os.stat(temp_file).st_size) / (1024 * 1024)))
-        new_lv = lvcreate_random('image-new_%s_%%s' % cdrom.cdrom_id, cdrom_size)
-        copy_file(temp_file, '/dev/xenvg/%s' % new_lv)
-        lvrename_random('image_%s' % cdrom.cdrom_id, 'image-old_%s_%%s' % cdrom.cdrom_id)
-        lvrename_random(new_lv, 'image_%s' % cdrom.cdrom_id)
+        new_lv = lv_random(lvcreate, f'image-new_{cdrom.cdrom_id}' + '_{}', cdrom_size)
+        copy_file(temp_file, f'/dev/xenvg/{new_lv}')
+        lv_random(lvrename, f'image-old_{cdrom.cdrom_id}' + '_{}', 'image_{cdrom.cdrom_id}')
+        lv_random(lvrename, f'image_{cdrom.cdrom_id}', new_lv)

         reap_images()
     finally:
         reap_images()
     finally:

-        os.unlink(temp_file)
+        os.remove(temp_file)

 
 def reap_images():
     """
     Remove stale cdrom images that are no longer in use
 
 def reap_images():
     """
     Remove stale cdrom images that are no longer in use

-    
+

     load_image doesn't attempt to remove the old image because it
     might still be in use. reap_images attempts to delete any LVs
     starting with 'image-old_', but ignores errors, in case they're
     still being used.
     """
     load_image doesn't attempt to remove the old image because it
     might still be in use. reap_images attempts to delete any LVs
     starting with 'image-old_', but ignores errors, in case they're
     still being used.
     """

-    lvm_list = subprocess.Popen(['lvs', '-o', 'lv_name', '--noheadings'],
-                               stdout=subprocess.PIPE,
-                               stdin=subprocess.PIPE)
-    lvm_list.wait()
-    
-    for lv in map(str.strip, lvm_list.stdout.read().splitlines()):
+
+    lvm_list = subprocess.run(['lvs', '-o', 'lv_name', '--noheadings'],
+                              check_output=True, encoding='utf-8', check=True)
+
+    for lv in (s.strip() for s in lvm_list.stdout.readlines()):

         if lv.startswith('image-old_'):
         if lv.startswith('image-old_'):

-            subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
-                            **getOutput())
-            subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
-                            **getOutput())
-            subprocess.call(['lvchange', '-a', 'ey', '/dev/xenvg/%s' % lv],
-                            **getOutput())
-            subprocess.call(['lvremove', '--force', '/dev/xenvg/%s' % lv],
-                            **getOutput())
+            subprocess.run(['lvchange', '-a', 'n', f'/dev/xenvg/{lv}'])
+            subprocess.run(['lvchange', '-a', 'n', f'/dev/xenvg/{lv}'])
+            subprocess.run(['lvchange', '-a', 'ey', f'/dev/xenvg/{lv}'])
+            subprocess.run(['lvremove', '--force', f'/dev/xenvg/{lv}'])
+
+def action_add_cdrom(args):
+    attrs = {key: vars(args)[key] for key in ('cdrom_id', 'description', 'mirror_id', 'uri_suffix')}
+
+    cdrom = database.CDROM(**attrs)
+    database.session.add(cdrom)
+    database.session.commit()
+
+    load_image(cdrom)
+
+def action_add_mirror(args):
+    attrs = {key: vars(args)[key] for key in ('mirror_id', 'uri_prefix')}
+
+    mirror = database.Mirror(**attrs)
+    database.session.add(mirror)
+    database.session.commit()
+
+def action_update(args):
+    if not args.names:
+        images = database.CDROM.query().all()
+    else:
+        images = [database.CDROM.query().get(arg) for arg in args.names]
+
+    for cdrom in images:
+        if cdrom is not None:
+            load_image(cdrom)
+
+def action_reap(args):
+    reap_images()

 
 def main():
 
 def main():

-    global verbosity
-    

     database.connect()
     database.connect()

-    
-    usage = """%prog [options] --add [--cdrom] cdrom_id description mirror_id uri_suffix
-       %prog [options] --add --mirror mirror_id uri_prefix
-
-       %prog [options] --update [short_name1 [short_name2 ...]]
-       %prog [options] --reap"""
-    
-    parser = op.OptionParser(usage=usage)
-    parser.set_defaults(verbosity=0,
-                        item='cdrom')
-    
-    parser.add_option('-a', '--add', action='store_const',
-                      dest='action', const='add',
-                      help='Add a new item to the database')
-    
-    parser.add_option('-u', '--update', action='store_const',
-                      dest='action', const='update',
-                      help='Update all cdrom images in the database with the latest version')
-    parser.add_option('-r', '--reap', action='store_const',
-                      dest='action', const='reap',
-                      help='Reap stale cdrom images that are no longer in use')
-    
-    a_group = op.OptionGroup(parser, 'Adding new items')
-    a_group.add_option('-c', '--cdrom', action='store_const',
-                       dest='item', const='cdrom',
-                       help='Add a new cdrom to the database')
-    a_group.add_option('-m', '--mirror', action='store_const',
-                       dest='item', const='mirror',
-                       help='Add a new mirror to the database')
-    parser.add_option_group(a_group)
-    
-    v_group = op.OptionGroup(parser, "Verbosity levels")
-    v_group.add_option("-q", "--quiet", action='store_const',
-                       dest='verbosity', const=0,
-                       help='Show no output from commands this script runs (default)')
-    v_group.add_option("-v", "--verbose", action='store_const',
-                       dest='verbosity', const=1,
-                       help='Show only errors from commands this script runs')
-    v_group.add_option("--noisy", action='store_const',
-                       dest='verbosity', const=2,
-                       help='Show all output from commands this script runs')
-    parser.add_option_group(v_group)
-    
-    (options, args) = parser.parse_args()
-    verbosity = options.verbosity
-    if options.action is None:
-        print parser.format_help()
-    elif options.action == 'add':
-        if options.item == 'cdrom':
-            attrs = dict(zip(('cdrom_id', 'description', 'mirror_id', 'uri_suffix'),
-                             args))
-            cdrom = database.CDROM(**attrs)
-            database.session.save(cdrom)
-            database.session.flush()
-            
-            load_image(cdrom)
-        
-        elif options.item == 'mirror':
-            attrs = dict(zip(('mirror_id', 'uri_prefix'),
-                             args))
-            mirror = database.Mirror(**attrs)
-            database.session.save(mirror)
-            database.session.flush()
-    elif options.action == 'update':
-        if len(args) > 0:
-            images = [database.CDROM.query().get(arg) for arg in args]
-        else:
-            images = database.CDROM.query().all()
-        for cdrom in images:
-            if cdrom is not None:
-                load_image(cdrom)
-    elif options.action == 'reap':
-        reap_images()
+    database.session.begin()
+
+    parser = argparse.ArgumentParser(description='Perform actions on the CD-ROM images in the database')
+    subparsers = parser.add_subparsers(help='Action to perform')
+
+    add_parser = subparsers.add_parser('add', help='Add new image to database')
+    add_parser.set_defaults(func=lambda args: add_parser.print_help())
+    add_subparsers = add_parser.add_subparsers()
+
+    add_cdrom_parser = add_subparsers.add_parser('cdrom')
+    add_cdrom_parser.add_argument('cdrom_id')
+    add_cdrom_parser.add_argument('description')
+    add_cdrom_parser.add_argument('mirror_id')
+    add_cdrom_parser.add_argument('uri_suffix')
+    add_cdrom_parser.set_defaults(func=action_add_cdrom)
+
+    add_mirror_parser = add_subparsers.add_parser('mirror')
+    add_mirror_parser.add_argument('mirror_id')
+    add_mirror_parser.add_argument('uri_prefix')
+    add_mirror_parser.set_defaults(func=action_add_mirror)
+
+    update_parser = subparsers.add_parser('update', help='Update images in database')
+    update_parser.add_argument('names', nargs='*', metavar='name', help='Shortnames of images to update')
+    update_parser.set_defaults(func=action_update)
+
+    reap_parser = subparsers.add_parser('reap', help='Reap old images in database')
+    reap_parser.set_defaults(func=action_reap)
+
+    args = parser.parse_args()
+    if 'func' in args:
+        args.func(args)
+    else:
+        parser.print_help()
+

 
 if __name__ == '__main__':
     main()
 
 if __name__ == '__main__':
     main()