diff --git a/src/slurmcostmanager.js b/src/slurmcostmanager.js index 90a4760..2fdc8e7 100644 --- a/src/slurmcostmanager.js +++ b/src/slurmcostmanager.js @@ -892,14 +892,10 @@ function Rates({ onRatesUpdated }) { try { let json; if (window.cockpit && window.cockpit.spawn) { - const { start, end } = getBillingPeriod(); const args = [ 'python3', `${PLUGIN_BASE}/slurmdb.py`, - '--start', - start, - '--end', - end, + '--accounts', '--output', '-', ]; diff --git a/src/slurmdb.py b/src/slurmdb.py index 06801b9..b9a3b6a 100644 --- a/src/slurmdb.py +++ b/src/slurmdb.py @@ -452,6 +452,38 @@ def aggregate_usage(self, start_time, end_time): job_entry['core_hours'] += cpus * dur_hours return agg, totals + def fetch_all_accounts(self): + """Return a sorted list of all accounts in the cluster.""" + self.connect() + accounts = set() + with self._conn.cursor() as cur: + try: + cur.execute("SHOW TABLES LIKE 'acct_table'") + if cur.fetchone(): + cur.execute("SELECT name FROM acct_table WHERE deleted = 0") + for row in cur.fetchall(): + name = row.get('name') + if name: + accounts.add(name) + return sorted(accounts) + except pymysql.err.ProgrammingError: + pass + + assoc_table = ( + f"{self.cluster}_assoc_table" if self.cluster else "assoc_table" + ) + try: + cur.execute( + f"SELECT DISTINCT acct FROM {assoc_table} WHERE deleted = 0" + ) + for row in cur.fetchall(): + acct = row.get('acct') + if acct: + accounts.add(acct) + except pymysql.err.ProgrammingError: + pass + return sorted(accounts) + def fetch_invoices(self, start_date=None, end_date=None): """Fetch invoice metadata from the database if present.""" if start_date: @@ -651,6 +683,7 @@ def export_summary(self, start_time, end_time): dest="slurm_conf", help="path to slurm.conf for auto cluster detection", ) + parser.add_argument("--accounts", action="store_true", help="list all accounts and exit") args = parser.parse_args() @@ -660,6 +693,17 @@ def export_summary(self, start_time, end_time): slurm_conf=args.slurm_conf, ) + if args.accounts: + data = {"accounts": db.fetch_all_accounts()} + if args.output in ("-", "/dev/stdout"): + json.dump(data, sys.stdout, indent=2, default=str) + sys.stdout.write("\n") + else: + with open(args.output, "w") as fh: + json.dump(data, fh, indent=2, default=str) + print(f"Wrote {args.output}") + sys.exit(0) + def _export_day(day): day_str = day.isoformat() data = db.export_summary(day_str, day_str) diff --git a/test/check-application b/test/check-application index bb17556..2b64f2b 100755 --- a/test/check-application +++ b/test/check-application @@ -10,3 +10,4 @@ PYTHONPATH=src python test/unit/slurm_schema_dump.test.py PYTHONPATH=src python test/unit/billing_summary.test.py PYTHONPATH=src python test/unit/invoice_retrieval.test.py PYTHONPATH=src python test/unit/auth_boundaries.test.py +PYTHONPATH=src python test/unit/accounts_listing.test.py diff --git a/test/unit/accounts_listing.test.py b/test/unit/accounts_listing.test.py new file mode 100644 index 0000000..22d7dc1 --- /dev/null +++ b/test/unit/accounts_listing.test.py @@ -0,0 +1,81 @@ +import unittest +from slurmdb import SlurmDB + + +class AccountListingTests(unittest.TestCase): + def test_fetch_accounts_from_acct_table(self): + db = SlurmDB() + db.connect = lambda: None + + class FakeCursor: + def __init__(self): + self.last_query = "" + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + pass + + def execute(self, query, params=None): + self.last_query = query + + def fetchone(self): + if "SHOW TABLES LIKE" in self.last_query: + return {"name": "acct_table"} + return None + + def fetchall(self): + if "SELECT name FROM acct_table" in self.last_query: + return [{"name": "acct1"}, {"name": "acct2"}] + return [] + + class FakeConn: + def cursor(self): + return FakeCursor() + + db._conn = FakeConn() + accounts = db.fetch_all_accounts() + self.assertEqual(accounts, ["acct1", "acct2"]) + + def test_fetch_accounts_from_assoc_table_when_acct_table_missing(self): + db = SlurmDB(cluster="localcluster") + db.connect = lambda: None + + class FakeCursor: + def __init__(self): + self.last_query = "" + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + pass + + def execute(self, query, params=None): + self.last_query = query + + def fetchone(self): + # acct_table does not exist + return None + + def fetchall(self): + if "localcluster_assoc_table" in self.last_query: + return [ + {"acct": "acct1"}, + {"acct": "acct2"}, + {"acct": "acct1"}, + ] + return [] + + class FakeConn: + def cursor(self): + return FakeCursor() + + db._conn = FakeConn() + accounts = db.fetch_all_accounts() + self.assertEqual(accounts, ["acct1", "acct2"]) + + +if __name__ == "__main__": + unittest.main()