1  from mlflow.exceptions import MlflowException
  2  from mlflow.server.auth.permissions import get_permission
  3  from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME, resolve_entity_workspace_name
  4  
  5  
  6  class User:
  7      def __init__(
  8          self,
  9          id_,
 10          username,
 11          password_hash,
 12          is_admin,
 13          experiment_permissions=None,
 14          registered_model_permissions=None,
 15          scorer_permissions=None,
 16      ):
 17          self._id = id_
 18          self._username = username
 19          self._password_hash = password_hash
 20          self._is_admin = is_admin
 21          self._experiment_permissions = experiment_permissions
 22          self._registered_model_permissions = registered_model_permissions
 23          self._scorer_permissions = scorer_permissions
 24  
 25      @property
 26      def id(self):
 27          return self._id
 28  
 29      @property
 30      def username(self):
 31          return self._username
 32  
 33      @property
 34      def password_hash(self):
 35          return self._password_hash
 36  
 37      @property
 38      def is_admin(self):
 39          return self._is_admin
 40  
 41      @is_admin.setter
 42      def is_admin(self, is_admin):
 43          self._is_admin = is_admin
 44  
 45      @property
 46      def experiment_permissions(self):
 47          return self._experiment_permissions
 48  
 49      @experiment_permissions.setter
 50      def experiment_permissions(self, experiment_permissions):
 51          self._experiment_permissions = experiment_permissions
 52  
 53      @property
 54      def registered_model_permissions(self):
 55          return self._registered_model_permissions
 56  
 57      @registered_model_permissions.setter
 58      def registered_model_permissions(self, registered_model_permissions):
 59          self._registered_model_permissions = registered_model_permissions
 60  
 61      @property
 62      def scorer_permissions(self):
 63          return self._scorer_permissions
 64  
 65      @scorer_permissions.setter
 66      def scorer_permissions(self, scorer_permissions):
 67          self._scorer_permissions = scorer_permissions
 68  
 69      def to_json(self):
 70          return {
 71              "id": self.id,
 72              "username": self.username,
 73              "is_admin": self.is_admin,
 74              "experiment_permissions": [p.to_json() for p in self.experiment_permissions],
 75              "registered_model_permissions": [
 76                  p.to_json() for p in self.registered_model_permissions
 77              ],
 78              "scorer_permissions": [p.to_json() for p in self.scorer_permissions],
 79          }
 80  
 81      @classmethod
 82      def from_json(cls, dictionary):
 83          return cls(
 84              id_=dictionary["id"],
 85              username=dictionary["username"],
 86              password_hash="REDACTED",
 87              is_admin=dictionary["is_admin"],
 88              experiment_permissions=[
 89                  ExperimentPermission.from_json(p) for p in dictionary["experiment_permissions"]
 90              ],
 91              registered_model_permissions=[
 92                  RegisteredModelPermission.from_json(p)
 93                  for p in dictionary["registered_model_permissions"]
 94              ],
 95              scorer_permissions=[
 96                  ScorerPermission.from_json(p) for p in dictionary["scorer_permissions"]
 97              ],
 98          )
 99  
