Add sort, order, and cursor options to workspace search methods (#182)

* Add sort, order, and cursor options to workspace search methods

* Add `Beaker.workspace.iter()`

CHANGELOG.md

@@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Added
+
+- Added `cursor`, `sort_by` and ordering options to `Beaker.workspace.*` search methods.
+
 ## [v1.12.1](https://github.com/allenai/beaker-py/releases/tag/v1.12.1) - 2022-12-08
 
 ### Fixed

beaker/config.py

@@ -35,7 +35,7 @@ class Config:
     The address of the Beaker server.
     """
 
-    default_org: Optional[str] = None
+    default_org: Optional[str] = "ai2"
     """
     Default Beaker organization to use.
     """

beaker/data_model/base.py

@@ -24,7 +24,7 @@
 logger = logging.getLogger("beaker")
 
 
-__all__ = ["BaseModel", "MappedSequence"]
+__all__ = ["BaseModel", "MappedSequence", "StrEnum"]
 
 
 BUG_REPORT_URL = (
@@ -144,3 +144,8 @@ def keys(self):
 
     def values(self):
         return self._mapping.values()
+
+
+class StrEnum(str, Enum):
+    def __str__(self) -> str:
+        return self.value

beaker/data_model/cluster.py

@@ -1,16 +1,15 @@
 from datetime import datetime
-from enum import Enum
 from typing import Optional, Tuple
 
 from pydantic import validator
 
-from .base import BaseModel
+from .base import BaseModel, StrEnum
 from .node import NodeResources, NodeUtilization
 
 __all__ = ["ClusterStatus", "Cluster", "ClusterUtilization", "ClusterSpec", "ClusterPatch"]
 
 
-class ClusterStatus(str, Enum):
+class ClusterStatus(StrEnum):
     """
     Current status of a cluster.
     """

beaker/data_model/dataset.py

@@ -4,7 +4,7 @@
 from pydantic import validator
 
 from .account import Account
-from .base import BaseModel
+from .base import BaseModel, StrEnum
 from .workspace import WorkspaceRef
 
 __all__ = [
@@ -18,6 +18,7 @@
     "DatasetsPage",
     "DatasetSpec",
     "DatasetPatch",
+    "DatasetSort",
 ]
 
 
@@ -168,3 +169,10 @@ class DatasetPatch(BaseModel):
     name: Optional[str] = None
     description: Optional[str] = None
     commit: Optional[bool] = None
+
+
+class DatasetSort(StrEnum):
+    created = "created"
+    author = "author"
+    dataset_name = "name"
+    dataset_name_or_description = "nameOrDescription"

beaker/data_model/experiment.py

@@ -4,11 +4,11 @@
 from pydantic import Field
 
 from .account import Account
-from .base import BaseModel, MappedSequence
+from .base import BaseModel, MappedSequence, StrEnum
 from .job import Job
 from .workspace import WorkspaceRef
 
-__all__ = ["Experiment", "Task", "Tasks", "ExperimentsPage", "ExperimentPatch"]
+__all__ = ["Experiment", "Task", "Tasks", "ExperimentsPage", "ExperimentPatch", "ExperimentSort"]
 
 
 class Experiment(BaseModel):
@@ -70,3 +70,10 @@ class ExperimentsPage(BaseModel):
 class ExperimentPatch(BaseModel):
     name: Optional[str] = None
     description: Optional[str] = None
+
+
+class ExperimentSort(StrEnum):
+    created = "created"
+    author = "author"
+    experiment_name = "name"
+    experiment_name_or_description = "nameOrDescription"

beaker/data_model/experiment_spec.py

@@ -1,11 +1,10 @@
-from enum import Enum
 from typing import Any, Dict, List, Optional
 
 from pydantic import Field, root_validator, validator
 
 from ..aliases import PathOrStr
 from ..exceptions import *
-from .base import BaseModel
+from .base import BaseModel, StrEnum
 
 __all__ = [
     "ImageSource",
@@ -238,7 +237,7 @@ class TaskResources(BaseModel, frozen=False):
     """
 
 
-class Priority(str, Enum):
+class Priority(StrEnum):
     urgent = "urgent"
     high = "high"
     normal = "normal"
@@ -305,7 +304,7 @@ class TaskSpec(BaseModel, frozen=False):
     Context describes how and where this task should run.
     """
 
-    constraints: Optional[Dict[str, List[str]]]
+    constraints: Optional[Dict[str, List[str]]] = None
     """
     Each task can have many constraints. And each constraint can have many values.
     Constraints are rules that change where a task is executed,
@@ -598,7 +597,7 @@ def with_constraint(self, **kwargs: List[str]) -> "TaskSpec":
         )
 
 
