python: create new BankFormatter subclass, restructure view-bank to use new class

src/bindings/python/fluxacct/accounting/bank_subcommands.py

@@ -21,63 +21,6 @@
 #                      Helper Functions                       #
 #                                                             #
 ###############################################################
-
-
-def print_user_rows(cur, rows, bank):
-    """Print user information in a table format."""
-    user_str = "\nUsers Under Bank {bank_name}:\n\n".format(bank_name=bank)
-    user_headers = [description[0] for description in cur.description]
-    # print column names of association_table
-    for header in user_headers:
-        user_str += header.ljust(18)
-    user_str += "\n"
-    for row in rows:
-        for col in list(row):
-            user_str += str(col).ljust(18)
-        user_str += "\n"
-
-    return user_str
-
-
-def get_bank_rows(cur, rows, bank):
-    """Print bank information in a table format."""
-    bank_str = ""
-    bank_headers = [description[0] for description in cur.description]
-    # bank has sub banks, so list them
-    for header in bank_headers:
-        bank_str += header.ljust(15)
-    bank_str += "\n"
-    for row in rows:
-        for col in list(row):
-            bank_str += str(col).ljust(15)
-        bank_str += "\n"
-
-    return bank_str
-
-
-def print_sub_banks(conn, bank, bank_str, indent=""):
-    """Traverse the bank table and print all sub banks ansd users."""
-    select_stmt = "SELECT bank FROM bank_table WHERE parent_bank=?"
-    cur = conn.cursor()
-    cur.execute(select_stmt, (bank,))
-    result = cur.fetchall()
-
-    # we've reached a bank with no sub banks
-    if len(result) == 0:
-        cur.execute("SELECT username FROM association_table WHERE bank=?", (bank,))
-        result = cur.fetchall()
-        if result:
-            for row in result:
-                bank_str += indent + " " + row[0] + "\n"
-    # else, delete all of its sub banks and continue traversing
-    else:
-        for row in result:
-            bank_str += indent + " " + row[0] + "\n"
-            bank_str = print_sub_banks(conn, row[0], bank_str, indent + " ")
-
-    return bank_str
-
-
 def validate_parent_bank(cur, parent_bank):
     try:
         cur.execute("SELECT shares FROM bank_table WHERE bank=?", (parent_bank,))
@@ -134,83 +77,6 @@ def reactivate_bank(conn, cur, bank, parent_bank):
     conn.commit()
 
 
-def print_hierarchy(cur, bank, hierarchy_str, indent=""):
-    # look for all sub banks under this parent bank
-    select_stmt = "SELECT bank,shares,job_usage FROM bank_table WHERE parent_bank=?"
-    cur.execute(select_stmt, (bank,))
-    sub_banks = cur.fetchall()
-
-    if len(sub_banks) == 0:
-        # we've reached a bank with no sub banks, so print out every user
-        # under this bank
-        cur.execute(
-            "SELECT username,shares,job_usage,fairshare FROM association_table WHERE bank=?",
-            (bank,),
-        )
-        users = cur.fetchall()
-        if users:
-            for user in users:
-                hierarchy_str += (
-                    indent
-                    + " "
-                    + bank.ljust(20)
-                    + str(user[0]).rjust(20 - (len(indent) + 1))
-                    + str(user[1]).rjust(20)
-                    + str(user[2]).rjust(20)
-                    + str(user[3]).rjust(20)
-                    + "\n"
-                )
-    else:
-        # continue traversing the hierarchy
-        for sub_bank in sub_banks:
-            hierarchy_str += (
-                indent
-                + " "
-                + str(sub_bank[0]).ljust(20)
-                + "".rjust(20 - (len(indent) + 1))  # this skips the "Username" column
-                + str(sub_bank[1]).rjust(20)
-                + str(sub_bank[2]).rjust(20)
-                + "\n"
-            )
-            hierarchy_str = print_hierarchy(
-                cur, sub_bank[0], hierarchy_str, indent + " "
-            )
-
-    return hierarchy_str
-
-
-def print_parsable_hierarchy(cur, bank, hierarchy_str, indent=""):
-    # look for all sub banks under this parent bank
-    select_stmt = "SELECT bank,shares,job_usage FROM bank_table WHERE parent_bank=?"
-    cur.execute(select_stmt, (bank,))
-    sub_banks = cur.fetchall()
-
-    if len(sub_banks) == 0:
-        # we've reached a bank with no sub banks, so print out every user
-        # under this bank
-        cur.execute(
-            "SELECT username,shares,job_usage,fairshare FROM association_table WHERE bank=?",
-            (bank,),
-        )
-        users = cur.fetchall()
-        if users:
-            for user in users:
-                hierarchy_str += (
-                    f"{indent} {bank}|{user[0]}|{user[1]}|{user[2]}|{user[3]}\n"
-                )
-    else:
-        # continue traversing the hierarchy
-        for sub_bank in sub_banks:
-            hierarchy_str += (
-                f"{indent} {str(sub_bank[0])}||{str(sub_bank[1])}|{str(sub_bank[2])}\n"
-            )
-            hierarchy_str = print_parsable_hierarchy(
-                cur, sub_bank[0], hierarchy_str, indent + " "
-            )
-
-    return hierarchy_str
-
-
 ###############################################################
 #                                                             #
 #                   Subcommand Functions                      #
