File size: 3,796 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 |
"""Public API: users."""
import requests
from wandb_gql import gql
import wandb
from wandb.apis.attrs import Attrs
class User(Attrs):
CREATE_USER_MUTATION = gql(
"""
mutation CreateUserFromAdmin($email: String!, $admin: Boolean) {
createUser(input: {email: $email, admin: $admin}) {
user {
id
name
username
email
admin
}
}
}
"""
)
DELETE_API_KEY_MUTATION = gql(
"""
mutation DeleteApiKey($id: String!) {
deleteApiKey(input: {id: $id}) {
success
}
}
"""
)
GENERATE_API_KEY_MUTATION = gql(
"""
mutation GenerateApiKey($description: String) {
generateApiKey(input: {description: $description}) {
apiKey {
id
name
}
}
}
"""
)
def __init__(self, client, attrs):
super().__init__(attrs)
self._client = client
self._user_api = None
@property
def user_api(self):
"""An instance of the api using credentials from the user."""
if self._user_api is None and len(self.api_keys) > 0:
self._user_api = wandb.Api(api_key=self.api_keys[0])
return self._user_api
@classmethod
def create(cls, api, email, admin=False):
"""Create a new user.
Args:
api: (`Api`) The api instance to use
email: (str) The name of the team
admin: (bool) Whether this user should be a global instance admin
Returns:
A `User` object
"""
res = api.client.execute(
cls.CREATE_USER_MUTATION,
{"email": email, "admin": admin},
)
return User(api.client, res["createUser"]["user"])
@property
def api_keys(self):
if self._attrs.get("apiKeys") is None:
return []
return [k["node"]["name"] for k in self._attrs["apiKeys"]["edges"]]
@property
def teams(self):
if self._attrs.get("teams") is None:
return []
return [k["node"]["name"] for k in self._attrs["teams"]["edges"]]
def delete_api_key(self, api_key):
"""Delete a user's api key.
Returns:
Boolean indicating success
Raises:
ValueError if the api_key couldn't be found
"""
idx = self.api_keys.index(api_key)
try:
self._client.execute(
self.DELETE_API_KEY_MUTATION,
{"id": self._attrs["apiKeys"]["edges"][idx]["node"]["id"]},
)
except requests.exceptions.HTTPError:
return False
return True
def generate_api_key(self, description=None):
"""Generate a new api key.
Returns:
The new api key, or None on failure
"""
try:
# We must make this call using credentials from the original user
key = self.user_api.client.execute(
self.GENERATE_API_KEY_MUTATION, {"description": description}
)["generateApiKey"]["apiKey"]
self._attrs["apiKeys"]["edges"].append({"node": key})
return key["name"]
except (requests.exceptions.HTTPError, AttributeError):
return None
def __repr__(self):
if "email" in self._attrs:
return f"<User {self._attrs['email']}>"
elif "username" in self._attrs:
return f"<User {self._attrs['username']}>"
elif "id" in self._attrs:
return f"<User {self._attrs['id']}>"
elif "name" in self._attrs:
return f"<User {self._attrs['name']!r}>"
else:
return "<User ???>"
|