-class SpecVersion(str, Enum):
+class SpecVersion(StrEnum):
     v2 = "v2"
     v2_alpha = "v2-alpha"
 

beaker/data_model/group.py

@@ -1,12 +1,19 @@
 from datetime import datetime
-from enum import Enum
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 from .account import Account
-from .base import BaseModel
+from .base import BaseModel, StrEnum
 from .workspace import WorkspaceRef
 
-__all__ = ["Group", "GroupSpec", "GroupParameterType", "GroupParameter", "GroupPatch"]
+__all__ = [
+    "Group",
+    "GroupSpec",
+    "GroupParameterType",
+    "GroupParameter",
+    "GroupPatch",
+    "GroupsPage",
+    "GroupSort",
+]
 
 
 class Group(BaseModel):
@@ -31,7 +38,7 @@ class GroupSpec(BaseModel):
     experiments: Optional[List[str]] = None
 
 
-class GroupParameterType(str, Enum):
+class GroupParameterType(StrEnum):
     metric = "metric"
     env = "env"
 
@@ -47,3 +54,17 @@ class GroupPatch(BaseModel):
     add_experiments: Optional[List[str]] = None
     remove_experiments: Optional[List[str]] = None
     parameters: Optional[List[GroupParameter]] = None
+
+
+class GroupsPage(BaseModel):
+    data: Tuple[Group, ...]
+    next_cursor: Optional[str] = None
+    next: Optional[str] = None
+
+
+class GroupSort(StrEnum):
+    created = "created"
+    modified = "modified"
+    author = "author"
+    group_name = "name"
+    group_name_or_description = "nameOrDescription"

beaker/data_model/image.py

@@ -1,11 +1,10 @@
 from datetime import datetime
-from enum import Enum
 from typing import Optional, Tuple
 
 from pydantic import validator
 
 from .account import Account
-from .base import BaseModel
+from .base import BaseModel, StrEnum
 from .workspace import WorkspaceRef
 
 __all__ = [
@@ -20,6 +19,7 @@
     "DockerLayerDownloadState",
     "ImageSpec",
     "ImagePatch",
+    "ImageSort",
 ]
 
 
@@ -73,15 +73,15 @@ class DockerLayerProgress(BaseModel):
     total: Optional[int] = None
 
 
-class DockerLayerUploadStatus(str, Enum):
+class DockerLayerUploadStatus(StrEnum):
     preparing = "preparing"
     waiting = "waiting"
     pushing = "pushing"
     pushed = "pushed"
     already_exists = "layer already exists"
 
 
-class DockerLayerDownloadStatus(str, Enum):
+class DockerLayerDownloadStatus(StrEnum):
     waiting = "waiting"
     downloading = "downloading"
     download_complete = "download complete"
@@ -122,3 +122,9 @@ class ImagePatch(BaseModel):
     name: Optional[str] = None
     description: Optional[str] = None
     commit: Optional[bool] = None
+
+
+class ImageSort(StrEnum):
+    created = "created"
+    author = "author"
+    image_name = "name"

beaker/data_model/job.py

@@ -5,7 +5,7 @@
 from pydantic import Field, validator
 
 from .account import Account