@@ -270,72 +136,39 @@ def add_bank(conn, bank, shares, parent_bank=""):
         raise sqlite3.IntegrityError(f"bank {bank} already exists in bank_table")
 
 
-def view_bank(conn, bank, tree=False, users=False, parsable=False):
-    cur = conn.cursor()
-    bank_str = ""
+def view_bank(conn, bank, tree=False, users=False, parsable=False, cols=None):
+    if tree and cols is not None:
+        # tree format cannot be combined with custom formatting, so raise an Exception
+        raise ValueError(f"--tree option does not support custom formatting")
+    if parsable and not tree:
+        # --parsable can only be called with --tree, so raise an Exception
+        raise ValueError(f"-P/--parsable can only be passed with -t/--tree")
+
+    # use all column names if none are passed in
+    cols = cols or fluxacct.accounting.BANK_TABLE
+
     try:
-        cur.execute("SELECT * FROM bank_table WHERE bank=?", (bank,))
-        result = cur.fetchall()
-
-        if result:
-            bank_str = get_bank_rows(cur, result, bank)
-        else:
-            raise ValueError(f"bank {bank} not found in bank_table")
-
-        name = result[0][1]
-        shares = result[0][4]
-        usage = result[0][5]
-
-        if parsable is True:
-            # print out the database hierarchy starting with the bank passed in
-            hierarchy_str = "Bank|Username|RawShares|RawUsage|Fairshare\n"
-            hierarchy_str += f"{name}||{str(shares)}|{str(round(usage, 2))}\n"
-            hierarchy_str = print_parsable_hierarchy(cur, bank, hierarchy_str, "")
-            return hierarchy_str
-        if tree is True:
-            # print out the hierarchy view with the specified bank as the root of the tree
-            hierarchy_str = (
-                "Bank".ljust(20)
-                + "Username".rjust(20)
-                + "RawShares".rjust(20)
-                + "RawUsage".rjust(20)
-                + "Fairshare".rjust(20)
-                + "\n"
-            )
-            # add the bank passed in to the hierarchy string
-            hierarchy_str += (
-                name.ljust(20)
-                + "".rjust(20)
-                + str(shares).rjust(20)
-                + str(round(usage, 2)).rjust(20)
-                + "\n"
-            )
+        cur = conn.cursor()
 
-            hierarchy_str = print_hierarchy(cur, name, hierarchy_str, "")
-            bank_str += "\n" + hierarchy_str
-        # if users is passed in, print out all potential users under
-        # the passed in bank
-        if users is True:
-            select_stmt = """
-                        SELECT username,userid,default_bank,shares,job_usage,
-                        fairshare,max_running_jobs,queues FROM association_table
-                        WHERE bank=?
-                        """
-            cur.execute(
-                select_stmt,
-                (bank,),
-            )
-            result = cur.fetchall()
+        sql.validate_columns(cols, fluxacct.accounting.BANK_TABLE)
+        # construct SELECT statement
+        select_stmt = f"SELECT {', '.join(cols)} FROM bank_table WHERE bank=?"
+        cur.execute(select_stmt, (bank,))
 
-            if result:
-                user_str = print_user_rows(cur, result, bank)
-                bank_str += user_str
-            else:
-                bank_str += "\nno users under {bank_name}".format(bank_name=bank)
+        # initialize BankFormatter object
+        formatter = fmt.BankFormatter(cur, bank)
 
-        return bank_str
-    except sqlite3.OperationalError as exc:
-        raise sqlite3.OperationalError(f"an sqlite3.OperationalError occurred: {exc}")
+        if tree:
+            if parsable:
+                return formatter.as_parsable_tree(bank)
+            return formatter.as_tree()
+        if users:
+            return formatter.with_users(bank)
+        return formatter.as_json()
+    except sqlite3.Error as err:
+        raise sqlite3.Error(f"view-bank: an sqlite3.Error occurred: {err}")
+    except ValueError as exc:
+        raise ValueError(f"view-bank: {exc}")
 
 
 def delete_bank(conn, bank):