diff --git a/src/slurmdb.py b/src/slurmdb.py index 9beb072..01a482a 100644 --- a/src/slurmdb.py +++ b/src/slurmdb.py @@ -241,8 +241,13 @@ def fetch_usage_records(self, start_time, end_time): if not cur.fetchone(): cpu_col = "cpus_req" + job_name_col = "job_name" + cur.execute(f"SHOW COLUMNS FROM {job_table} LIKE %s", (job_name_col,)) + if not cur.fetchone(): + job_name_col = "name" + query = ( - f"SELECT j.id_job AS jobid, j.name AS job_name, j.account, j.partition, " + f"SELECT j.id_job AS jobid, j.{job_name_col} AS job_name, j.account, j.partition, " f"a.user AS user_name, j.time_start, j.time_end, j.tres_req, j.tres_alloc, " f"j.{cpu_col} AS cpus_alloc, j.state FROM {job_table} AS j " f"LEFT JOIN {assoc_table} AS a ON j.id_assoc = a.id_assoc " diff --git a/test/unit/slurmdb_validation.test.py b/test/unit/slurmdb_validation.test.py index 5326525..36192d4 100644 --- a/test/unit/slurmdb_validation.test.py +++ b/test/unit/slurmdb_validation.test.py @@ -1,7 +1,7 @@ import unittest import json from slurmdb import SlurmDB -from slurm_schema import extract_schema +from slurm_schema import extract_schema, extract_schema_from_dump class SlurmDBValidationTests(unittest.TestCase): def test_invalid_cluster_rejected(self): @@ -181,7 +181,14 @@ def fake_connect(): def test_fetch_usage_records_uses_cpus_req_if_alloc_missing(self): with open('test/example_slurm_schema_for_testing.json') as fh: schema = json.load(fh) + schema_sql = extract_schema_from_dump('test/example_slurmdb_for_testing.sql') job_cols = schema.get('localcluster_job_table', []) + job_cols_sql = schema_sql.get('localcluster_job_table', []) + # ensure schema JSON and SQL dump agree on CPU columns + self.assertIn('cpus_req', job_cols) + self.assertIn('cpus_req', job_cols_sql) + self.assertNotIn('cpus_alloc', job_cols) + self.assertNotIn('cpus_alloc', job_cols_sql) class FakeCursor: def __init__(self): @@ -222,7 +229,60 @@ def cursor(self): db.connect = lambda: None db.fetch_usage_records(0, 0) queries = db._conn.cursor_obj.queries - self.assertIn("j.cpus_req AS cpus_alloc", queries[1]) + self.assertIn("j.cpus_req AS cpus_alloc", queries[-1]) + + def test_fetch_usage_records_uses_job_name_column(self): + with open('test/example_slurm_schema_for_testing.json') as fh: + schema = json.load(fh) + schema_sql = extract_schema_from_dump('test/example_slurmdb_for_testing.sql') + job_cols = schema.get('localcluster_job_table', []) + job_cols_sql = schema_sql.get('localcluster_job_table', []) + # ensure schema JSON and SQL dump agree on job name column + self.assertIn('job_name', job_cols) + self.assertIn('job_name', job_cols_sql) + self.assertNotIn('name', job_cols) + self.assertNotIn('name', job_cols_sql) + + class FakeCursor: + def __init__(self): + self.queries = [] + + def execute(self, query, params=None): + self.queries.append(query) + if query.lower().startswith("show columns"): + column = params[0] if params else None + if column in job_cols: + self._fetchone = {'Field': column} + else: + self._fetchone = None + else: + self._fetchall = [] + + def fetchone(self): + return getattr(self, "_fetchone", None) + + def fetchall(self): + return getattr(self, "_fetchall", []) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + pass + + class FakeConn: + def __init__(self): + self.cursor_obj = FakeCursor() + + def cursor(self): + return self.cursor_obj + + db = SlurmDB(cluster="localcluster") + db._conn = FakeConn() + db.connect = lambda: None + db.fetch_usage_records(0, 0) + queries = db._conn.cursor_obj.queries + self.assertIn("j.job_name AS job_name", queries[-1]) if __name__ == '__main__': unittest.main()