Skip to content

Commit

Permalink
pydantic the new user quota usage APIs
Browse files Browse the repository at this point in the history
  • Loading branch information
jmchilton committed Feb 27, 2023
1 parent dabdb95 commit c4671c7
Show file tree
Hide file tree
Showing 5 changed files with 171 additions and 53 deletions.
77 changes: 77 additions & 0 deletions client/src/schema/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1189,6 +1189,14 @@ export interface paths {
*/
post: operations["set_beacon_api_users__user_id__beacon_post"];
};
"/api/users/{user_id}/usage": {
/** Return the user's quota usage summary broken down by quota source */
get: operations["get_user_usage_api_users__user_id__usage_get"];
};
"/api/users/{user_id}/usage/{label}": {
/** Return the user's quota usage summary for a given quota source label */
get: operations["get_user_usage_for_label_api_users__user_id__usage__label__get"];
};
"/api/version": {
/**
* Return Galaxy version information: major/minor version, optional extra info
Expand Down Expand Up @@ -7386,6 +7394,19 @@ export interface components {
*/
user: components["schemas"]["UserModel"];
};
/** UserQuotaUsage */
UserQuotaUsage: {
/** Quota */
quota?: string;
/** Quota Bytes */
quota_bytes?: number;
/** Quota Percent */
quota_percent?: number;
/** Quota Source Label */
quota_source_label?: string;
/** Total Disk Usage */
total_disk_usage: number;
};
/** ValidationError */
ValidationError: {
/** Location */
Expand Down Expand Up @@ -14149,6 +14170,62 @@ export interface operations {
};
};
};
get_user_usage_api_users__user_id__usage_get: {
/** Return the user's quota usage summary broken down by quota source */
parameters: {
/** @description The user ID that will be used to effectively make this API call. Only admins and designated users can make API calls on behalf of other users. */
header?: {
"run-as"?: string;
};
/** @description The ID of the user to get or __current__. */
path: {
user_id: string;
};
};
responses: {
/** @description Successful Response */
200: {
content: {
"application/json": components["schemas"]["UserQuotaUsage"][];
};
};
/** @description Validation Error */
422: {
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
get_user_usage_for_label_api_users__user_id__usage__label__get: {
/** Return the user's quota usage summary for a given quota source label */
parameters: {
/** @description The user ID that will be used to effectively make this API call. Only admins and designated users can make API calls on behalf of other users. */
header?: {
"run-as"?: string;
};
/** @description The ID of the user to get or __current__. */
/** @description The label corresponding to the quota source to fetch usage information about. */
path: {
user_id: string;
label: string;
};
};
responses: {
/** @description Successful Response */
200: {
content: {
"application/json": components["schemas"]["UserQuotaUsage"];
};
};
/** @description Validation Error */
422: {
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
version_api_version_get: {
/**
* Return Galaxy version information: major/minor version, optional extra info
Expand Down
43 changes: 30 additions & 13 deletions lib/galaxy/managers/users.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
base,
deletable,
)
from galaxy.model import UserQuotaUsage
from galaxy.security.validate_user_input import (
VALID_EMAIL_RE,
validate_email,
Expand Down Expand Up @@ -650,22 +651,38 @@ def add_serializers(self):
}
)

def serialize_disk_usage(self, user: model.User) -> List[Dict[str, Any]]:
rval = user.dictify_usage(self.app.object_store)
for usage in rval:
quota_source_label = usage["quota_source_label"]
usage["quota_percent"] = self.user_manager.quota(user, quota_source_label=quota_source_label)
usage["quota"] = self.user_manager.quota(user, total=True, quota_source_label=quota_source_label)
usage["quota_bytes"] = self.user_manager.quota_bytes(user, quota_source_label=quota_source_label)
def serialize_disk_usage(self, user: model.User) -> List[UserQuotaUsage]:
usages = user.dictify_usage(self.app.object_store)
rval: List[UserQuotaUsage] = []
for usage in usages:
quota_source_label = usage.quota_source_label
quota_percent = self.user_manager.quota(user, quota_source_label=quota_source_label)
quota = self.user_manager.quota(user, total=True, quota_source_label=quota_source_label)
quota_bytes = self.user_manager.quota_bytes(user, quota_source_label=quota_source_label)
rval.append(
UserQuotaUsage(
quota_source_label=quota_source_label,
total_disk_usage=usage.total_disk_usage,
quota_percent=quota_percent,
quota=quota,
quota_bytes=quota_bytes,
)
)
return rval

def serialize_disk_usage_for(self, user: model.User, label: Optional[str]) -> Dict[str, Any]:
def serialize_disk_usage_for(self, user: model.User, label: Optional[str]) -> UserQuotaUsage:
usage = user.dictify_usage_for(label)
quota_source_label = usage["quota_source_label"]
usage["quota_percent"] = self.user_manager.quota(user, quota_source_label=quota_source_label)
usage["quota"] = self.user_manager.quota(user, total=True, quota_source_label=quota_source_label)
usage["quota_bytes"] = self.user_manager.quota_bytes(user, quota_source_label=quota_source_label)
return usage
quota_source_label = usage.quota_source_label
quota_percent = self.user_manager.quota(user, quota_source_label=quota_source_label)
quota = self.user_manager.quota(user, total=True, quota_source_label=quota_source_label)
quota_bytes = self.user_manager.quota_bytes(user, quota_source_label=quota_source_label)
return UserQuotaUsage(
quota_source_label=quota_source_label,
total_disk_usage=usage.total_disk_usage,
quota_percent=quota_percent,
quota=quota,
quota_bytes=quota_bytes,
)


class UserDeserializer(base.ModelDeserializer):
Expand Down
70 changes: 42 additions & 28 deletions lib/galaxy/model/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@

import sqlalchemy
from boltons.iterutils import remap
from pydantic import BaseModel
from social_core.storage import (
AssociationMixin,
CodeMixin,
Expand Down Expand Up @@ -634,6 +635,19 @@ def calculate_user_disk_usage_statements(user_id, quota_source_map, for_sqlite=F
return statements


# move these to galaxy.schema.schema once galaxy-data depends on
# galaxy-schema.
class UserQuotaBasicUsage(BaseModel):
quota_source_label: Optional[str]
total_disk_usage: float


class UserQuotaUsage(UserQuotaBasicUsage):
quota_percent: Optional[float]
quota_bytes: Optional[int]
quota: Optional[str]


class User(Base, Dictifiable, RepresentById):
"""
Data for a Galaxy user or admin and relations to their
Expand Down Expand Up @@ -1025,57 +1039,57 @@ def attempt_create_private_role(self):
session.add(assoc)
session.flush()

def dictify_usage(self, object_store=None) -> List[Dict[str, Any]]:
def dictify_usage(self, object_store=None) -> List[UserQuotaBasicUsage]:
"""Include object_store to include empty/unused usage info."""
used_labels: Set[Union[str, None]] = set()
rval: List[Dict[str, Any]] = [
{
"quota_source_label": None,
"total_disk_usage": float(self.disk_usage or 0),
}
rval: List[UserQuotaBasicUsage] = [
UserQuotaBasicUsage(
quota_source_label=None,
total_disk_usage=float(self.disk_usage or 0),
)
]
used_labels.add(None)
for quota_source_usage in self.quota_source_usages:
label = quota_source_usage.quota_source_label
rval.append(
{
"quota_source_label": label,
"total_disk_usage": float(quota_source_usage.disk_usage),
}
UserQuotaBasicUsage(
quota_source_label=label,
total_disk_usage=float(quota_source_usage.disk_usage),
)
)
used_labels.add(label)

if object_store is not None:
for label in object_store.get_quota_source_map().ids_per_quota_source().keys():
if label not in used_labels:
rval.append(
{
"quota_source_label": label,
"total_disk_usage": 0.0,
}
UserQuotaBasicUsage(
quota_source_label=label,
total_disk_usage=0.0,
)
)

return rval

def dictify_usage_for(self, quota_source_label: Optional[str]) -> Dict[str, Any]:
rval: Dict[str, Any]
def dictify_usage_for(self, quota_source_label: Optional[str]) -> UserQuotaBasicUsage:
rval: UserQuotaBasicUsage
if quota_source_label is None:
rval = {
"quota_source_label": None,
"total_disk_usage": float(self.disk_usage or 0),
}
rval = UserQuotaBasicUsage(
quota_source_label=None,
total_disk_usage=float(self.disk_usage or 0),
)
else:
quota_source_usage = self.quota_source_usage_for(quota_source_label)
if quota_source_usage is None:
rval = {
"quota_source_label": quota_source_label,
"total_disk_usage": 0.0,
}
rval = UserQuotaBasicUsage(
quota_source_label=quota_source_label,
total_disk_usage=0.0,
)
else:
rval = {
"quota_source_label": quota_source_label,
"total_disk_usage": float(quota_source_usage.disk_usage),
}
rval = UserQuotaBasicUsage(
quota_source_label=quota_source_label,
total_disk_usage=float(quota_source_usage.disk_usage),
)

return rval

Expand Down
23 changes: 15 additions & 8 deletions lib/galaxy/quota/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Galaxy Quotas"""
import logging
from typing import Optional

from sqlalchemy.sql import text

Expand All @@ -23,10 +24,10 @@ class QuotaAgent: # metaclass=abc.ABCMeta
"""

# TODO: make abstractmethod after they work better with mypy
def get_quota(self, user, quota_source_label=None):
def get_quota(self, user, quota_source_label=None) -> Optional[int]:
"""Return quota in bytes or None if no quota is set."""

def get_quota_nice_size(self, user, quota_source_label=None):
def get_quota_nice_size(self, user, quota_source_label=None) -> Optional[str]:
"""Return quota as a human-readable string or 'unlimited' if no quota is set."""
quota_bytes = self.get_quota(user, quota_source_label=quota_source_label)
if quota_bytes is not None:
Expand All @@ -36,10 +37,12 @@ def get_quota_nice_size(self, user, quota_source_label=None):
return quota_str

# TODO: make abstractmethod after they work better with mypy
def get_percent(self, trans=None, user=False, history=False, usage=False, quota=False, quota_source_label=None):
def get_percent(
self, trans=None, user=False, history=False, usage=False, quota=False, quota_source_label=None
) -> Optional[int]:
"""Return the percentage of any storage quota applicable to the user/transaction."""

def get_usage(self, trans=None, user=False, history=False, quota_source_label=None):
def get_usage(self, trans=None, user=False, history=False, quota_source_label=None) -> Optional[float]:
if trans:
user = trans.user
history = trans.history
Expand Down Expand Up @@ -73,14 +76,16 @@ class NoQuotaAgent(QuotaAgent):
def __init__(self):
pass

def get_quota(self, user, quota_source_label=None):
def get_quota(self, user, quota_source_label=None) -> Optional[int]:
return None

@property
def default_quota(self):
return None

def get_percent(self, trans=None, user=False, history=False, usage=False, quota=False, quota_source_label=None):
def get_percent(
self, trans=None, user=False, history=False, usage=False, quota=False, quota_source_label=None
) -> Optional[int]:
return None

def is_over_quota(self, app, job, job_destination):
Expand All @@ -94,7 +99,7 @@ def __init__(self, model):
self.model = model
self.sa_session = model.context

def get_quota(self, user, quota_source_label=None):
def get_quota(self, user, quota_source_label=None) -> Optional[int]:
"""
Calculated like so:
Expand Down Expand Up @@ -220,7 +225,9 @@ def set_default_quota(self, default_type, quota):
self.sa_session.add(target_default)
self.sa_session.flush()

def get_percent(self, trans=None, user=False, history=False, usage=False, quota=False, quota_source_label=None):
def get_percent(
self, trans=None, user=False, history=False, usage=False, quota=False, quota_source_label=None
) -> Optional[int]:
"""
Return the percentage of any storage quota applicable to the user/transaction.
"""
Expand Down
11 changes: 7 additions & 4 deletions lib/galaxy/webapps/galaxy/api/users.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,10 @@
import json
import logging
import re
from typing import Any, Dict, List, Optional

from typing import (
List,
Optional,
)

from fastapi import (
Body,
Expand Down Expand Up @@ -35,6 +37,7 @@
from galaxy.model import (
User,
UserAddress,
UserQuotaUsage,
)
from galaxy.schema import APIKeyModel
from galaxy.schema.fields import DecodedDatabaseIdField
Expand Down Expand Up @@ -158,7 +161,7 @@ def usage(
self,
trans: ProvidesUserContext = DependsOnTrans,
user_id: str = FlexibleUserIdPathParam,
) -> List[Dict[str, Any]]:
) -> List[UserQuotaUsage]:
user = get_user_full(trans, user_id, False)
if user:
rval = self.user_serializer.serialize_disk_usage(user)
Expand All @@ -176,7 +179,7 @@ def usage_for(
trans: ProvidesUserContext = DependsOnTrans,
user_id: str = FlexibleUserIdPathParam,
label: str = QuotaSourceLabelPathParam,
) -> Optional[Dict[str, Any]]:
) -> Optional[UserQuotaUsage]:
user = get_user_full(trans, user_id, False)
effective_label: Optional[str] = label
if label == "__null__":
Expand Down

0 comments on commit c4671c7

Please sign in to comment.