feat: Add queryfilter support (#171)

achimnol · web-flow · commit 74cd26cd6369 · 2021-07-19T16:56:12.000+09:00
* feat: Add `filter` and `order` args to functional API and CLI,
  while keeping backward compatibility with older API versions.
* feat: Add explicit `BackendAPIVersionError` to indicate mismatch
  functionality of different API versions (e.g., when `filter` is
  given for legacy API sessions).
* fix: Handle standard-compliant GQL erros like before
diff --git a/changes/171.feature b/changes/171.feature
@@ -0,0 +1 @@
+Add support for queryfilter (`--filter`) and ordering (`--order`) arguments for paginated lists (e.g., `ps`, `admin agents`, `admin users` commands)
diff --git a/src/ai/backend/client/cli/admin/agents.py b/src/ai/backend/client/cli/admin/agents.py
@@ -118,7 +118,11 @@ def agent(agent_id):
               help='Filter agents by the given status.')
 @click.option('--scaling-group', '--sgroup', type=str, default=None,
               help='Filter agents by the scaling group.')
-def agents(status, scaling_group):
+@click.option('--filter', 'filter_', default=None,
+              help='Set the query filter expression.')
+@click.option('--order', default=None,
+              help='Set the query ordering expression.')
+def agents(status, scaling_group, filter_, order):
     '''
     List and manage agents.
     (super-admin privilege required)
@@ -151,6 +155,8 @@ def format_item(item):
                     scaling_group,
                     fields=[f[1] for f in fields],
                     page_size=page_size,
+                    filter=filter_,
+                    order=order,
                 )
                 echo_via_pager(
                     tabulate_items(items, fields,
diff --git a/src/ai/backend/client/cli/admin/keypairs.py b/src/ai/backend/client/cli/admin/keypairs.py
@@ -56,7 +56,11 @@ def keypair():
               help='Show keypairs of this given user. [default: show all]')
 @click.option('--is-active', type=bool, default=None,
               help='Filter keypairs by activation.')
-def keypairs(ctx, user_id, is_active):
+@click.option('--filter', 'filter_', default=None,
+              help='Set the query filter expression.')
+@click.option('--order', default=None,
+              help='Set the query ordering expression.')
+def keypairs(ctx, user_id, is_active, filter_, order):
     '''
     List and manage keypairs.
     To show all keypairs or other user's, your access key must have the admin
@@ -96,6 +100,8 @@ def format_item(item):
                     is_active,
                     fields=[f[1] for f in fields],
                     page_size=page_size,
