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
7 changes: 6 additions & 1 deletion src/slurmdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 "
Expand Down
64 changes: 62 additions & 2 deletions test/unit/slurmdb_validation.test.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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()
Loading