From e78c96c45d3021120384c12286b7208a53358f6e Mon Sep 17 00:00:00 2001
From: Brett Smith <brett.smith@curii.com>
Date: Sat, 26 Aug 2023 23:59:21 -0400
Subject: [PATCH] 20839: Add order_key to select in keyset_list_all

Like uuid, the main loop requires this field.

Arvados-DCO-1.1-Signed-off-by: Brett Smith <brett.smith@curii.com>
---
 sdk/python/arvados/util.py    | 10 ++++++--
 sdk/python/tests/test_util.py | 43 ++++++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/sdk/python/arvados/util.py b/sdk/python/arvados/util.py
index 5bce8d3f83..2f145e7343 100644
--- a/sdk/python/arvados/util.py
+++ b/sdk/python/arvados/util.py
@@ -397,8 +397,14 @@ def keyset_list_all(fn, order_key="created_at", num_retries=0, ascending=True, *
     kwargs["order"] = ["%s %s" % (order_key, asc), "uuid %s" % asc]
     other_filters = kwargs.get("filters", [])
 
-    if "select" in kwargs and "uuid" not in kwargs["select"]:
-        kwargs["select"].append("uuid")
+    try:
+        select = set(kwargs['select'])
+    except KeyError:
+        pass
+    else:
+        select.add(order_key)
+        select.add('uuid')
+        kwargs['select'] = list(select)
 
     nextpage = []
     tot = 0
diff --git a/sdk/python/tests/test_util.py b/sdk/python/tests/test_util.py
index 4dba9ce3dc..f111abcf4a 100644
--- a/sdk/python/tests/test_util.py
+++ b/sdk/python/tests/test_util.py
@@ -2,10 +2,14 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import itertools
 import os
+import parameterized
 import subprocess
 import unittest
 
+from unittest import mock
+
 import arvados
 import arvados.util
 
@@ -54,6 +58,12 @@ class KeysetTestHelper:
         self.n += 1
         return self.expect[self.n-1][1]
 
+_SELECT_FAKE_ITEM = {
+    'uuid': 'zzzzz-zyyyz-zzzzzyyyyywwwww',
+    'name': 'KeysetListAllTestCase.test_select mock',
+    'created_at': '2023-08-28T12:34:56.123456Z',
+}
+
 class KeysetListAllTestCase(unittest.TestCase):
     def test_empty(self):
         ks = KeysetTestHelper([[
@@ -163,7 +173,6 @@ class KeysetListAllTestCase(unittest.TestCase):
         ls = list(arvados.util.keyset_list_all(ks.fn, filters=[["foo", ">", "bar"]]))
         self.assertEqual(ls, [{"created_at": "1", "uuid": "1"}, {"created_at": "2", "uuid": "2"}])
 
-
     def test_onepage_desc(self):
         ks = KeysetTestHelper([[
             {"limit": 1000, "count": "none", "order": ["created_at desc", "uuid desc"], "filters": []},
@@ -175,3 +184,35 @@ class KeysetListAllTestCase(unittest.TestCase):
 
         ls = list(arvados.util.keyset_list_all(ks.fn, ascending=False))
         self.assertEqual(ls, [{"created_at": "2", "uuid": "2"}, {"created_at": "1", "uuid": "1"}])
+
+    @parameterized.parameterized.expand(zip(
+        itertools.cycle(_SELECT_FAKE_ITEM),
+        itertools.chain.from_iterable(
+            itertools.combinations(_SELECT_FAKE_ITEM, count)
+            for count in range(len(_SELECT_FAKE_ITEM) + 1)
+        ),
+    ))
+    def test_select(self, order_key, select):
+        # keyset_list_all must have both uuid and order_key to function.
+        # Test that it selects those fields along with user-specified ones.
+        expect_select = {'uuid', order_key, *select}
+        item = {
+            key: value
+            for key, value in _SELECT_FAKE_ITEM.items()
+            if key in expect_select
+        }
+        list_func = mock.Mock()
+        list_func().execute = mock.Mock(
+            side_effect=[
+                {'items': [item]},
+                {'items': []},
+                {'items': []},
+            ],
+        )
+        list_func.reset_mock()
+        actual = list(arvados.util.keyset_list_all(list_func, order_key, select=list(select)))
+        self.assertEqual(actual, [item])
+        calls = list_func.call_args_list
+        self.assertTrue(len(calls) >= 2, "list_func() not called enough to exhaust items")
+        for call in calls:
+            self.assertEqual(set(call.kwargs.get('select', ())), expect_select)
-- 
2.30.2