+                    filter=filter_,
+                    order=order,
                 )
                 echo_via_pager(
                     tabulate_items(items, fields, item_formatter=format_item)
diff --git a/src/ai/backend/client/cli/admin/sessions.py b/src/ai/backend/client/cli/admin/sessions.py
@@ -85,10 +85,12 @@ def transform_fields(item: SessionItem, *, in_row: bool = True) -> SessionItem:
 @click.option('--running', is_flag=True,
               help='Filter only scheduled and running sessions. Ignores --status option.')
 @click.option('--detail', is_flag=True, help='Show more details using more columns.')
-@click.option('-f', '--format', default=None,  help='Display only specified fields.')
+@click.option('-f', '--format', default=None, help='Display only specified fields.')
+@click.option('--filter', 'filter_', default=None, help='Set the query filter expression.')
+@click.option('--order', default=None, help='Set the query ordering expression.')
 @click.option('--plain', is_flag=True,
               help='Display the session list without decorative line drawings and the header.')
-def sessions(status, access_key, name_only, dead, running, detail, plain, format):
+def sessions(status, access_key, name_only, dead, running, detail, filter_, order, plain, format):
     '''
     List and manage compute sessions.
     '''
@@ -169,6 +171,8 @@ def sessions(status, access_key, name_only, dead, running, detail, plain, format
                     status, access_key,
                     fields=[f[1] for f in fields],
                     page_size=page_size,
+                    filter=filter_,
+                    order=order,
                 )
                 if name_only:
                     echo_via_pager(
diff --git a/src/ai/backend/client/cli/admin/storage.py b/src/ai/backend/client/cli/admin/storage.py
@@ -51,7 +51,11 @@ def storage(vfolder_host):
 
 
 @admin.command()
-def storage_list():
+@click.option('--filter', 'filter_', default=None,
+              help='Set the query filter expression.')
+@click.option('--order', default=None,
+              help='Set the query ordering expression.')
+def storage_list(filter_, order):
     """
     List storage volumes.
     (super-admin privilege required)
@@ -68,6 +72,8 @@ def storage_list():
                 items = session.Storage.paginated_list(
                     fields=[f[1] for f in fields],
                     page_size=page_size,
+                    filter=filter_,
+                    order=order,
                 )
                 echo_via_pager(
                     tabulate_items(items, fields)
diff --git a/src/ai/backend/client/cli/admin/users.py b/src/ai/backend/client/cli/admin/users.py
@@ -61,7 +61,11 @@ def user(email):
               help='Filter users in a specific state (active, inactive, deleted, before-verification).')
 @click.option('-g', '--group', type=str, default=None,
               help='Filter by group ID.')
-def users(ctx, status, group) -> None:
+@click.option('--filter', 'filter_', default=None,
+              help='Set the query filter expression.')
+@click.option('--order', default=None,
+              help='Set the query ordering expression.')
+def users(ctx, status, group, filter_, order) -> None:
     '''
     List and manage users.
     (admin privilege required)
@@ -94,6 +98,8 @@ def format_item(item):
                     status, group,
                     fields=[f[1] for f in fields],
                     page_size=page_size,
+                    filter=filter_,
+                    order=order,
                 )
                 echo_via_pager(
                     tabulate_items(items, fields,
diff --git a/src/ai/backend/client/cli/admin/vfolders.py b/src/ai/backend/client/cli/admin/vfolders.py
@@ -22,7 +22,11 @@
                    '(only works if you are a super-admin)')
 @click.option('-g', '--group', type=str, default=None,
               help='Filter by group ID.')
-def vfolders(ctx, access_key, group):
+@click.option('--filter', 'filter_', default=None,
+              help='Set the query filter expression.')
+@click.option('--order', default=None,
+              help='Set the query ordering expression.')
+def vfolders(ctx, access_key, group, filter_, order):
     '''
     List and manage virtual folders.
     '''
@@ -50,6 +54,8 @@ def vfolders(ctx, access_key, group):
                     group, access_key,
                     fields=[f[1] for f in fields],
                     page_size=page_size,
+                    filter=filter_,
+                    order=order,
                 )
                 echo_via_pager(
                     tabulate_items(items, fields)
diff --git a/src/ai/backend/client/cli/ps.py b/src/ai/backend/client/cli/ps.py
@@ -21,10 +21,12 @@
               help='Filter only scheduled and running sessions. Ignores --status option.')
 @click.option('--detail', is_flag=True, help='Show more details using more columns.')
 @click.option('-f', '--format', default=None,  help='Display only specified fields.')
+@click.option('--filter', 'filter_', default=None, help='Set the query filter expression.')
+@click.option('--order', default=None, help='Set the query ordering expression.')
 @click.option('--plain', is_flag=True,
               help='Display the session list without decorative line drawings and the header.')
 @click.pass_context
-def ps(ctx, status, name_only, dead, running, detail, plain, format):
+def ps(ctx, status, name_only, dead, running, detail, filter_, order, plain, format):
     '''
     Lists the current running compute sessions for the current keypair.
     This is an alias of the "admin sessions --status=RUNNING" command.
diff --git a/src/ai/backend/client/exceptions.py b/src/ai/backend/client/exceptions.py
@@ -44,6 +44,13 @@ def data(self) -> Any:
         return self.args[2]
 
 
+class BackendAPIVersionError(BackendError):
+    """
+    Exception indicating that the given operation/argument is not supported
+    in the currently negotiated server API version.
+    """
+
+
 class BackendClientError(BackendError):
     """
     Exceptions from the client library, such as argument validation
diff --git a/src/ai/backend/client/func/admin.py b/src/ai/backend/client/func/admin.py
@@ -61,7 +61,11 @@ async def _query(
                 response = await resp.json()
                 errors = response.get("errors", [])
                 if errors:
-                    raise BackendAPIError(400, reason="GraphQL-generated error", data=errors)
+                    raise BackendAPIError(400, reason="Bad request", data={
+                        'type': 'https://api.backend.ai/probs/graphql-error',
+                        'title': 'GraphQL-generated error',
+                        'data': errors,
+                    })
                 else:
                     return response["data"]
         else:
diff --git a/src/ai/backend/client/func/agent.py b/src/ai/backend/client/func/agent.py
@@ -57,6 +57,8 @@ async def paginated_list(
         *,
         fields: Sequence[str] = _default_list_fields,
         page_size: int = 20,
+        filter: str = None,
+        order: str = None,
     ) -> AsyncIterator[dict]:
         """
         Lists the keypairs.
@@ -67,6 +69,8 @@ async def paginated_list(
             {
                 'status': (status, 'String'),
                 'scaling_group': (scaling_group, 'String'),
+                'filter': (filter, 'String'),
+                'order': (order, 'String'),
             },
             fields,
             page_size=page_size,
diff --git a/src/ai/backend/client/func/keypair.py b/src/ai/backend/client/func/keypair.py
@@ -34,12 +34,15 @@ def __init__(self, access_key: str):
 
     @api_function
     @classmethod
-    async def create(cls, user_id: Union[int, str],
-                     is_active: bool = True,
-                     is_admin: bool = False,
-                     resource_policy: str = None,
-                     rate_limit: int = None,
-                     fields: Iterable[str] = None) -> dict:
+    async def create(
+        cls,
+        user_id: Union[int, str],
+        is_active: bool = True,
+        is_admin: bool = False,
+        resource_policy: str = None,
+        rate_limit: int = None,
+        fields: Iterable[str] = None,
+    ) -> dict:
         """
         Creates a new keypair with the given options.
         You need an admin privilege for this operation.
@@ -67,11 +70,14 @@ async def create(cls, user_id: Union[int, str],
 
     @api_function
     @classmethod
-    async def update(cls, access_key: str,
-                     is_active: bool = None,
-                     is_admin: bool = None,
-                     resource_policy: str = None,
-                     rate_limit: int = None) -> dict:
+    async def update(
+        cls,
+        access_key: str,
+        is_active: bool = None,
+        is_admin: bool = None,
+        resource_policy: str = None,
+        rate_limit: int = None,
+    ) -> dict:
         """
         Creates a new keypair with the given options.
         You need an admin privilege for this operation.
@@ -113,7 +119,8 @@ async def delete(cls, access_key: str):
     @api_function
     @classmethod
     async def list(
-        cls, user_id: Union[int, str] = None,
+        cls,
+        user_id: Union[int, str] = None,
         is_active: bool = None,
         fields: Sequence[str] = _default_list_fields,
     ) -> Sequence[dict]:
@@ -153,6 +160,8 @@ async def paginated_list(
         user_id: str = None,
         fields: Sequence[str] = _default_list_fields,
         page_size: int = 20,
+        filter: str = None,
+        order: str = None,
     ) -> AsyncIterator[dict]:
         """
         Lists the keypairs.
@@ -161,6 +170,8 @@ async def paginated_list(
         variables = {
             'is_active': (is_active, 'Boolean'),
             'domain_name': (domain_name, 'String'),
+            'filter': (filter, 'String'),
+            'order': (order, 'String'),
         }
         if user_id is not None:
             variables['email'] = (user_id, 'String')
diff --git a/src/ai/backend/client/func/session.py b/src/ai/backend/client/func/session.py
@@ -89,6 +89,8 @@ async def paginated_list(
         *,
         fields: Sequence[str] = None,
         page_size: int = 20,
+        filter: str = None,
+        order: str = None,
     ) -> AsyncIterator[dict]:
         """
         Fetches the list of users. Domain admins can only get domain users.
@@ -111,6 +113,8 @@ async def paginated_list(
             {
                 'status': (status, 'String'),
                 'access_key': (access_key, 'String'),
+                'filter': (filter, 'String'),
+                'order': (order, 'String'),
             },
             fields,
             page_size=page_size,
diff --git a/src/ai/backend/client/func/storage.py b/src/ai/backend/client/func/storage.py
@@ -48,14 +48,19 @@ async def paginated_list(
         *,
         fields: Sequence[str] = _default_list_fields,
         page_size: int = 20,
+        filter: str = None,
+        order: str = None,
     ) -> AsyncIterator[dict]:
         """
         Lists the keypairs.
         You need an admin privilege for this operation.
         """
         async for item in generate_paginated_results(
             'storage_volume_list',
-            {},
+            {
+                'filter': (filter, 'String'),
+                'order': (order, 'String'),
+            },
             fields,
             page_size=page_size,
         ):
diff --git a/src/ai/backend/client/func/user.py b/src/ai/backend/client/func/user.py
@@ -93,6 +93,8 @@ async def paginated_list(
         *,
         fields: Sequence[str] = _default_list_fields,
         page_size: int = 20,
+        filter: str = None,
+        order: str = None,
     ) -> AsyncIterator[dict]:
         """
         Fetches the list of users. Domain admins can only get domain users.
@@ -107,6 +109,8 @@ async def paginated_list(
             {
                 'status': (status, 'String'),
                 'group_id': (group, 'UUID'),
+                'filter': (filter, 'String'),
+                'order': (order, 'String'),
             },
             fields,
             page_size=page_size,
diff --git a/src/ai/backend/client/func/vfolder.py b/src/ai/backend/client/func/vfolder.py
@@ -91,6 +91,8 @@ async def paginated_list(
         *,
         fields: Sequence[str] = _default_list_fields,
         page_size: int = 20,
+        filter: str = None,
+        order: str = None,
     ) -> AsyncIterator[dict]:
         """
         Fetches the list of vfolders. Domain admins can only get domain vfolders.
@@ -103,6 +105,8 @@ async def paginated_list(
             {
                 'group_id': (group, 'UUID'),
                 'access_key': (access_key, 'String'),
+                'filter': (filter, 'String'),
+                'order': (order, 'String'),
             },
             fields,
             page_size=page_size,
diff --git a/src/ai/backend/client/pagination.py b/src/ai/backend/client/pagination.py
@@ -3,7 +3,7 @@
     cast,
     Any,
     AsyncIterator,
-    Mapping,
+    Dict,
     Sequence,
     Tuple,
 )
@@ -12,7 +12,7 @@
     TypedDict,
 )
 
-from .exceptions import NoItems
+from .exceptions import NoItems, BackendAPIVersionError
 from .session import api_session
 
 MAX_PAGE_SIZE: Final = 100
@@ -25,7 +25,7 @@ class PaginatedResult(TypedDict):
 
 async def execute_paginated_query(
     root_field: str,
-    variables: Mapping[str, Tuple[Any, str]],
+    variables: Dict[str, Tuple[Any, str]],
     fields: Sequence[str],
     *,
     limit: int,
@@ -63,13 +63,21 @@ async def execute_paginated_query(
 
 async def generate_paginated_results(
     root_field: str,
-    variables: Mapping[str, Tuple[Any, str]],
+    variables: Dict[str, Tuple[Any, str]],
     fields: Sequence[str],
     *,
     page_size: int,
 ) -> AsyncIterator[Any]:
     if page_size > MAX_PAGE_SIZE:
         raise ValueError(f"The page size cannot exceed {MAX_PAGE_SIZE}")
+    if api_session.get().api_version < (6, '20210815'):
+        if variables['filter'][0] is not None or variables['order'][0] is not None:
+            raise BackendAPIVersionError(
+                "filter and order arguments for paginated lists require v6.20210815 or later."
+            )
+        # should remove to work with older managers
+        variables.pop('filter')
+        variables.pop('order')
     offset = 0
     total_count = -1
     while True:

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Add support for queryfilter (`--filter`) and ordering (`--order`) arguments for paginated lists (e.g., `ps`, `admin agents`, `admin users` commands)