Merge pull request #109 from NessieCanCode/update-accounts-dropdown-in-settings

List all cluster accounts in settings

Author: NessieCanCode

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',
             '-',
           ];

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)

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

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