-from .base import BaseModel
+from .base import BaseModel, StrEnum
 from .experiment_spec import DataMount, EnvVar, ImageSource, Priority, TaskSpec
 
 __all__ = [
@@ -25,7 +25,7 @@
 ]
 
 
-class CurrentJobStatus(str, Enum):
+class CurrentJobStatus(StrEnum):
     """
     The status of a job.
     """
@@ -41,7 +41,7 @@ class CurrentJobStatus(str, Enum):
     preempted = "preempted"
 
 
-class CanceledCode(int, Enum):
+class CanceledCode(Enum):
     not_set = 0
     system_preemption = 1
     user_preemption = 2
@@ -125,7 +125,7 @@ class JobExecution(BaseModel):
     workspace: Optional[str] = None
 
 
-class JobKind(str, Enum):
+class JobKind(StrEnum):
     """
     The kind of job.
     """

beaker/data_model/organization.py

@@ -1,9 +1,8 @@
 from datetime import datetime
-from enum import Enum
 from typing import Optional
 
 from .account import Account
-from .base import BaseModel
+from .base import BaseModel, StrEnum
 
 __all__ = ["Organization", "OrganizationRole", "OrganizationMember"]
 
@@ -17,7 +16,7 @@ class Organization(BaseModel):
     pronouns: Optional[str] = None
 
 
-class OrganizationRole(str, Enum):
+class OrganizationRole(StrEnum):
     admin = "admin"
     member = "member"
 

beaker/data_model/workspace.py

@@ -1,9 +1,8 @@
 from datetime import datetime
-from enum import Enum
 from typing import Dict, List, Optional, Tuple
 
 from .account import Account
-from .base import BaseModel
+from .base import BaseModel, StrEnum
 
 __all__ = [
     "WorkspaceSize",
@@ -17,6 +16,7 @@
     "WorkspacePatch",
     "WorkspacePermissionsPatch",
     "WorkspaceClearResult",
+    "WorkspaceSort",
 ]
 
 
@@ -64,7 +64,7 @@ class WorkspaceTransferSpec(BaseModel):
     ids: List[str]
 
 
-class Permission(str, Enum):
+class Permission(StrEnum):
     """
     Workspace permission levels.
     """
@@ -101,3 +101,9 @@ class WorkspaceClearResult(BaseModel):
     images_deleted: int = 0
     datasets_deleted: int = 0
     secrets_deleted: int = 0
+
+
+class WorkspaceSort(StrEnum):
+    created = "created"
+    modified = "modified"
+    workspace_name = "name"

beaker/exceptions.py

@@ -55,6 +55,7 @@
     "TaskStoppedError",
     "JobFailedError",
     "JobTimeoutError",
+    "ExperimentSpecError",
 ]
 
 

beaker/services/__init__.py

@@ -10,3 +10,18 @@
 from .secret import SecretClient
 from .service_client import ServiceClient
 from .workspace import WorkspaceClient
+
+__all__ = [
+    "AccountClient",
+    "ClusterClient",
+    "DatasetClient",
+    "ExperimentClient",
+    "GroupClient",
+    "ImageClient",
+    "JobClient",
+    "NodeClient",
+    "OrganizationClient",
+    "SecretClient",
+    "ServiceClient",
+    "WorkspaceClient",
+]

beaker/services/dataset.py

@@ -2,7 +2,16 @@
 import os
 from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING, Deque, Dict, Generator, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    ClassVar,
+    Deque,
+    Dict,
+    Generator,
+    Optional,
+    Tuple,
+    Union,
+)
 
 from ..aliases import PathOrStr
 from ..data_model import *
@@ -27,7 +36,7 @@ class DatasetClient(ServiceClient):
     HEADER_LAST_MODIFIED = "Last-Modified"
     HEADER_CONTENT_LENGTH = "Content-Length"
 
-    REQUEST_SIZE_LIMIT = 32 * 1024 * 1024
+    REQUEST_SIZE_LIMIT: ClassVar[int] = 32 * 1024 * 1024
 
     def get(self, dataset: str) -> Dataset:
         """