Changeset 266e866

Timestamp:

Dec 3, 2008 2:24:15 PM (16 years ago)

Author:

Ted Faber <faber@…>

Branches:

axis_example, compt_changes, info-ops, master, version-1.30, version-2.00, version-3.01, version-3.02

Children:

dceeef9

Parents:

b312431b

Message:

Add user databas manipulation commands. Collapse some code

File:

• fedd/fixed_resource.py (modified) (6 diffs)

fedd/fixed_resource.py

@@ -4 +4 @@
 import MySQLdb
 from federation.fixed_resource import read_key_db, write_key_db, \
-        read_project_db, write project_db
+        read_project_db, write_project_db, read_user_db, write_user_db
 from optparse import OptionParser
 
@@ -27 +27 @@
                 "[opts] (--help for details)", version="0.1")
         self.add_option('-t','--type', dest='type', type='choice', 
-                choices=['keys','projects'], help="database type")
+                choices=['keys','projects', 'users'], help="database type")
         self.add_option('-u','--user', dest='user', default=None,
                 action='store', help='user to add/delete')
@@ -41 +41 @@
                 help='database file (synonym for --database)')
 
+def process_db(action, file, key, read_db, write_db, query):
+    """
+    The logic of insertion/deletion/initialization is the same regardless of
+    the DB.  This function encapsulates that control flow.
+    """
+
+    if action != 'init' and action != 'initall':
+        keys = read_db(file)
+    else:
+        keys = set()
+
+    # init action falls through to write an empty DB
+    if action == 'initall':
+        # Add all keys returned by the query from the Emulab DB
+        try:
+            db = MySQLdb.connect(db="tbdb")
+        except:
+            sys.exit("Cannot access the Emulab database")
+
+        c = db.cursor()
+        c.execute(query)
+        for k in c.fetchall():
+            keys.add(k)
+        c.close()
+        db.close()
+    elif action == 'add':
+        keys.add(key)
+    elif action == 'delete': 
+        if key in keys:
+            keys.remove(key)
+        else:
+            print >>sys.stderr, "Cannot delete %s: not in db" % (key,)
+
+    write_db(file, keys)
+
 
 parser = key_opts()
 action = sys.argv[1]
+init_actions = set(['init', 'initall'])
 
 # Check the action
@@ -49 +85 @@
 else: sys.exit("Bad action, must be one of %s" % ", ".join(actions))
 
-# Parse and check the the options for consistency
 (opts, args) = parser.parse_args()
 if not opts.file:
@@ -56 +91 @@
     sys.exit("Must specify database type (--type)")
 elif opts.type == 'keys':
-    if action != "init" and action != "initall" and action != "addall":
+    if action not in init_actions:
         if opts.user and (opts.key or opts.keyfile):
             user = opts.user
@@ -68 +103 @@
         else:
             sys.exit("Must specify user and key")
+    else: 
+        key = None
+        user = None
+    process_db(action, opts.file, (user, key), read_key_db, write_key_db,
+            "SELECT uid, pubkey FROM user_pubkeys")
+
 elif opts.type == 'projects':
-    if action != "init" and action != "initall" and action != "addall":
+    if action not in init_actions:
         if opts.project: project = opts.project
         else: sys.exit("Must specify project")
+    else: project = None
+    process_db(action, opts.file, project, read_project_db, write_project_db, 
+            "SELECT pid FROM projects")
+elif opts.type == 'users':
+    if action not in init_actions:
+        if opts.user: user = opts.user
+        else: sys.exit("Must specify user")
+    else: user = None
+    process_db(action, opts.file, user, read_user_db, write_user_db, 
+            "SELECT uid FROM users")
 else:
     sys.exit("Invalid --type field (how'd you do that?)")
 
-if opts.type == 'keys':
-    if action != 'init' and action != 'initall':
-        keys = read_key_db(opts.file)
-    else:
-        keys = set()
-    
-    if action == 'initall':
-        # Add all users from the Emulab DB
-        try:
-            db = MySQLdb.connect(db="tbdb")
-        except:
-            sys.exit("Cannot access the Emulab database")
-
-        c = db.cursor()
-        c.execute("SELECT uid, pubkey FROM user_pubkeys")
-        for u, k in c.fetchall():
-            keys.add((u, k))
-        c.close()
-        db.close()
-    elif action == 'add':
-        keys.add((user, key))
-    elif action == 'delete': 
-        if (user, key) in keys:
-            keys.remove((user, key))
-        else:
-            print >>sys.stderr, "Cannot delete (%s, %s): not in db" % \
-                    (user, key)
-    # init action falls through to write an empty DB
-    write_key_db(opts.file, keys)
-else:
-    if action != 'init' and action != 'initall':
-        projects = read_project_db(opts.file)
-    else:
-        projects = set()
-    
-    if action == 'initall':
-        # Add all projects from the Emulab DB
-        try:
-            db = MySQLdb.connect(db="tbdb")
-        except:
-            sys.exit("Cannot access the Emulab database")
-
-        c = db.cursor()
-        c.execute("SELECT pid FROM projects")
-        for p in c.fetchall():
-            projects.add(p)
-        c.close()
-        db.close()
-    elif action == 'add':
-        projects.add(project)
-    elif action == 'delete': 
-        if project in projects:
-            projects.remove(project)
-        else:
-            print >>sys.stderr, "Cannot delete %s: not in db" % project
-
-    # init action falls through to write an empty DB
-    write_project_db(opts.file, projects)
-
 sys.exit(0)