From 627d0341c09a7fcd42b6c5113da690ad12a3ebc1 Mon Sep 17 00:00:00 2001
From: Jaime Frey <jfrey@cs.wisc.edu>
Date: Thu, 14 Nov 2024 10:52:45 -0600
Subject: [PATCH] HTCONDOR-2688 Fix ClassAd passing in condor_ce_jobmetrics

v2 ClassAds can't be pickled, so unparse them to move betwen processes.
Also fix a few imports of v1 bindings missed in previous round and stop
using xquery().
---
 src/condor_ce_jobmetrics    | 10 +++++++---
 src/htcondorce/web.py       |  6 +++---
 src/htcondorce/web_utils.py |  2 +-
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/condor_ce_jobmetrics b/src/condor_ce_jobmetrics
index 7e064adc..9c064c97 100755
--- a/src/condor_ce_jobmetrics
+++ b/src/condor_ce_jobmetrics
@@ -81,7 +81,7 @@ def process_one_schedd(ad):
     except RuntimeError as e:
         print(f"Failure to query CE {ad.get('Name', 'Unknown')}: {e}", file=sys.stderr)
     return [total_results, vo_results, job_count,
-            gpu_total_results, gpu_job_count, ad]
+            gpu_total_results, gpu_job_count, str(ad)]
 
 
 def get_ads(spooldir):
@@ -119,7 +119,7 @@ def list_job_results_final(job_results):
 def main():
     opts = parse_opts()
 
-    htcondor = __import__("htcondor")
+    htcondor = __import__("htcondor2")
     # Check if the CEVIEW daemon is in the daemon list, exit if it is not
     if htcondor.param["DAEMON_LIST"].find("CEVIEW") == -1:
         print("CEVIEW is not in DAEMON_LIST, exiting...", file=sys.stderr)
@@ -136,8 +136,10 @@ def main():
 
     map_results = pool.map(process_one_schedd, ads, 1)
 
+    import classad2 as classad
+
     for (schedd_total_results, schedd_vo_results, schedd_job_results,
-         schedd_gpu_total_results, schedd_gpu_job_results, ad) in map_results:
+         schedd_gpu_total_results, schedd_gpu_job_results, ad_str) in map_results:
 
         dict_sum1(schedd_total_results,     total_results)
         dict_sum1(schedd_gpu_total_results, gpu_total_results)
@@ -146,6 +148,8 @@ def main():
         dict_sum2(schedd_job_results,       job_results)
         dict_sum2(schedd_gpu_job_results,   gpu_job_results)
 
+        ad = classad.ClassAd(ad_str)
+
         total_fname = rrd.check_rrd(environ, ad['Name'], "jobs")
         if os.isatty(2):
             print(f"Schedd {ad['Name']} totals: {schedd_total_results}", file=sys.stderr)
diff --git a/src/htcondorce/web.py b/src/htcondorce/web.py
index 105b76b3..266964a7 100644
--- a/src/htcondorce/web.py
+++ b/src/htcondorce/web.py
@@ -171,7 +171,7 @@ def totals_ce_json(environ, start_response):
     objs = htcondorce.web_utils.get_schedd_objs(environ)
     results = {"Running": 0, "Idle": 0, "Held": 0, "UpdateDate": time.time()}
     for schedd, name in objs:
-        for job in schedd.xquery("true", ["JobStatus"]):
+        for job in schedd.query("true", ["JobStatus"]):
             if job.get("JobStatus") == 1:
                 results['Idle'] += 1
             elif job.get("JobStatus") == 2:
@@ -194,7 +194,7 @@ def pilots_ce_json(environ, start_response):
     objs = htcondorce.web_utils.get_schedd_objs(environ)
     job_count = {}
     for schedd, name in objs:
-        for job in schedd.xquery('true', ['x509UserProxyVOName', 'x509UserProxyFirstFQAN', 'JobStatus', 'x509userproxysubject']):
+        for job in schedd.query('true', ['x509UserProxyVOName', 'x509UserProxyFirstFQAN', 'JobStatus', 'x509userproxysubject']):
             DN = job.get("x509userproxysubject", 'Unknown')
             VO = job.get('x509UserProxyVOName', 'Unknown')
             VOMS = job.get('x509UserProxyFirstFQAN', '').replace("/Capability=NULL", "").replace("/Role=NULL", "")
@@ -225,7 +225,7 @@ def vos_ce_json(environ, start_response):
     objs = htcondorce.web_utils.get_schedd_objs(environ)
     job_count = {}
     for schedd, name in objs:
-        for job in schedd.xquery('true', ['x509UserProxyVOName', 'JobStatus']):
+        for job in schedd.query('true', ['x509UserProxyVOName', 'JobStatus']):
             VO = job.get('x509UserProxyVOName', 'Unknown')
             job_key = VO
             if job_key not in job_count:
diff --git a/src/htcondorce/web_utils.py b/src/htcondorce/web_utils.py
index 2bf88080..aea9520c 100644
--- a/src/htcondorce/web_utils.py
+++ b/src/htcondorce/web_utils.py
@@ -11,7 +11,7 @@ def check_htcondor():
     global htcondor
     if not htcondor:
         os.environ.setdefault('CONDOR_CONFIG', "/etc/condor-ce/condor_config")
-        htcondor = __import__("htcondor")
+        htcondor = __import__("htcondor2")
     return htcondor