Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 1 addition & 5 deletions src/slurmcostmanager.js
Original file line number Diff line number Diff line change
Expand Up @@ -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',
'-',
];
Expand Down
44 changes: 44 additions & 0 deletions src/slurmdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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()

Expand All @@ -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)
Expand Down
1 change: 1 addition & 0 deletions test/check-application
Original file line number Diff line number Diff line change
Expand Up @@ -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
81 changes: 81 additions & 0 deletions test/unit/accounts_listing.test.py
Original file line number Diff line number Diff line change
@@ -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()
Loading