Skip to content

Commit

Permalink
Allow cluster defaults to apply during user creation
Browse files Browse the repository at this point in the history
  • Loading branch information
romasku committed Jan 11, 2022
1 parent 58dd813 commit 731743f
Showing 1 changed file with 21 additions and 12 deletions.
33 changes: 21 additions & 12 deletions neuro_admin_client/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -897,6 +897,20 @@ async def get_cluster_user(
raw_user = await resp.json()
return self._parse_cluster_user(cluster_name, raw_user)

def _quota_to_payload(self, quota: Quota) -> dict[str, Any]:
res = {}
if quota.total_running_jobs is not None:
res["total_running_jobs"] = quota.total_running_jobs
return res

def _balance_to_payload(self, balance: Balance) -> dict[str, Any]:
res = {}
if balance.credits is not None:
res["credits"] = str(balance.credits)
if balance.spent_credits is not None:
res["spent_credits"] = str(balance.spent_credits)
return res

@overload
async def create_cluster_user(
self,
Expand Down Expand Up @@ -933,20 +947,15 @@ async def create_cluster_user(
with_user_info: bool = False,
org_name: str | None = None,
) -> ClusterUser | ClusterUserWithInfo:
payload: dict[str, Any] = {
"user_name": user_name,
"role": role.value,
"quota": {},
"balance": {},
}
payload: dict[str, Any] = {"user_name": user_name, "role": role.value}
if org_name:
payload["org_name"] = org_name
if quota.total_running_jobs is not None:
payload["quota"]["total_running_jobs"] = quota.total_running_jobs
if balance.credits is not None:
payload["balance"]["credits"] = str(balance.credits)
if balance.spent_credits is not None:
payload["balance"]["spent_credits"] = str(balance.spent_credits)
quota_payload = self._quota_to_payload(quota)
if quota_payload:
payload["quota"] = quota_payload
balance_payload = self._balance_to_payload(balance)
if quota_payload:
payload["balance"] = balance_payload

async with self._request(
"POST",
Expand Down

0 comments on commit 731743f

Please sign in to comment.