Merge pull request #486 from cmoussa1/move.export-db

`cmd`: add `export-db` as a flux account command

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

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

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:

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)

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")
 

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

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
 '