From 6eebd4499d7f33d2f544ca25d83adb4dcf707b40 Mon Sep 17 00:00:00 2001 From: Chris Barnett Date: Thu, 5 Dec 2024 14:51:26 -0500 Subject: [PATCH] added PATCH methods added additional models added folder endpoints --- vast_api_client/abstract_client.py | 16 +++++- vast_api_client/models.py | 36 +++++++++++++- vast_api_client/vast_api_client.py | 80 +++++++++++++++++++++++++----- 3 files changed, 117 insertions(+), 15 deletions(-) diff --git a/vast_api_client/abstract_client.py b/vast_api_client/abstract_client.py index 703973d..8fc698e 100644 --- a/vast_api_client/abstract_client.py +++ b/vast_api_client/abstract_client.py @@ -40,13 +40,23 @@ def _send_get_request(self, endpoint, params=None, retries=2): return r.json() def _send_post_request(self, endpoint, payload, headers=None, skip_auth=False): + return self._send_body('POST', endpoint, payload, headers, skip_auth) + + def _send_patch_request(self, endpoint, payload, headers=None, skip_auth=False): + return self._send_body('PATCH', endpoint, payload, headers, skip_auth) + + def _send_body(self, http_method, endpoint, payload, headers=None, skip_auth=False): if not skip_auth and self.token is None and self.refresh_token is not None: self.renew_token(self.refresh_token) + # list of valid http methods + if http_method not in ['POST', 'PUT', 'PATCH', 'DELETE']: + raise ValueError(f'Invalid http method: {http_method}') + headers = headers if headers is not None else {} headers.update({'Content-Type': 'application/json'}) - r = requests.post(Path(self.url, endpoint).as_posix(), + r = requests.request(http_method, Path(self.url, endpoint).as_posix(), json=payload, headers=self._get_headers(headers, skip_auth=skip_auth), verify=False, @@ -54,7 +64,9 @@ def _send_post_request(self, endpoint, payload, headers=None, skip_auth=False): r.raise_for_status() return r.json() - def _send_delete_request(self, endpoint): + def _send_delete_request(self, endpoint, body=None): + if body is not None: + return self._send_body('DELETE', endpoint, body) if self.token is None and self.refresh_token is not None: self.renew_token(self.refresh_token) diff --git a/vast_api_client/models.py b/vast_api_client/models.py index 1b869d6..0e0f170 100644 --- a/vast_api_client/models.py +++ b/vast_api_client/models.py @@ -2,7 +2,7 @@ from typing import Set, Optional from enum import Enum, IntEnum import re -from pydantic import (BaseModel, ConfigDict, PositiveInt, field_validator, +from pydantic import (BaseModel, ConfigDict, Field, PositiveInt, field_validator, model_validator, field_serializer, InstanceOf) @@ -223,5 +223,39 @@ def serialize_source_dir(self, source_dir: Path, _info): return source_dir.as_posix() +class PathBody(BaseModel): + model_config = ConfigDict(extra='allow', str_strip_whitespace=True) + path: Path + + @field_validator("path") + @classmethod + def is_valid_unix_path(cls, path: Path) -> Path: + return validate_path(path) + + +class FolderCreateOrUpdate(BaseModel): + model_config = ConfigDict(extra='forbid', str_strip_whitespace=True, frozen=True) + path: Path + owner_is_group: bool = False + user: str = None + group: str = None + + @field_validator("path") + @classmethod + def is_valid_unix_path(cls, path: Path) -> Path: + return validate_path(path) + + +class QuotaUpdate(BaseModel): + model_config = ConfigDict(extra='forbid', str_strip_whitespace=True) + soft_limit: Optional[PositiveInt] = None + hard_limit: PositiveInt + @model_validator(mode="after") + def soft_limit_below_hard_limit(self) -> 'QuotaUpdate': + if self.soft_limit is None: + self.soft_limit = self.hard_limit + if self.hard_limit < self.soft_limit: + raise ValueError("'soft_limit' cannot be larger than 'hard_limit'") + return self diff --git a/vast_api_client/vast_api_client.py b/vast_api_client/vast_api_client.py index a03f0af..38077c0 100644 --- a/vast_api_client/vast_api_client.py +++ b/vast_api_client/vast_api_client.py @@ -68,18 +68,17 @@ def get_quotas(self, path: Path = None): :return: list of quotas """ if path is not None: - return self._send_get_request('quotas/', params={'path': path.as_posix()}) + body = PathBody(path=path).model_dump() + return self._send_get_request('quotas/', params=body) else: return self._send_get_request('quotas/') def get_views(self, path: Path = None): if path is not None: - return self._send_get_request('views/', params={'path': path.as_posix()}) + body = PathBody(path=path).model_dump() + return self._send_get_request('views/', params=body) return self._send_get_request('views/') - def get_status(self): - return self._send_get_request('latest/dashboard/status/') - def add_view(self, name: str, path: Path, protocols: set[ProtocolEnum] = {}, policy_id: PolicyEnum = None, @@ -130,9 +129,71 @@ def add_quota(self, name: str, path: Path, else: return self._send_post_request('quotas/', qc.model_dump()) + def add_folder(self, path: Path, group: str, user: str = None, owner_is_group: bool = False): + """ + Add a folder to the storage system + :param path: path of the folder + :param group: group that owns the folder + :param user: user that owns the folder + :param owner_is_group: if True, the owner is a group + """ + folder = FolderCreateOrUpdate(path=path, owner_is_group=owner_is_group, user=user, group=group) + return self._send_post_request('folders/create_folder/', folder.model_dump()) + + def modify_folder(self, path: Path, group: str, user: str = None, owner_is_group: bool = None): + """ + update folder + :param path: path of the folder + :param group: group that owns the folder + :param user: user that owns the folder + :param owner_is_group: if True, the owner is a group + """ + if user is None and owner_is_group is None: + folder = FolderCreateOrUpdate(path=path, group=group) + elif user is None: + folder = FolderCreateOrUpdate(path=path, group=group, owner_is_group=owner_is_group) + elif owner_is_group is None: + folder = FolderCreateOrUpdate(path=path, group=group, user=user) + else: + folder = FolderCreateOrUpdate(path=path, group=group, user=user, owner_is_group=owner_is_group) + return self._send_post_request('folders/create_folder/', folder.model_dump(exclude_unset=True, exclude_defaults=True)) + + def delete_folder(self, path: Path, tenant_id: int = None): + """ + DELETE /folders/delete_folder/ + """ + body = PathBody(path=path).model_dump() + if tenant_id is not None: + body['tenant_id'] = tenant_id + return self._send_delete_request('folders/delete_folder/', body) + + def get_owner_and_group(self, path: Path, tenant_id: int = None): + """ + response: + { + "owning_user": "string", + "owning_uid": "string", + "owning_group": "string", + "owning_gid": "string", + "has_default_acl": true, + "is_directory": true, + "children": 0 + } + """ + body = PathBody(path=path).model_dump() + if tenant_id is not None: + body['tenant_id'] = tenant_id + return self._send_post_request('folders/stat_path/', body) + def update_quota_size(self, quota_id: int, new_size: int): + """ + :param quota_id: int + :param new_size: int in bytes + :return: message indicating success or failure + """ if isinstance(quota_id, int) and isinstance(new_size, int): - return self._send_put_request(f'quotas/{str(quota_id)}/', {'size': new_size}) + body = QuotaUpdate(soft_limit=new_size, hard_limit=new_size) + return self._send_patch_request(f'quotas/{str(quota_id)}/', body.model_dump()) else: raise TypeError('quota_id and size must be of type int') @@ -142,16 +203,11 @@ def delete_quota(self, quota_id: int): else: raise TypeError('quota_id must be of type int') - def is_base10(self): - r = self.get_status() - return r['vms'][0]['capacity_base_10'] - def get_total_capacity(self): """ :return: """ - r = self.get_status() - return r['vms'][0] + return self._send_get_request('capacity/') def get_protected_paths(self, source_dir: Path = None): if source_dir is None: