Merge pull request #527 from cmoussa1/restructure.association.views

`view-user`: create new `AssociationFormatter` subclass for viewing associations

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

@@ -286,3 +286,23 @@ def with_users(self, bank):
             return info
         except ValueError:
             return info + f"\n\nno users under {bank}"
+
+
+class AssociationFormatter(AccountingFormatter):
+    """
+    Subclass of AccountingFormatter, specific to associations in the flux-accounting
+    database.
+    """
+
+    def __init__(self, cursor, username):
+        """
+        Initialize an AssociationFormatter object with a SQLite cursor.
+
+        Args:
+            cursor: a SQLite Cursor object that has the results of a SQL query.
+            username: the username of the association.
+        """
+        self.username = username
+        super().__init__(
+            cursor, error_msg=f"user {self.username} not found in association_table"
+        )

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

@@ -12,7 +12,10 @@
 import sqlite3
 import time
 import pwd
-import json
+
+import fluxacct.accounting
+from fluxacct.accounting import formatter as fmt
+from fluxacct.accounting import sql_util as sql
 
 ###############################################################
 #                                                             #
@@ -85,86 +88,6 @@ def set_default_project(projects):
     return "*"
 
 
-def create_json_object(conn, user):
-    cur = conn.cursor()
-    main_headers = ["username", "userid", "default_bank"]
-    secondary_headers = [
-        "bank",
-        "active",
-        "shares",
-        "job_usage",
-        "fairshare",
-        "max_running_jobs",
-        "max_active_jobs",
-        "max_nodes",
-        "queues",
-        "projects",
-        "default_project",
-    ]
-
-    cur.execute(
-        """SELECT username, userid, default_bank
-        FROM association_table WHERE username=?""",
-        (user,),
-    )
-    result = cur.fetchall()
-    user_info_dict = dict(zip(main_headers, list(result)[0]))
-
-    cur.execute(
-        """SELECT bank, active, shares, job_usage, fairshare,
-        max_running_jobs, max_active_jobs, max_nodes,
-        queues, projects, default_project FROM association_table
-        WHERE username=?""",
-        (user,),
-    )
-    result = cur.fetchall()
-
-    # store all information pertaining to each bank as a separate
-    # entry in a list
-    user_info_dict["banks"] = []
-    for _ in result:
-        user_info_dict["banks"].append(dict(zip(secondary_headers, list(_))))
-
-    user_info_json = json.dumps(user_info_dict, indent=4)
-    return user_info_json
-
-
-def get_user_rows(conn, user, headers, rows, parsable, json_fmt):
-    if parsable is True:
-        # fetch column names and determine width of each column
-        col_widths = [
-            max(len(str(value)) for value in [col] + [row[i] for row in rows])
-            for i, col in enumerate(headers)
-        ]
-
-        def format_row(row):
-            return " | ".join(
-                [f"{str(value).ljust(col_widths[i])}" for i, value in enumerate(row)]
-            )
-
-        header = format_row(headers)
-        separator = "-+-".join(["-" * width for width in col_widths])
-        data_rows = "\n".join([format_row(row) for row in rows])
-        table = f"{header}\n{separator}\n{data_rows}"
-
-        return table
-
-    user_str = ""
-    if json_fmt is True:
-        user_str += create_json_object(conn, user)
-
-        return user_str
-
-    for row in rows:
-        # iterate through column names of association_table and
-        # print out its associated value
-        for key, value in zip(headers, list(row)):
-            user_str += key + ": " + str(value) + "\n"
-        user_str += "\n"
-
-    return user_str
-
-
 def set_default_bank(cur, username, bank):
     """
     Check the default bank of the user being added; if the user is new, set
@@ -310,25 +233,36 @@ def clear_projects(conn, username, bank=None):
 #                   Subcommand Functions                      #
 #                                                             #
 ###############################################################
-def view_user(conn, user, parsable=False, json_fmt=False):
-    cur = conn.cursor()
+def view_user(conn, user, parsable=False, cols=None):
+    # use all column names if none are passed in
+    cols = cols or fluxacct.accounting.ASSOCIATION_TABLE
+
     try:
-        # get the information pertaining to a user in the DB
-        cur.execute("SELECT * FROM association_table where username=?", (user,))
-        result = cur.fetchall()
-        headers = [description[0] for description in cur.description]  # column names
-        if not result:
-            raise ValueError(f"User {user} not found in association_table")
+        cur = conn.cursor()
+
+        sql.validate_columns(cols, fluxacct.accounting.ASSOCIATION_TABLE)
+        # construct SELECT statement
+        select_stmt = (
+            f"SELECT {', '.join(cols)} FROM association_table WHERE username=?"
+        )
+        cur.execute(select_stmt, (user,))
 
-        user_str = get_user_rows(conn, user, headers, result, parsable, json_fmt)
+        # initialize AssociationFormatter object
+        formatter = fmt.AssociationFormatter(cur, user)
 
-        return user_str
+        if parsable:
+            return formatter.as_table()
+        return formatter.as_json()
     # this kind of exception is raised for errors related to the DB's operation,
     # not necessarily under the control of the programmer, e.g DB path cannot be
     # found or transaction could not be processed
     # (https://docs.python.org/3/library/sqlite3.html#sqlite3.OperationalError)
     except sqlite3.OperationalError as exc:
-        raise sqlite3.OperationalError(f"an sqlite3.OperationalError occurred: {exc}")
+        raise sqlite3.OperationalError(
+            f"view-user: an sqlite3.OperationalError occurred: {exc}"
+        )
+    except ValueError as exc:
+        raise ValueError(f"view-user: {exc}")
 
 
 def add_user(

src/cmd/flux-account-service.py

@@ -149,7 +149,7 @@ def view_user(self, handle, watcher, msg, arg):
                 self.conn,
                 msg.payload["username"],
                 msg.payload["parsable"],
-                msg.payload["json"],
+                msg.payload["fields"].split(",") if msg.payload.get("fields") else None,
             )
 
             payload = {"view_user": val}

src/cmd/flux-account.py

@@ -47,11 +47,15 @@ def add_view_user_arg(subparsers):
         metavar="PARSABLE",
     )
     subparser_view_user.add_argument(
-        "--json",
-        action="store_const",
-        const=True,
-        help="print all information of an association in JSON format",
-        metavar="JSON",
+        "--fields",
+        type=str,
+        help="list of fields to include in JSON output",
+        default=None,
+        metavar=(
+            "CREATION_TIME,MOD_TIME,ACTIVE,USERNAME,USERID,BANK,DEFAULT_BANK,"
+            "SHARES,JOB_USAGE,FAIRSHARE,MAX_RUNNING_JOBS,MAX_ACTIVE_JOBS,MAX_NODES,"
+            "QUEUES,PROJECTS,DEFAULT_PROJECT"
+        ),
     )
 
 

t/Makefile.am

@@ -53,6 +53,7 @@ TESTSCRIPTS = \
 	python/t1006_job_archive.py \
 	python/t1007_formatter.py \
 	python/t1008_banks_output.py
+	python/t1009_users_output.py
 
 dist_check_SCRIPTS = \
 	$(TESTSCRIPTS) \

t/python/t1009_users_output.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+
+###############################################################
+# Copyright 2024 Lawrence Livermore National Security, LLC
+# (c.f. AUTHORS, NOTICE.LLNS, COPYING)
+#
+# This file is part of the Flux resource manager framework.
+# For details, see https://github.com/flux-framework.
+#
+# SPDX-License-Identifier: LGPL-3.0
+###############################################################
+import unittest
+import os
+import sqlite3
+import textwrap
+
+import fluxacct.accounting
+from fluxacct.accounting import create_db as c
+from fluxacct.accounting import bank_subcommands as b
+from fluxacct.accounting import user_subcommands as u
+from fluxacct.accounting import formatter as fmt
+
+
+class TestAccountingCLI(unittest.TestCase):
+    @classmethod
+    def setUpClass(self):
+        # create test accounting database
+        c.create_db("test_view_associations.db")
+        global conn
+        global cur
+
+        conn = sqlite3.connect("test_view_associations.db")
+        cur = conn.cursor()
+
+    # add some associations, initialize formatter
+    def test_AccountingFormatter_with_associations(self):
+        b.add_bank(conn, bank="root", shares=1)
+        b.add_bank(conn, bank="A", shares=1, parent_bank="root")
+        u.add_user(conn, username="user1", bank="A")
+
+        cur.execute("SELECT * FROM association_table")
+        formatter = fmt.AssociationFormatter(cur, "user1")
+
+        self.assertIsInstance(formatter, fmt.AssociationFormatter)
+
+    def test_default_columns_association_table(self):
+        cur.execute("PRAGMA table_info (association_table)")
+        columns = cur.fetchall()
+        association_table = [column[1] for column in columns]
+
+        self.assertEqual(fluxacct.accounting.ASSOCIATION_TABLE, association_table)
+
+    def test_view_association_noexist(self):
+        with self.assertRaises(ValueError):
+            u.view_user(conn, user="foo")
+
+    # test JSON output for viewing an association
+    def test_view_association(self):
+        expected = textwrap.dedent(
+            """\
+        [
+          {
+            "active": 1,
+            "username": "user1",
+            "bank": "A",
+            "shares": 1
+          }
+        ]
+        """
+        )
+        test = u.view_user(
+            conn, user="user1", cols=["active", "username", "bank", "shares"]
+        )
+        self.assertEqual(expected.strip(), test.strip())
+
+    def test_view_association_table(self):
+        expected = textwrap.dedent(
+            """\
+        active | username | bank | shares
+        -------+----------+------+-------
+        1      | user1    | A    | 1  
+        """
+        )
+        test = u.view_user(
+            conn,
+            user="user1",
+            parsable=True,
+            cols=["active", "username", "bank", "shares"],
+        )
+        self.assertEqual(expected.strip(), test.strip())
+
+    # remove database and log file
+    @classmethod
+    def tearDownClass(self):
+        conn.close()
+        os.remove("test_view_associations.db")
+
+
+def suite():
+    suite = unittest.TestSuite()
+
+    return suite
+
+
+if __name__ == "__main__":
+    from pycotap import TAPTestRunner
+
+    unittest.main(testRunner=TAPTestRunner())

t/t1007-flux-account-users.t

@@ -67,19 +67,16 @@ test_expect_success 'trying to add an association that already exists should rai
 
 test_expect_success 'view some user information' '
 	flux account view-user user5011 > user_info.out &&
-	grep -w "username: user5011\|userid: 5011\|bank: A" user_info.out
+	grep "\"username\": \"user5011\"" user_info.out &&
+	grep "\"userid\": 5011" user_info.out &&
+	grep "\"bank\": \"A\"" user_info.out
 '
 
 test_expect_success 'view some user information with --parsable' '
 	flux account view-user --parsable user5011 > user_info_parsable.out &&
 	grep -w "user5011\|5011\|A" user_info_parsable.out
 '
 
-test_expect_success 'view some user information with --json' '
-	flux account view-user --json user5014 > user_info_json.out &&
-	grep -w "\"username\": \"user5014\"\|\"userid\": 5014\|\"bank\": \"C\"" user_info_json.out
-'
-
 test_expect_success 'edit a userid for a user' '
 	flux account edit-user user5011 --userid=12345 &&
 	flux account view-user user5011 > edit_userid.out &&
@@ -95,7 +92,7 @@ test_expect_success 'edit the max_active_jobs of an existing user' '
 
 test_expect_success 'trying to view a user who does not exist in the DB should raise a ValueError' '
 	test_must_fail flux account view-user user9999 > user_nonexistent.out 2>&1 &&
-	grep "User user9999 not found in association_table" user_nonexistent.out
+	grep "view-user: user user9999 not found in association_table" user_nonexistent.out
 '
 
 test_expect_success 'trying to view a user that does exist in the DB should return some information' '

t/t1011-job-archive-interface.t

@@ -60,9 +60,9 @@ test_expect_success 'run scripts to update job usage and fair-share' '
 '
 
 test_expect_success 'check that usage does not get affected by canceled jobs' '
-	flux account view-user --json $username > user.json &&
+	flux account view-user $username > user.json &&
 	test_debug "jq -S . <user.json" &&
-	jq -e ".banks[0].job_usage == 0.0" <user.json
+	jq -e ".[0].job_usage == 0.0" <user.json
 '
 
 test_expect_success 'check that no jobs show up under user' '
@@ -133,9 +133,9 @@ test_expect_success 'run update-usage and update-fshare commands' '
 '
 
 test_expect_success 'check that job usage and fairshare values get updated' '
-	flux account -p ${DB_PATH} view-user $username --json > query1.json &&
+	flux account -p ${DB_PATH} view-user $username > query1.json &&
 	test_debug "jq -S . <query1.json" &&
-	jq -e ".banks[1].job_usage >= 4" <query1.json
+	jq -e ".[1].job_usage >= 4" <query1.json
 '
 
 # if update-usage is called in the same half-life period when no jobs are found
@@ -144,9 +144,9 @@ test_expect_success 'check that job usage and fairshare values get updated' '
 test_expect_success 'call update-usage in the same half-life period where no jobs are run' '
 	flux account -p ${DB_PATH} update-usage &&
 	flux account-update-fshare -p ${DB_PATH} &&
-	flux account -p ${DB_PATH} view-user $username --json > query2.json &&
+	flux account -p ${DB_PATH} view-user $username > query2.json &&
 	test_debug "jq -S . <query2.json" &&
-	jq -e ".banks[1].job_usage >= 4" <query2.json
+	jq -e ".[1].job_usage >= 4" <query2.json
 '
 
 test_expect_success 'remove flux-accounting DB' '

t/t1023-flux-account-banks.t

@@ -115,9 +115,9 @@ test_expect_success 'delete user default bank row' '
 
 test_expect_success 'check that user default bank gets updated to other bank' '
 	flux account view-user user5015 > new_default_bank.out &&
-	grep "username: user5015" new_default_bank.out &&
-	grep "bank: F" new_default_bank.out &&
-	grep "default_bank: F" new_default_bank.out
+	grep "\"username\": \"user5015\"" new_default_bank.out
+	grep "\"bank\": \"F\"" new_default_bank.out &&
+	grep "\"default_bank\": \"F\"" new_default_bank.out
 '
 
 test_expect_success 'trying to add a user to a nonexistent bank should raise a ValueError' '