100  
101  class ExperimentPermission:
102      def __init__(
103          self,
104          experiment_id,
105          user_id,
106          permission,
107      ):
108          self._experiment_id = experiment_id
109          self._user_id = user_id
110          self._permission = permission
111  
112      @property
113      def experiment_id(self):
114          return self._experiment_id
115  
116      @property
117      def user_id(self):
118          return self._user_id
119  
120      @property
121      def permission(self):
122          return self._permission
123  
124      @permission.setter
125      def permission(self, permission):
126          self._permission = permission
127  
128      def to_json(self):
129          return {
130              "experiment_id": self.experiment_id,
131              "user_id": self.user_id,
132              "permission": self.permission,
133          }
134  
135      @classmethod
136      def from_json(cls, dictionary):
137          return cls(
138              experiment_id=dictionary["experiment_id"],
139              user_id=dictionary["user_id"],
140              permission=dictionary["permission"],
141          )
142  
143  
144  class RegisteredModelPermission:
145      def __init__(
146          self,
147          name,
148          user_id,
149          permission,
150          workspace=None,
151      ):
152          self._workspace = resolve_entity_workspace_name(workspace)
153          self._name = name
154          self._user_id = user_id
155          self._permission = permission
156  
157      @property
158      def workspace(self):
159          return self._workspace
160  
161      @property
162      def name(self):
163          return self._name
164  
165      @property
166      def user_id(self):
167          return self._user_id
168  
169      @property
170      def permission(self):
171          return self._permission
172  
173      @permission.setter
174      def permission(self, permission):
175          self._permission = permission
176  
177      def to_json(self):
178          return {
179              "workspace": self.workspace,
180              "name": self.name,
181              "user_id": self.user_id,
182              "permission": self.permission,
183          }
184  
185      @classmethod
186      def from_json(cls, dictionary):
187          return cls(
188              name=dictionary["name"],
189              user_id=dictionary["user_id"],
190              permission=dictionary["permission"],
191              workspace=dictionary.get("workspace"),
192          )
193  
194  
195  class ScorerPermission:
196      def __init__(
197          self,
198          experiment_id,
199          scorer_name,
200          user_id,
201          permission,
202      ):
203          self._experiment_id = experiment_id
204          self._scorer_name = scorer_name
205          self._user_id = user_id
206          self._permission = permission
207  
208      @property
209      def experiment_id(self):
210          return self._experiment_id
211  
212      @property
213      def scorer_name(self):
214          return self._scorer_name
215  
216      @property
217      def user_id(self):
218          return self._user_id
219  
220      @property
221      def permission(self):
222          return self._permission
223  
224      @permission.setter
225      def permission(self, permission):
226          self._permission = permission
227  
228      def to_json(self):
229          return {
230              "experiment_id": self.experiment_id,
231              "scorer_name": self.scorer_name,
232              "user_id": self.user_id,
233              "permission": self.permission,
234          }
235  
236      @classmethod
237      def from_json(cls, dictionary):
238          return cls(
239              experiment_id=dictionary["experiment_id"],
240              scorer_name=dictionary["scorer_name"],
241              user_id=dictionary["user_id"],
242              permission=dictionary["permission"],
243          )
244  
245  
246  class GatewaySecretPermission:
247      def __init__(self, secret_id, user_id, permission):
248          self._secret_id = secret_id
249          self._user_id = user_id
250          self._permission = permission
251  
252      @property
253      def secret_id(self):
254          return self._secret_id
255  
256      @property
257      def user_id(self):
258          return self._user_id
259  
260      @property
261      def permission(self):
262          return self._permission
263  
264      @permission.setter
265      def permission(self, permission):
266          self._permission = permission
267  
268      def to_json(self):
269          return {
270              "secret_id": self.secret_id,
271              "user_id": self.user_id,
272              "permission": self.permission,
273          }
274  
275      @classmethod
276      def from_json(cls, dictionary):
277          return cls(
278              secret_id=dictionary["secret_id"],
279              user_id=dictionary["user_id"],
280              permission=dictionary["permission"],
281          )
282  
283  
284  class GatewayEndpointPermission:
285      def __init__(self, endpoint_id, user_id, permission):
286          self._endpoint_id = endpoint_id
287          self._user_id = user_id
288          self._permission = permission
289  
290      @property
291      def endpoint_id(self):
292          return self._endpoint_id
293  
294      @property
295      def user_id(self):
296          return self._user_id
297  
298      @property
299      def permission(self):
300          return self._permission
301  
302      @permission.setter
303      def permission(self, permission):
304          self._permission = permission
305  
306      def to_json(self):
307          return {
308              "endpoint_id": self.endpoint_id,
309              "user_id": self.user_id,
310              "permission": self.permission,
311          }
312  
313      @classmethod
314      def from_json(cls, dictionary):
315          return cls(
316              endpoint_id=dictionary["endpoint_id"],
317              user_id=dictionary["user_id"],
318              permission=dictionary["permission"],
319          )
320  
321  
322  class GatewayModelDefinitionPermission:
323      def __init__(self, model_definition_id, user_id, permission):
324          self._model_definition_id = model_definition_id
325          self._user_id = user_id
326          self._permission = permission
327  
328      @property
329      def model_definition_id(self):
330          return self._model_definition_id
331  
332      @property
333      def user_id(self):
334          return self._user_id
335  
336      @property
337      def permission(self):
338          return self._permission
339  
340      @permission.setter
341      def permission(self, permission):
342          self._permission = permission
343  
344      def to_json(self):
345          return {
346              "model_definition_id": self.model_definition_id,
347              "user_id": self.user_id,
348              "permission": self.permission,
349          }
350  
351      @classmethod
352      def from_json(cls, dictionary):
353          return cls(
354              model_definition_id=dictionary["model_definition_id"],
355              user_id=dictionary["user_id"],
356              permission=dictionary["permission"],
357          )
358  
359  
360  class Role:
361      def __init__(
362          self,
363          id_,
364          name,
365          workspace,
366          description=None,
367          permissions=None,
368      ):
369          self._id = id_
370          self._name = name
371          self._workspace = workspace
372          self._description = description
373          self._permissions = permissions or []
374  
375      @property
376      def id(self):
377          return self._id
378  
379      @property
380      def name(self):
381          return self._name
382  
383      @name.setter
384      def name(self, name):
385          self._name = name
386  
387      @property
388      def workspace(self):
389          return self._workspace
390  
391      @property
392      def description(self):
393          return self._description
394  
395      @description.setter
396      def description(self, description):
397          self._description = description
398  
399      @property
400      def permissions(self):
401          return self._permissions
402  
403      def to_json(self):
404          return {
405              "id": self.id,
406              "name": self.name,
407              "workspace": self.workspace,
408              "description": self.description,
409              "permissions": [p.to_json() for p in self.permissions],
410          }
411  
412      @classmethod
413      def from_json(cls, dictionary):
414          return cls(
415              id_=dictionary["id"],
416              name=dictionary["name"],
417              workspace=dictionary.get("workspace", DEFAULT_WORKSPACE_NAME),
418              description=dictionary.get("description"),
419              permissions=[RolePermission.from_json(p) for p in dictionary.get("permissions", [])],
420          )
421  
422  
423  class RolePermission:
424      def __init__(self, id_, role_id, resource_type, resource_pattern, permission):
425          self._id = id_
426          self._role_id = role_id
427          self._resource_type = resource_type
428          self._resource_pattern = resource_pattern
429          self._permission = permission
430  
431      @property
432      def id(self):
433          return self._id
434  
435      @property
436      def role_id(self):
437          return self._role_id
438  
439      @property
440      def resource_type(self):
441          return self._resource_type
442  
443      @property
444      def resource_pattern(self):
445          return self._resource_pattern
446  
447      @property
448      def permission(self):
449          return self._permission
450  
451      @permission.setter
452      def permission(self, permission):
453          self._permission = permission
454  
455      def to_json(self):
456          return {
457              "id": self.id,
458              "role_id": self.role_id,
459              "resource_type": self.resource_type,
460              "resource_pattern": self.resource_pattern,
461              "permission": self.permission,
462          }
463  
464      @classmethod
465      def from_json(cls, dictionary):
466          return cls(
467              id_=dictionary["id"],
468              role_id=dictionary["role_id"],
469              resource_type=dictionary["resource_type"],
470              resource_pattern=dictionary["resource_pattern"],
471              permission=dictionary["permission"],
472          )
473  
474  
475  class UserRoleAssignment:
476      def __init__(self, id_, user_id, role_id):
477          self._id = id_
478          self._user_id = user_id
479          self._role_id = role_id
480  
481      @property
482      def id(self):
483          return self._id
484  
485      @property
486      def user_id(self):
487          return self._user_id
488  
489      @property
490      def role_id(self):
491          return self._role_id
492  
493      def to_json(self):
494          return {
495              "id": self.id,
496              "user_id": self.user_id,
497              "role_id": self.role_id,
498          }
499  
500      @classmethod
501      def from_json(cls, dictionary):
502          return cls(
503              id_=dictionary["id"],
504              user_id=dictionary["user_id"],
505              role_id=dictionary["role_id"],
506          )
507  
508  
509  class WorkspacePermission:
510      def __init__(self, workspace, user_id, permission):
511          if workspace is None or user_id is None or permission is None:
512              raise MlflowException.invalid_parameter_value(
513                  "workspace, user_id, and permission are required."
514              )
515          self._workspace = workspace
516          self._user_id = user_id
517          self._permission = permission
518  
519      @property
520      def workspace(self):
521          return self._workspace
522  
523      @property
524      def user_id(self):
525          return self._user_id
526  
527      @property
528      def permission(self):
529          return self._permission
530  
531      @property
532      def can_use(self):
533          return get_permission(self.permission).can_use
534  
535      def to_json(self):
536          return {
537              "workspace": self.workspace,
538              "user_id": self.user_id,
539              "permission": self.permission,
540              "can_use": self.can_use,
541          }
542  
543      @classmethod
544      def from_json(cls, dictionary):
545          required_fields = ["workspace", "user_id", "permission"]
546          if missing := [field for field in required_fields if field not in dictionary]:
547              raise MlflowException.invalid_parameter_value(
548                  f"Missing required fields: {', '.join(missing)}"
549              )
550          return cls(
551              workspace=dictionary["workspace"],
552              user_id=dictionary["user_id"],
553              permission=dictionary["permission"],
554          )