From a39f2c5a48ee2302342f2a90eb55408caa120951 Mon Sep 17 00:00:00 2001 From: cmoussa1 Date: Thu, 1 Aug 2024 11:25:28 -0700 Subject: [PATCH] python: add list-banks command Problem: There is currently no way to list the current banks defined in the bank_table in the flux-accounting database, especially in a convenient format like JSON. Add a new command to the flux-accounting command suite called "list-banks". When called, this command will create a JSON object containing all of the active banks in the bank_table. Passing --inactive will include *all* banks, even those which have been disabled. The fields output in the JSON object can be customized with a list of fields to be included with --fields. --- .../fluxacct/accounting/bank_subcommands.py | 48 +++++++++++++++++++ src/cmd/flux-account-service.py | 25 +++++++++- src/cmd/flux-account.py | 30 ++++++++++++ 3 files changed, 102 insertions(+), 1 deletion(-) diff --git a/src/bindings/python/fluxacct/accounting/bank_subcommands.py b/src/bindings/python/fluxacct/accounting/bank_subcommands.py index 9c18e307..b4e7c8be 100644 --- a/src/bindings/python/fluxacct/accounting/bank_subcommands.py +++ b/src/bindings/python/fluxacct/accounting/bank_subcommands.py @@ -10,6 +10,7 @@ # SPDX-License-Identifier: LGPL-3.0 ############################################################### import sqlite3 +import json from fluxacct.accounting import user_subcommands as u @@ -410,3 +411,50 @@ def edit_bank( conn.commit() return 0 + + +def list_banks( + conn, + inactive=False, + fields=None, +): + """ + List all banks in the bank_table in JSON format. + + Args: + inactive: whether to include inactive banks. By default, only banks that are + active will be included in the output. + + fields: a list of fields to include in the output. By default, all fields are + included. + """ + default_fields = {"bank_id", "bank", "active", "parent_bank", "shares", "job_usage"} + # if fields is None, just use the default fields + fields = fields or default_fields + + try: + cur = conn.cursor() + + # validate the fields passed in + invalid_fields = [field for field in fields if field not in default_fields] + if invalid_fields: + raise ValueError(f"invalid fields: {', '.join(invalid_fields)}") + + # construct SELECT statement + select_fields = ", ".join(fields) + select_stmt = f"SELECT {select_fields} FROM bank_table" + if not inactive: + select_stmt += " WHERE active=1" + + cur.execute(select_stmt) + result = cur.fetchall() + + # create individual object for each row in the query result + banks = [ + {field: row[idx] for idx, field in enumerate(fields)} for row in result + ] + + json_string = json.dumps(banks, indent=2) + return json_string + except sqlite3.Error as err: + raise sqlite3.Error(f"an sqlite3.Error occurred: {err}") diff --git a/src/cmd/flux-account-service.py b/src/cmd/flux-account-service.py index 2d3f3ee8..eb18df60 100755 --- a/src/cmd/flux-account-service.py +++ b/src/cmd/flux-account-service.py @@ -66,7 +66,7 @@ def check_db_version(conn): sys.exit(1) -# pylint: disable=broad-except +# pylint: disable=broad-except, too-many-public-methods class AccountingService: def __init__(self, flux_handle, conn): @@ -87,6 +87,7 @@ def __init__(self, flux_handle, conn): general_endpoints = [ "view_user", "view_bank", + "list_banks", "view_job_records", "view_queue", "view_project", @@ -321,6 +322,28 @@ def edit_bank(self, handle, watcher, msg, arg): msg, 0, f"a non-OSError exception was caught: {str(exc)}" ) + def list_banks(self, handle, watcher, msg, arg): + try: + val = b.list_banks( + self.conn, + msg.payload["inactive"], + msg.payload["fields"], + ) + + payload = {"list_banks": val} + + handle.respond(msg, payload) + except KeyError as exc: + handle.respond_error(msg, 0, f"missing key in payload: {exc}") + except ValueError as exc: + handle.respond_error(msg, 0, f"{exc}") + except sqlite3.Error as exc: + handle.respond_error(msg, 0, f"a SQLite error occured: {exc}") + except Exception as exc: + handle.respond_error( + msg, 0, f"a non-OSError exception was caught: {str(exc)}" + ) + # pylint: disable=no-self-use def view_job_records(self, handle, watcher, msg, arg): try: diff --git a/src/cmd/flux-account.py b/src/cmd/flux-account.py index c68a8f36..79f9d599 100755 --- a/src/cmd/flux-account.py +++ b/src/cmd/flux-account.py @@ -343,6 +343,28 @@ def add_edit_bank_arg(subparsers): ) +def add_list_banks_arg(subparsers): + subparser_list_banks = subparsers.add_parser( + "list-banks", + help="list all banks in the flux-accounting DB", + formatter_class=flux.util.help_formatter(), + ) + subparser_list_banks.set_defaults(func="list_banks") + subparser_list_banks.add_argument( + "--inactive", + action="store_const", + const=True, + help="include inactive banks in output", + ) + subparser_list_banks.add_argument( + "--fields", + type=str, + help="list of fields to include in JSON output", + default="bank_id,bank,parent_bank,shares,job_usage", + metavar="BANK_ID,BANK,ACTIVE,PARENT_BANK,SHARES,JOB_USAGE", + ) + + def add_update_usage_arg(subparsers): subparser_update_usage = subparsers.add_parser( "update-usage", @@ -523,6 +545,7 @@ def add_arguments_to_parser(parser, subparsers): add_view_bank_arg(subparsers) add_delete_bank_arg(subparsers) add_edit_bank_arg(subparsers) + add_list_banks_arg(subparsers) add_update_usage_arg(subparsers) add_add_queue_arg(subparsers) add_view_queue_arg(subparsers) @@ -642,6 +665,13 @@ def select_accounting_function(args, output_file, parser): "parent_bank": args.parent_bank, } return_val = flux.Flux().rpc("accounting.edit_bank", data).get() + elif args.func == "list_banks": + data = { + "path": args.path, + "inactive": args.inactive, + "fields": args.fields.split(","), + } + return_val = flux.Flux().rpc("accounting.list_banks", data).get() elif args.func == "update_usage": data = { "path": args.path,