From bd840d81e002a36e350f269f0337f6adaadc2713 Mon Sep 17 00:00:00 2001
From: Mike McLean <mikem@redhat.com>
Date: Jan 06 2021 21:33:24 +0000
Subject: handle plugins and generator results in count and countAndFilterResults


Fixes: https://pagure.io/koji/issue/2632

---

diff --git a/hub/kojihub.py b/hub/kojihub.py
index f13d071..8e7187d 100644
--- a/hub/kojihub.py
+++ b/hub/kojihub.py
@@ -42,6 +42,7 @@ import sys
 import tarfile
 import tempfile
 import time
+import types
 import traceback
 from urllib.parse import parse_qs
 import xmlrpc.client
@@ -12656,15 +12657,16 @@ class RootExports(object):
         A method return value of None will return O, a return value of type "list", "tuple", or
         "dict" will return len(value), and a return value of any other type will return 1. An
         invalid methodName will raise GenericError."""
-        try:
-            method = getattr(self, methodName)
-        except AttributeError:
+        handler = context.handlers.get(methodName)
+        if not handler:
             raise koji.GenericError("method %s doesn't exist" % methodName)
-        result = method(*args, **kw)
+        result = handler(*args, **kw)
         if result is None:
             return 0
         elif isinstance(result, (list, tuple, dict)):
             return len(result)
+        elif isinstance(result, types.GeneratorType):
+            return sum(1 for r in result)
         else:
             return 1
 
@@ -12718,15 +12720,18 @@ class RootExports(object):
         """
         filterOpts = kw.pop('filterOpts', {})
 
-        try:
-            method = getattr(self, methodName)
-        except AttributeError:
+        handler = context.handlers.get(methodName)
+        if not handler:
             raise koji.GenericError("method %s doesn't exist" % methodName)
         try:
-            results = method(*args, **kw)
+            results = handler(*args, **kw)
         except Exception as ex:
             raise koji.GenericError("method %s raised an exception (%s)" % (methodName, str(ex)))
 
+        if isinstance(results, types.GeneratorType):
+            # unfortunately due to the nature of the return, we have to generate the full list
+            results = list(results)
+
         if results is None:
             return 0, None
         elif isinstance(results, list):