Skip to content

Commit

Permalink
cmd: add export-db as a flux account command
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
cmoussa1 committed Aug 27, 2024
1 parent eea54d9 commit 5ea3321
Show file tree
Hide file tree
Showing 8 changed files with 117 additions and 147 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 0 additions & 1 deletion src/Makefile.am
Original file line number Diff line number Diff line change
Expand Up @@ -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
1 change: 1 addition & 0 deletions src/bindings/python/fluxacct/accounting/Makefile.am
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
48 changes: 48 additions & 0 deletions src/bindings/python/fluxacct/accounting/db_info_subcommands.py
Original file line number Diff line number Diff line change
@@ -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)
143 changes: 0 additions & 143 deletions src/cmd/flux-account-export-db.py

This file was deleted.

20 changes: 20 additions & 0 deletions src/cmd/flux-account-service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -107,6 +108,7 @@ def __init__(self, flux_handle, conn):
"add_project",
"delete_project",
"scrub_old_jobs",
"export_db",
"shutdown_service",
]

Expand Down Expand Up @@ -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")

Expand Down
45 changes: 45 additions & 0 deletions src/cmd/flux-account.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions t/t1016-export-db.t
Original file line number Diff line number Diff line change
Expand Up @@ -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' '
Expand Down Expand Up @@ -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
'
Expand Down

0 comments on commit 5ea3321

Please sign in to comment.