fix: Regression in "admin vfolders" command (#140)

Author: achimnol

changes/140.fix

@@ -0,0 +1 @@
+Fix regression of the `admin vfolders` command since introduction of paginated GraphQL queries

src/ai/backend/client/cli/admin/vfolders.py

@@ -5,16 +5,24 @@
 from tabulate import tabulate
 
 from . import admin
-from ...session import Session
+from ..pagination import (
+    get_preferred_page_size,
+    echo_via_pager,
+    tabulate_items,
+)
 from ..pretty import print_error
+from ...session import Session
+from ...exceptions import NoItems
 
 
 @admin.group(invoke_without_command=True)
 @click.pass_context
 @click.option('--access-key', type=str, default=None,
               help='Get vfolders for the given access key '
                    '(only works if you are a super-admin)')
-def vfolders(ctx, access_key):
+@click.option('-g', '--group', type=str, default=None,
+              help='Filter by group ID.')
+def vfolders(ctx, access_key, group):
     '''
     List and manage virtual folders.
     '''
@@ -23,25 +31,34 @@ def vfolders(ctx, access_key):
 
     fields = [
         ('Name', 'name'),
+        ('Host', 'host'),
+        ('Group', 'group'),
+        ('Creator', 'creator'),
+        ('Permission', 'permission'),
+        ('Usage Mode', 'usage_mode'),
+        ('Ownership Type', 'ownership_type'),
         ('Created At', 'created_at'),
         ('Last Used', 'last_used'),
-        ('Max Files', 'max_files'),
         ('Max Size', 'max_size'),
     ]
-    if access_key is None:
-        q = 'query { vfolders { $fields } }'
-    else:
-        q = 'query($ak:String) { vfolders(access_key:$ak) { $fields } }'
-    q = q.replace('$fields', ' '.join(item[1] for item in fields))
-    v = {'ak': access_key}
-    with Session() as session:
-        try:
-            resp = session.Admin.query(q, v)
-        except Exception as e:
-            print_error(e)
-            sys.exit(1)
-        print(tabulate((item.values() for item in resp['vfolders']),
-                       headers=(item[0] for item in fields)))
+
+    try:
+        with Session() as session:
+            page_size = get_preferred_page_size()
+            try:
+                items = session.VFolder.paginated_list(
+                    group, access_key,
+                    fields=[f[1] for f in fields],
+                    page_size=page_size,
+                )
+                echo_via_pager(
+                    tabulate_items(items, fields)
+                )
+            except NoItems:
+                print("There are no matching vfolders.")
+    except Exception as e:
+        print_error(e)
+        sys.exit(1)
 
 
 @vfolders.command()

src/ai/backend/client/func/vfolder.py

@@ -1,6 +1,7 @@
 import asyncio
 from pathlib import Path
 from typing import (
+    AsyncIterator,
     Sequence,
     Union
 )
@@ -15,12 +16,23 @@
 from .base import api_function, BaseFunction
 from ..compat import current_loop
 from ..config import DEFAULT_CHUNK_SIZE, MAX_INFLIGHT_CHUNKS
+from ..pagination import generate_paginated_results
 from ..request import Request
 
 __all__ = (
     'VFolder',
 )
 
+_default_list_fields = (
+    'host',
+    'name',
+    'created_at',
+    'creator',
+    'group',
+    'permission',
+    'ownership_type',
+)
+
 
 class VFolder(BaseFunction):
 
@@ -70,6 +82,33 @@ async def list(cls, list_all=False):
         async with rqst.fetch() as resp:
             return await resp.json()
 
+    @api_function
+    @classmethod
+    async def paginated_list(
+        cls,
+        group: str = None,
+        access_key: str = None,
+        *,
+        fields: Sequence[str] = _default_list_fields,
+        page_size: int = 20,
+    ) -> AsyncIterator[dict]:
+        """
+        Fetches the list of vfolders. Domain admins can only get domain vfolders.
+
+        :param group: Fetch vfolders in a specific group.
+        :param fields: Additional per-vfolder query fields to fetch.
+        """
+        async for item in generate_paginated_results(
+            'vfolder_list',
+            {
+                'group_id': (group, 'UUID'),
+                'access_key': (access_key, 'String'),
+            },
+            fields,
+            page_size=page_size,
+        ):
+            yield item
+
     @api_function
     @classmethod
     async def list_hosts(cls):