File size: 5,474 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
"""Public API: teams."""
import requests
from wandb_gql import gql
from wandb.apis.attrs import Attrs
class Member(Attrs):
DELETE_MEMBER_MUTATION = gql(
"""
mutation DeleteInvite($id: String, $entityName: String) {
deleteInvite(input: {id: $id, entityName: $entityName}) {
success
}
}
"""
)
def __init__(self, client, team, attrs):
super().__init__(attrs)
self._client = client
self.team = team
def delete(self):
"""Remove a member from a team.
Returns:
Boolean indicating success
"""
try:
return self._client.execute(
self.DELETE_MEMBER_MUTATION, {"id": self.id, "entityName": self.team}
)["deleteInvite"]["success"]
except requests.exceptions.HTTPError:
return False
def __repr__(self):
return f"<Member {self.name} ({self.account_type})>"
class Team(Attrs):
CREATE_TEAM_MUTATION = gql(
"""
mutation CreateTeam($teamName: String!, $teamAdminUserName: String) {
createTeam(input: {teamName: $teamName, teamAdminUserName: $teamAdminUserName}) {
entity {
id
name
available
photoUrl
limits
}
}
}
"""
)
CREATE_INVITE_MUTATION = gql(
"""
mutation CreateInvite($entityName: String!, $email: String, $username: String, $admin: Boolean) {
createInvite(input: {entityName: $entityName, email: $email, username: $username, admin: $admin}) {
invite {
id
name
email
createdAt
toUser {
name
}
}
}
}
"""
)
TEAM_QUERY = gql(
"""
query Entity($name: String!) {
entity(name: $name) {
id
name
available
photoUrl
readOnly
readOnlyAdmin
isTeam
privateOnly
storageBytes
codeSavingEnabled
defaultAccess
isPaid
members {
id
admin
pending
email
username
name
photoUrl
accountType
apiKey
}
}
}
"""
)
CREATE_SERVICE_ACCOUNT_MUTATION = gql(
"""
mutation CreateServiceAccount($entityName: String!, $description: String!) {
createServiceAccount(
input: {description: $description, entityName: $entityName}
) {
user {
id
}
}
}
"""
)
def __init__(self, client, name, attrs=None):
super().__init__(attrs or {})
self._client = client
self.name = name
self.load()
@classmethod
def create(cls, api, team, admin_username=None):
"""Create a new team.
Args:
api: (`Api`) The api instance to use
team: (str) The name of the team
admin_username: (str) optional username of the admin user of the team, defaults to the current user.
Returns:
A `Team` object
"""
try:
api.client.execute(
cls.CREATE_TEAM_MUTATION,
{"teamName": team, "teamAdminUserName": admin_username},
)
except requests.exceptions.HTTPError:
pass
return Team(api.client, team)
def invite(self, username_or_email, admin=False):
"""Invite a user to a team.
Args:
username_or_email: (str) The username or email address of the user you want to invite
admin: (bool) Whether to make this user a team admin, defaults to False
Returns:
True on success, False if user was already invited or didn't exist
"""
variables = {"entityName": self.name, "admin": admin}
if "@" in username_or_email:
variables["email"] = username_or_email
else:
variables["username"] = username_or_email
try:
self._client.execute(self.CREATE_INVITE_MUTATION, variables)
except requests.exceptions.HTTPError:
return False
return True
def create_service_account(self, description):
"""Create a service account for the team.
Args:
description: (str) A description for this service account
Returns:
The service account `Member` object, or None on failure
"""
try:
self._client.execute(
self.CREATE_SERVICE_ACCOUNT_MUTATION,
{"description": description, "entityName": self.name},
)
self.load(True)
return self.members[-1]
except requests.exceptions.HTTPError:
return None
def load(self, force=False):
if force or not self._attrs:
response = self._client.execute(self.TEAM_QUERY, {"name": self.name})
self._attrs = response["entity"]
self._attrs["members"] = [
Member(self._client, self.name, member)
for member in self._attrs["members"]
]
return self._attrs
def __repr__(self):
return f"<Team {self.name}>"
|