Merge pull request #11404 from vkuznet/fix-multirequest

Fix multirequest API to encode input parameters and properly handle gzip'ed content
dmwm · Dec 20, 2022 · fb39cd4 · fb39cd4
2 parents cd5d4c1 + 8d35eb0
commit fb39cd4
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 15 deletions.
diff --git a/src/python/WMCore/Services/pycurl_manager.py b/src/python/WMCore/Services/pycurl_manager.py
@@ -63,6 +63,27 @@
 from Utils.TokenManager import TokenManager
 
 
+def getException(url, params, headers, header, data):
+    """
+    Return HTTP exception for a given set of parameters:
+    :param url: string
+    :param params: dict
+    :param headers: dict
+    :param header: ResponseHeader
+    :param data: HTTP body
+    """
+    msg = 'url=%s, code=%s, reason=%s, headers=%s, result=%s' \
+          % (url, header.status, header.reason, header.header, data)
+    exc = http.client.HTTPException(msg)
+    setattr(exc, 'req_data', params)
+    setattr(exc, 'req_headers', headers)
+    setattr(exc, 'url', url)
+    setattr(exc, 'result', data)
+    setattr(exc, 'status', header.status)
+    setattr(exc, 'reason', header.reason)
+    setattr(exc, 'headers', header.header)
+    return exc
+
 def decompress(body, headers):
     """
     Helper function to decompress given body if HTTP headers contains gzip encoding
@@ -326,16 +347,7 @@ def request(self, url, params, headers=None, verb='GET',
             else:
                 data = self.parse_body(data, decode)
         else:
-            msg = 'url=%s, code=%s, reason=%s, headers=%s, result=%s' \
-                  % (url, header.status, header.reason, header.header, data)
-            exc = http.client.HTTPException(msg)
-            setattr(exc, 'req_data', params)
-            setattr(exc, 'req_headers', headers)
-            setattr(exc, 'url', url)
-            setattr(exc, 'result', data)
-            setattr(exc, 'status', header.status)
-            setattr(exc, 'reason', header.reason)
-            setattr(exc, 'headers', header.header)
+            exc = getException(url, params, headers, header, data)
             bbuf.flush()
             hbuf.flush()
             raise exc
@@ -361,15 +373,16 @@ def getheader(self, url, params, headers=None, verb='GET',
         return header
 
     @portForward(8443)
-    def multirequest(self, url, parray, headers=None,
-                     ckey=None, cert=None, verbose=None, cookie=None):
+    def multirequest(self, url, parray, headers=None, verb='GET',
+                     ckey=None, cert=None, verbose=None, cookie=None,
+                     encode=False, decode=False):
         """Fetch data for given set of parameters"""
         multi = pycurl.CurlMulti()
         for params in parray:
             curl = pycurl.Curl()
             bbuf, hbuf = \
                 self.set_opts(curl, url, params, headers, ckey=ckey, cert=cert,
-                              verbose=verbose, cookie=cookie)
+                              verbose=verbose, cookie=cookie, encode=encode)
             multi.add_handle(curl)
             while True:
                 ret, num_handles = multi.perform()
@@ -385,8 +398,20 @@ def multirequest(self, url, parray, headers=None,
                         break
             dummyNumq, response, dummyErr = multi.info_read()
             for _respItem in response:
-                data = decodeBytesToUnicode(bbuf.getvalue())
-                data = json.loads(data)
+                header = self.parse_header(hbuf.getvalue())
+                data = bbuf.getvalue()
+                data = decompress(data, header.header)
+                data = decodeBytesToUnicode(data)
+                if header.status < 300:
+                    if verb == 'HEAD':
+                        data = ''
+                    else:
+                        data = self.parse_body(data, decode)
+                else:
+                    exc = getException(url, params, headers, header, data)
+                    bbuf.flush()
+                    hbuf.flush()
+                    raise exc
                 if isinstance(data, dict):
                     data.update(params)
                     yield data

diff --git a/test/python/WMCore_t/Services_t/pycurl_manager_t.py b/test/python/WMCore_t/Services_t/pycurl_manager_t.py
@@ -288,6 +288,29 @@ def testHTTPResponse(self):
         obj = ResponseHeader(response)
         self.assertEqual(obj.status, 400)
 
+    def testMultirequest(self):
+        """
+        Test multirequest function
+        """
+        tfile = tempfile.NamedTemporaryFile()
+        url = "https://cmsweb-prod.cern.ch/dbs/prod/global/DBSReader/filelumis"
+        blocks = [
+            '/LQToDEle_M-4000_single_TuneCP2_13TeV-madgraph-pythia8/RunIISummer20UL17MiniAODv2-106X_mc2017_realistic_v9-v1/MINIAODSIM#471e5596-af04-4423-a850-5ef9091f154f',
+            '/LQToDEle_M-4000_single_TuneCP2_13TeV-madgraph-pythia8/RunIISummer20UL17MiniAODv2-106X_mc2017_realistic_v9-v1/MINIAODSIM#6eb03689-167a-472f-8b09-f4bfadad6a8a',
+            '/LQToDEle_M-4000_single_TuneCP2_13TeV-madgraph-pythia8/RunIISummer20UL17MiniAODv2-106X_mc2017_realistic_v9-v1/MINIAODSIM#b8cdec8f-b664-49a6-ab2d-bb2a89893581',
+            '/LQToDEle_M-4000_single_TuneCP2_13TeV-madgraph-pythia8/RunIISummer20UL17MiniAODv2-106X_mc2017_realistic_v9-v1/MINIAODSIM#ff78bb73-0e8c-41cb-9e51-381cfbdf15e2'
+        ]
+        cern_sso_cookie(url, tfile.name, self.cert, self.ckey)
+        parray = [{'block_name':b} for b in blocks]
+        headers = {'Accept': 'application/json'}
+        cookie = {url: tfile.name}
+        mgr = RequestHandler()
+        data = mgr.multirequest(url, parray, headers=headers, ckey=self.ckey, cert=self.cert, cookie=cookie, encode=True, decode=True)
+        pairs = set()
+        for row in data:
+            pair = (row['lumi_section_num'], row['run_num'])
+            pairs.add(pair)
+        self.assertTrue(len(pairs), 100)
 
 if __name__ == "__main__":
     unittest.main()