From 5ea3321a261c06fe9d6640099dbf1b923f64321f Mon Sep 17 00:00:00 2001 From: cmoussa1 Date: Tue, 27 Aug 2024 15:39:24 -0700 Subject: [PATCH] cmd: add export-db as a flux account command Problem: The "export-db" command is located in a separate file in the src/cmd/ directory and thus requires a different method of calling it than the other flux-accounting commands, i.e you need to write "flux account-export-db" instead of "flux account export-db". Move this file to the directory containing the other Python bindings for flux-accounting and add it as a regular command like the other bindings so that you can just call it like "flux account export-db". Remove the methods and command line argument that deal with establishing a connection to the flux-accounting DB since the systemd service will handle establishing that connection. Adjust the calls in the sharness test and top-level README for export-db to account for the change. --- README.md | 2 +- src/Makefile.am | 1 - .../python/fluxacct/accounting/Makefile.am | 1 + .../accounting/db_info_subcommands.py | 48 ++++++ src/cmd/flux-account-export-db.py | 143 ------------------ src/cmd/flux-account-service.py | 20 +++ src/cmd/flux-account.py | 45 ++++++ t/t1016-export-db.t | 4 +- 8 files changed, 117 insertions(+), 147 deletions(-) create mode 100644 src/bindings/python/fluxacct/accounting/db_info_subcommands.py delete mode 100755 src/cmd/flux-account-export-db.py diff --git a/README.md b/README.md index 8b323dd1a..095c1a9ca 100644 --- a/README.md +++ b/README.md @@ -206,7 +206,7 @@ and the `flux account-pop-db` command. Run `flux account-pop-db --help` for `.csv` formatting instructions. User and bank information can also be exported from the database using the -`flux account-export-db` command, which will extract information from both the +`flux account export-db` command, which will extract information from both the user and bank tables and place them into `.csv` files. #### Release diff --git a/src/Makefile.am b/src/Makefile.am index 587277cf7..b969c1e53 100644 --- a/src/Makefile.am +++ b/src/Makefile.am @@ -125,7 +125,6 @@ dist_fluxcmd_SCRIPTS = \ cmd/flux-account-update-fshare \ cmd/flux-account-priority-update.py \ cmd/flux-account-pop-db.py \ - cmd/flux-account-export-db.py \ cmd/flux-account-update-db.py \ cmd/flux-account-service.py \ cmd/flux-account-fetch-job-records.py diff --git a/src/bindings/python/fluxacct/accounting/Makefile.am b/src/bindings/python/fluxacct/accounting/Makefile.am index ac2965949..457bfc09d 100644 --- a/src/bindings/python/fluxacct/accounting/Makefile.am +++ b/src/bindings/python/fluxacct/accounting/Makefile.am @@ -6,6 +6,7 @@ acctpy_PYTHON = \ project_subcommands.py \ job_usage_calculation.py \ jobs_table_subcommands.py \ + db_info_subcommands.py \ create_db.py clean-local: diff --git a/src/bindings/python/fluxacct/accounting/db_info_subcommands.py b/src/bindings/python/fluxacct/accounting/db_info_subcommands.py new file mode 100644 index 000000000..c9882d8b8 --- /dev/null +++ b/src/bindings/python/fluxacct/accounting/db_info_subcommands.py @@ -0,0 +1,48 @@ +#!/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 csv + + +def export_db_info(conn, users=None, banks=None): + try: + cur = conn.cursor() + select_users_stmt = """ + SELECT username, userid, bank, shares, max_running_jobs, max_active_jobs, + max_nodes, queues FROM association_table + """ + cur.execute(select_users_stmt) + table = cur.fetchall() + + # open a .csv file for writing + users_filepath = users if users else "users.csv" + users_file = open(users_filepath, "w") + with users_file: + writer = csv.writer(users_file) + + for row in table: + writer.writerow(row) + + select_banks_stmt = """ + SELECT bank, parent_bank, shares FROM bank_table + """ + cur.execute(select_banks_stmt) + table = cur.fetchall() + + banks_filepath = banks if banks else "banks.csv" + banks_file = open(banks_filepath, "w") + with banks_file: + writer = csv.writer(banks_file) + + for row in table: + writer.writerow(row) + except IOError as err: + print(err) diff --git a/src/cmd/flux-account-export-db.py b/src/cmd/flux-account-export-db.py deleted file mode 100755 index 019f4992b..000000000 --- a/src/cmd/flux-account-export-db.py +++ /dev/null @@ -1,143 +0,0 @@ -#!/usr/bin/env python3 - -############################################################### -# Copyright 2020 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 argparse -import csv -import os -import sqlite3 -import sys - -from argparse import RawDescriptionHelpFormatter - -import fluxacct.accounting - - -def set_db_loc(args): - path = args.path if args.path else fluxacct.accounting.db_path - - return path - - -def est_sqlite_conn(path): - # try to open database file; will exit with -1 if database file not found - if not os.path.isfile(path): - print(f"Database file does not exist: {path}", file=sys.stderr) - sys.exit(1) - - db_uri = "file:" + path + "?mode=rw" - try: - conn = sqlite3.connect(db_uri, uri=True) - # set foreign keys constraint - conn.execute("PRAGMA foreign_keys = 1") - except sqlite3.OperationalError as exc: - print(f"Unable to open database file: {db_uri}", file=sys.stderr) - print(f"Exception: {exc}") - sys.exit(1) - - # check version of database; if not up to date, output message - # and exit - cur = conn.cursor() - cur.execute("PRAGMA user_version") - db_version = cur.fetchone()[0] - if db_version < fluxacct.accounting.db_schema_version: - print( - "flux-accounting database out of date; please update DB with 'flux account-update-db' before running commands" - ) - sys.exit(1) - - return conn - - -def export_db_info(path, users=None, banks=None): - conn = est_sqlite_conn(path) - cur = conn.cursor() - - try: - select_users_stmt = """ - SELECT username, userid, bank, shares, max_running_jobs, max_active_jobs, - max_nodes, queues FROM association_table - """ - cur.execute(select_users_stmt) - table = cur.fetchall() - - # open a .csv file for writing - users_filepath = users if users else "users.csv" - users_file = open(users_filepath, "w") - with users_file: - writer = csv.writer(users_file) - - for row in table: - writer.writerow(row) - - select_banks_stmt = """ - SELECT bank, parent_bank, shares FROM bank_table - """ - cur.execute(select_banks_stmt) - table = cur.fetchall() - - banks_filepath = banks if banks else "banks.csv" - banks_file = open(banks_filepath, "w") - with banks_file: - writer = csv.writer(banks_file) - - for row in table: - writer.writerow(row) - except IOError as err: - print(err) - - -def main(): - parser = argparse.ArgumentParser( - description=""" - Description: Extract flux-accounting database information into two .csv files. - - Order of columns extracted from association_table: - - Username,UserID,Bank,Shares,MaxRunningJobs,MaxActiveJobs,MaxNodes,Queues - - If no custom path is specified, this will create a file in the - current working directory called users.csv. - - ---------------- - - Order of columns extracted from bank_table: - - Bank,ParentBank,Shares - - If no custom path is specified, this will create a file in the - current working directory called banks.csv. - - Use these two files to populate a new flux-accounting DB with: - - flux account-pop-db -p path/to/db -b banks.csv -u users.csv - """, - formatter_class=RawDescriptionHelpFormatter, - ) - - parser.add_argument( - "-p", "--path", dest="path", help="specify location of database file" - ) - parser.add_argument( - "-u", "--users", help="path to a .csv file containing user information" - ) - parser.add_argument( - "-b", "--banks", help="path to a .csv file containing bank information" - ) - - args = parser.parse_args() - - path = set_db_loc(args) - - export_db_info(path, args.users, args.banks) - - -if __name__ == "__main__": - main() diff --git a/src/cmd/flux-account-service.py b/src/cmd/flux-account-service.py index eb18df60b..fef4c5dfd 100755 --- a/src/cmd/flux-account-service.py +++ b/src/cmd/flux-account-service.py @@ -25,6 +25,7 @@ from fluxacct.accounting import queue_subcommands as qu from fluxacct.accounting import project_subcommands as p from fluxacct.accounting import jobs_table_subcommands as j +from fluxacct.accounting import db_info_subcommands as d def establish_sqlite_connection(path): @@ -107,6 +108,7 @@ def __init__(self, flux_handle, conn): "add_project", "delete_project", "scrub_old_jobs", + "export_db", "shutdown_service", ] @@ -523,6 +525,24 @@ def scrub_old_jobs(self, handle, watcher, msg, arg): msg, 0, f"a non-OSError exception was caught: {str(exc)}" ) + def export_db(self, handle, watcher, msg, arg): + try: + val = d.export_db_info( + self.conn, + msg.payload["users"], + msg.payload["banks"], + ) + + payload = {"export_db": val} + + handle.respond(msg, payload) + except KeyError as exc: + handle.respond_error(msg, 0, f"missing key in payload: {exc}") + except Exception as exc: + handle.respond_error( + msg, 0, f"a non-OSError exception was caught: {str(exc)}" + ) + LOGGER = logging.getLogger("flux-uri") diff --git a/src/cmd/flux-account.py b/src/cmd/flux-account.py index 79f9d599e..e2520617f 100755 --- a/src/cmd/flux-account.py +++ b/src/cmd/flux-account.py @@ -532,6 +532,43 @@ def add_scrub_job_records_arg(subparsers): ) +def add_export_db_arg(subparsers): + subparser = subparsers.add_parser( + "export-db", + help=""" + Extract flux-accounting database information into two .csv files. + + Order of columns extracted from association_table: + + Username,UserID,Bank,Shares,MaxRunningJobs,MaxActiveJobs,MaxNodes,Queues + + If no custom path is specified, this will create a file in the + current working directory called users.csv. + + ---------------- + + Order of columns extracted from bank_table: + + Bank,ParentBank,Shares + + If no custom path is specified, this will create a file in the + current working directory called banks.csv. + + Use these two files to populate a new flux-accounting DB with: + + flux account-pop-db -p path/to/db -b banks.csv -u users.csv + """, + formatter_class=flux.util.help_formatter(), + ) + subparser.set_defaults(func="export_db") + subparser.add_argument( + "-u", "--users", help="path to a .csv file containing user information" + ) + subparser.add_argument( + "-b", "--banks", help="path to a .csv file containing bank information" + ) + + def add_arguments_to_parser(parser, subparsers): add_path_arg(parser) add_output_file_arg(parser) @@ -555,6 +592,7 @@ def add_arguments_to_parser(parser, subparsers): add_view_project_arg(subparsers) add_delete_project_arg(subparsers) add_scrub_job_records_arg(subparsers) + add_export_db_arg(subparsers) def set_db_location(args): @@ -734,6 +772,13 @@ def select_accounting_function(args, output_file, parser): "num_weeks": args.num_weeks, } return_val = flux.Flux().rpc("accounting.scrub_old_jobs", data).get() + elif args.func == "export_db": + data = { + "path": args.path, + "users": args.users, + "banks": args.banks, + } + return_val = flux.Flux().rpc("accounting.export_db", data).get() else: print(parser.print_usage()) return diff --git a/t/t1016-export-db.t b/t/t1016-export-db.t index 6f7bc3f98..482f0a7e2 100755 --- a/t/t1016-export-db.t +++ b/t/t1016-export-db.t @@ -39,7 +39,7 @@ test_expect_success 'add some users to the DB' ' ' test_expect_success 'export DB information into .csv files' ' - flux account-export-db -p ${DB_PATHv1} + flux account export-db ' test_expect_success 'compare banks.csv' ' @@ -95,7 +95,7 @@ test_expect_success 'compare DB hierarchies to make sure they are the same' ' ' test_expect_success 'specify a different filename for exported users and banks .csv files' ' - flux account-export-db -p ${DB_PATHv2} --users foo.csv --banks bar.csv && + flux account export-db --users foo.csv --banks bar.csv && test_cmp -b users_expected.csv foo.csv && test_cmp -b banks_expected.csv bar.csv '