wallaroo.workspace

  1import json
  2from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
  3
  4import requests
  5
  6from wallaroo.wallaroo_ml_ops_api_client.models.workspaces_create_response_200 import (
  7    WorkspacesCreateResponse200,
  8)
  9
 10from . import queries
 11from .models import Models, ModelsList
 12from .object import *
 13from .pipeline import Pipeline
 14from .user import User
 15from .user_type import UserType
 16from .version import _user_agent
 17from .wallaroo_ml_ops_api_client.api.workspace import workspaces_create
 18from .wallaroo_ml_ops_api_client.models import workspaces_create_json_body
 19
 20if TYPE_CHECKING:
 21    # Imports that happen below in methods to fix circular import dependency
 22    # issues need to also be specified here to satisfy mypy type checking.
 23    from .client import Client
 24
 25
 26class Workspace(Object):
 27    """Workspace provides a user and visibility context for access to models and pipelines."""
 28
 29    def __init__(
 30        self, client: Optional["Client"], data: Dict[str, Any], standalone=False
 31    ) -> None:
 32        assert client is not None
 33        self.client = client
 34        #        data = {"name": workspace_name, "owner": owner_user_id}
 35        super().__init__(gql_client=client._gql_client, data=data)
 36
 37    def __repr__(self) -> str:
 38        return str(self.to_json())
 39
 40    def _fetch_attributes(self) -> Dict[str, Any]:
 41        """Fetches all member data from the GraphQL API."""
 42        data = self._gql_client.execute(
 43            gql.gql(queries.named("WorkspaceById")),
 44            variable_values={"id": self.id()},
 45        )["workspace_by_pk"]
 46
 47        return data
 48
 49    def _fill(self, data: Dict[str, Any]) -> None:
 50        """Fills an object given a response dictionary from the GraphQL API.
 51
 52        Only the primary key member must be present; other members will be
 53        filled in via rehydration if their corresponding member function is
 54        called.
 55        """
 56
 57        for required_attribute in ["id"]:
 58            if required_attribute not in data:
 59                raise RequiredAttributeMissing(
 60                    self.__class__.__name__, required_attribute
 61                )
 62        # Required
 63        self._id = data["id"]
 64
 65        # Optional
 66        self._name = value_if_present(data, "name")
 67        self._archived = value_if_present(data, "archived")
 68        self._created_at = value_if_present(data, "created_at")
 69        self._created_by = value_if_present(data, "created_by")
 70        self._models = value_if_present(data, "models")
 71        self._pipelines = value_if_present(data, "pipelines")
 72        self._users = value_if_present(data, "users")
 73
 74    def to_json(self):
 75        return {
 76            "name": self.name(),
 77            "id": self.id(),
 78            "archived": self.archived(),
 79            "created_by": self.created_by(),
 80            "created_at": self.created_at(),
 81            "models": self.models(),
 82            "pipelines": self.pipelines(),
 83        }
 84
 85    @staticmethod
 86    def _get_user_default_workspace(client):
 87
 88        res = client._gql_client.execute(
 89            gql.gql(queries.named("UserDefaultWorkspace")),
 90            variable_values={"user_id": client.auth.user_id()},
 91        )
 92        if res["user_default_workspace"] == []:
 93            return None
 94        else:
 95            return Workspace(client, res["user_default_workspace"][0]["workspace"])
 96
 97    @staticmethod
 98    def _create_user_default_workspace(client) -> "Workspace":
 99
100        res = client._gql_client.execute(
101            gql.gql(queries.named("CreateDefaultUserWorkspace")),
102            variable_values={
103                "user_id": client.auth.user_id(),
104                "workspace_name": f"{client.auth.user_email()} - Default Workspace",
105                "user_type": UserType.OWNER.name,
106            },
107        )
108        return Workspace(client, res["insert_workspace_one"])
109
110    @staticmethod
111    def _create_workspace(client: "Client", name: str):
112        body = workspaces_create_json_body.WorkspacesCreateJsonBody(name)
113        res = workspaces_create.sync(client=client.mlops(), json_body=body)
114
115        if not isinstance(res, WorkspacesCreateResponse200):
116            raise Exception("Failed to create workspace.")
117
118        if res is None:
119            raise Exception("Failed to create workspace.")
120
121        return Workspace(client, {"id": res.workspace_id})
122
123    def id(self) -> int:
124        return self._id
125
126    @rehydrate("_name")
127    def name(self) -> str:
128        return cast(str, self._name)
129
130    @rehydrate("_archived")
131    def archived(self) -> bool:
132        return cast(bool, self._archived)
133
134    @rehydrate("_created_at")
135    def created_at(self) -> str:
136        return cast(str, self._created_at)
137
138    @rehydrate("_created_by")
139    def created_by(self) -> str:
140        return cast(str, self._created_by)
141
142    @rehydrate("_models")
143    def models(self) -> ModelsList:
144        """
145        Returns a List of Models objects that have been created in this workspace.
146        """
147        self._fill(self._fetch_attributes())
148        return cast(
149            ModelsList,
150            list(
151                map(
152                    lambda m: Models(self.client, data=m),
153                    cast(List[Dict], self._models),
154                )
155            ),
156        )
157
158    @rehydrate("_pipelines")
159    def pipelines(self) -> List[Pipeline]:
160        self._fill(self._fetch_attributes())
161        return cast(
162            List[Pipeline],
163            list(
164                map(
165                    lambda m: Pipeline(self.client, data=m),
166                    cast(List[Dict], self._pipelines),
167                )
168            ),
169        )
170
171    @rehydrate("_users")
172    def users(self) -> List[User]:
173        return cast(
174            List[User],
175            list(
176                map(
177                    lambda u: self.client._get_user_by_id(u["user_id"]),
178                    cast(List[Dict], self._users),
179                )
180            ),
181        )
182
183    def add_user(self, user_email: str) -> "Workspace":
184        """Add a user to workspace as participant"""
185        user = self.client.get_user_by_email(user_email)
186        if user is None:
187            raise EntityNotFoundError("User", {"email": user_email})
188        return self._add_user(user, UserType.COLLABORATOR)
189
190    def add_owner(self, user_email: str) -> "Workspace":
191        """Add a user to workspace as owner"""
192
193        user = self.client.get_user_by_email(user_email)
194        if user is None:
195            raise EntityNotFoundError("User", {"email": user_email})
196
197        return self._add_user(user, UserType.OWNER)
198
199    def _add_user(self, user: User, user_type: UserType) -> "Workspace":
200        headers = {
201            "authorization": self.client.auth._bearer_token_str(),
202            "user-agent": _user_agent,
203        }
204        invite = requests.post(
205            f"{self.client.api_endpoint}/v1/api/workspaces/add_user",
206            data=json.dumps(
207                {
208                    "email": user.email(),
209                    "workspace_id": self._id,
210                    "user_id": user.id(),
211                    "url": "http://localhost:3000",
212                    "user_type": user_type,
213                }
214            ),
215            headers=headers,
216        )
217        if invite.status_code > 299:
218            raise Exception(f"Failed to invite: {user.id()} to workspace: {self._id}")
219        self._fill(data={"id": self._id})
220        return self
221
222    def remove_user(self, user_email: str):
223
224        user = self.client.get_user_by_email(user_email)
225        assert user is not None
226        if user.id() == self.client.auth.user_id():
227            default_workspace = self._get_user_default_workspace(self.client)
228            if default_workspace.id() == self.id():
229                raise Exception("User cannot be removed from their default workspace")
230
231        if len(self.users()) <= 1:
232            raise Exception("Last user of a workspace cannot be removed")
233
234        res = self.client._gql_client.execute(
235            gql.gql(queries.named("DeleteWorkspaceUser")),
236            variable_values={"user_id": user.id(), "workspace_id": self.id()},
237        )
238        if user.id() == self.client.auth.user_id():
239            self._fill({"id": default_workspace.id()})
240        else:
241            self._rehydrate()
242
243
244class Workspaces(List[Workspace]):
245    """Wraps a list of workspaces for display."""
246
247    @staticmethod
248    def _format_users(users: List[User]) -> List[str]:
249        return [u.email() for u in users]
250
251    def _repr_html_(self) -> str:
252        rows = [
253            f"""
254        <tr >
255            <td>{r.name()}</td>
256            <td>{r.created_at()[:19].replace("T", " ")}</td>
257            <td>{Workspaces._format_users(r.users())}</td>
258            <td>{len(r.models())}</td>
259            <td>{len(r.pipelines())}</td>
260        </tr>
261        """
262            for r in self
263        ]
264        table = """
265        <table>
266            <tr>
267                <th>Name</th>
268                <th>Created At</th>
269                <th>Users</th>
270                <th>Models</th>
271                <th>Pipelines</th>
272            </tr>
273            {0}
274        </table>
275        """.format(
276            "\n".join(rows)
277        )
278        return table
class Workspace(wallaroo.object.Object):
 27class Workspace(Object):
 28    """Workspace provides a user and visibility context for access to models and pipelines."""
 29
 30    def __init__(
 31        self, client: Optional["Client"], data: Dict[str, Any], standalone=False
 32    ) -> None:
 33        assert client is not None
 34        self.client = client
 35        #        data = {"name": workspace_name, "owner": owner_user_id}
 36        super().__init__(gql_client=client._gql_client, data=data)
 37
 38    def __repr__(self) -> str:
 39        return str(self.to_json())
 40
 41    def _fetch_attributes(self) -> Dict[str, Any]:
 42        """Fetches all member data from the GraphQL API."""
 43        data = self._gql_client.execute(
 44            gql.gql(queries.named("WorkspaceById")),
 45            variable_values={"id": self.id()},
 46        )["workspace_by_pk"]
 47
 48        return data
 49
 50    def _fill(self, data: Dict[str, Any]) -> None:
 51        """Fills an object given a response dictionary from the GraphQL API.
 52
 53        Only the primary key member must be present; other members will be
 54        filled in via rehydration if their corresponding member function is
 55        called.
 56        """
 57
 58        for required_attribute in ["id"]:
 59            if required_attribute not in data:
 60                raise RequiredAttributeMissing(
 61                    self.__class__.__name__, required_attribute
 62                )
 63        # Required
 64        self._id = data["id"]
 65
 66        # Optional
 67        self._name = value_if_present(data, "name")
 68        self._archived = value_if_present(data, "archived")
 69        self._created_at = value_if_present(data, "created_at")
 70        self._created_by = value_if_present(data, "created_by")
 71        self._models = value_if_present(data, "models")
 72        self._pipelines = value_if_present(data, "pipelines")
 73        self._users = value_if_present(data, "users")
 74
 75    def to_json(self):
 76        return {
 77            "name": self.name(),
 78            "id": self.id(),
 79            "archived": self.archived(),
 80            "created_by": self.created_by(),
 81            "created_at": self.created_at(),
 82            "models": self.models(),
 83            "pipelines": self.pipelines(),
 84        }
 85
 86    @staticmethod
 87    def _get_user_default_workspace(client):
 88
 89        res = client._gql_client.execute(
 90            gql.gql(queries.named("UserDefaultWorkspace")),
 91            variable_values={"user_id": client.auth.user_id()},
 92        )
 93        if res["user_default_workspace"] == []:
 94            return None
 95        else:
 96            return Workspace(client, res["user_default_workspace"][0]["workspace"])
 97
 98    @staticmethod
 99    def _create_user_default_workspace(client) -> "Workspace":
100
101        res = client._gql_client.execute(
102            gql.gql(queries.named("CreateDefaultUserWorkspace")),
103            variable_values={
104                "user_id": client.auth.user_id(),
105                "workspace_name": f"{client.auth.user_email()} - Default Workspace",
106                "user_type": UserType.OWNER.name,
107            },
108        )
109        return Workspace(client, res["insert_workspace_one"])
110
111    @staticmethod
112    def _create_workspace(client: "Client", name: str):
113        body = workspaces_create_json_body.WorkspacesCreateJsonBody(name)
114        res = workspaces_create.sync(client=client.mlops(), json_body=body)
115
116        if not isinstance(res, WorkspacesCreateResponse200):
117            raise Exception("Failed to create workspace.")
118
119        if res is None:
120            raise Exception("Failed to create workspace.")
121
122        return Workspace(client, {"id": res.workspace_id})
123
124    def id(self) -> int:
125        return self._id
126
127    @rehydrate("_name")
128    def name(self) -> str:
129        return cast(str, self._name)
130
131    @rehydrate("_archived")
132    def archived(self) -> bool:
133        return cast(bool, self._archived)
134
135    @rehydrate("_created_at")
136    def created_at(self) -> str:
137        return cast(str, self._created_at)
138
139    @rehydrate("_created_by")
140    def created_by(self) -> str:
141        return cast(str, self._created_by)
142
143    @rehydrate("_models")
144    def models(self) -> ModelsList:
145        """
146        Returns a List of Models objects that have been created in this workspace.
147        """
148        self._fill(self._fetch_attributes())
149        return cast(
150            ModelsList,
151            list(
152                map(
153                    lambda m: Models(self.client, data=m),
154                    cast(List[Dict], self._models),
155                )
156            ),
157        )
158
159    @rehydrate("_pipelines")
160    def pipelines(self) -> List[Pipeline]:
161        self._fill(self._fetch_attributes())
162        return cast(
163            List[Pipeline],
164            list(
165                map(
166                    lambda m: Pipeline(self.client, data=m),
167                    cast(List[Dict], self._pipelines),
168                )
169            ),
170        )
171
172    @rehydrate("_users")
173    def users(self) -> List[User]:
174        return cast(
175            List[User],
176            list(
177                map(
178                    lambda u: self.client._get_user_by_id(u["user_id"]),
179                    cast(List[Dict], self._users),
180                )
181            ),
182        )
183
184    def add_user(self, user_email: str) -> "Workspace":
185        """Add a user to workspace as participant"""
186        user = self.client.get_user_by_email(user_email)
187        if user is None:
188            raise EntityNotFoundError("User", {"email": user_email})
189        return self._add_user(user, UserType.COLLABORATOR)
190
191    def add_owner(self, user_email: str) -> "Workspace":
192        """Add a user to workspace as owner"""
193
194        user = self.client.get_user_by_email(user_email)
195        if user is None:
196            raise EntityNotFoundError("User", {"email": user_email})
197
198        return self._add_user(user, UserType.OWNER)
199
200    def _add_user(self, user: User, user_type: UserType) -> "Workspace":
201        headers = {
202            "authorization": self.client.auth._bearer_token_str(),
203            "user-agent": _user_agent,
204        }
205        invite = requests.post(
206            f"{self.client.api_endpoint}/v1/api/workspaces/add_user",
207            data=json.dumps(
208                {
209                    "email": user.email(),
210                    "workspace_id": self._id,
211                    "user_id": user.id(),
212                    "url": "http://localhost:3000",
213                    "user_type": user_type,
214                }
215            ),
216            headers=headers,
217        )
218        if invite.status_code > 299:
219            raise Exception(f"Failed to invite: {user.id()} to workspace: {self._id}")
220        self._fill(data={"id": self._id})
221        return self
222
223    def remove_user(self, user_email: str):
224
225        user = self.client.get_user_by_email(user_email)
226        assert user is not None
227        if user.id() == self.client.auth.user_id():
228            default_workspace = self._get_user_default_workspace(self.client)
229            if default_workspace.id() == self.id():
230                raise Exception("User cannot be removed from their default workspace")
231
232        if len(self.users()) <= 1:
233            raise Exception("Last user of a workspace cannot be removed")
234
235        res = self.client._gql_client.execute(
236            gql.gql(queries.named("DeleteWorkspaceUser")),
237            variable_values={"user_id": user.id(), "workspace_id": self.id()},
238        )
239        if user.id() == self.client.auth.user_id():
240            self._fill({"id": default_workspace.id()})
241        else:
242            self._rehydrate()

Workspace provides a user and visibility context for access to models and pipelines.

Workspace( client: Optional[wallaroo.client.Client], data: Dict[str, Any], standalone=False)
30    def __init__(
31        self, client: Optional["Client"], data: Dict[str, Any], standalone=False
32    ) -> None:
33        assert client is not None
34        self.client = client
35        #        data = {"name": workspace_name, "owner": owner_user_id}
36        super().__init__(gql_client=client._gql_client, data=data)

Base constructor.

Each object requires:

  • a GraphQL client - in order to fill its missing members dynamically
  • an initial data blob - typically from unserialized JSON, contains at
  • least the data for required members (typically the object's primary key) and optionally other data members.
def to_json(self):
75    def to_json(self):
76        return {
77            "name": self.name(),
78            "id": self.id(),
79            "archived": self.archived(),
80            "created_by": self.created_by(),
81            "created_at": self.created_at(),
82            "models": self.models(),
83            "pipelines": self.pipelines(),
84        }
def id(self) -> int:
124    def id(self) -> int:
125        return self._id
def name(*args, **kwargs):
41        def wrapper(*args, **kwargs):
42            obj = args[0]
43            if not getattr(obj, "_standalone", None):
44                present = getattr(obj, attr) != DehydratedValue()
45                # Uncomment to debug while testing
46                # print(
47                #    "rehydrate: {} -> {}".format(
48                #        attr, "present" if present else "not present"
49                #    )
50                # )
51                if not present:
52                    obj._rehydrate()
53            result = fn(*args, **kwargs)
54            return result
def archived(*args, **kwargs):
41        def wrapper(*args, **kwargs):
42            obj = args[0]
43            if not getattr(obj, "_standalone", None):
44                present = getattr(obj, attr) != DehydratedValue()
45                # Uncomment to debug while testing
46                # print(
47                #    "rehydrate: {} -> {}".format(
48                #        attr, "present" if present else "not present"
49                #    )
50                # )
51                if not present:
52                    obj._rehydrate()
53            result = fn(*args, **kwargs)
54            return result
def created_at(*args, **kwargs):
41        def wrapper(*args, **kwargs):
42            obj = args[0]
43            if not getattr(obj, "_standalone", None):
44                present = getattr(obj, attr) != DehydratedValue()
45                # Uncomment to debug while testing
46                # print(
47                #    "rehydrate: {} -> {}".format(
48                #        attr, "present" if present else "not present"
49                #    )
50                # )
51                if not present:
52                    obj._rehydrate()
53            result = fn(*args, **kwargs)
54            return result
def created_by(*args, **kwargs):
41        def wrapper(*args, **kwargs):
42            obj = args[0]
43            if not getattr(obj, "_standalone", None):
44                present = getattr(obj, attr) != DehydratedValue()
45                # Uncomment to debug while testing
46                # print(
47                #    "rehydrate: {} -> {}".format(
48                #        attr, "present" if present else "not present"
49                #    )
50                # )
51                if not present:
52                    obj._rehydrate()
53            result = fn(*args, **kwargs)
54            return result
def models(*args, **kwargs):
41        def wrapper(*args, **kwargs):
42            obj = args[0]
43            if not getattr(obj, "_standalone", None):
44                present = getattr(obj, attr) != DehydratedValue()
45                # Uncomment to debug while testing
46                # print(
47                #    "rehydrate: {} -> {}".format(
48                #        attr, "present" if present else "not present"
49                #    )
50                # )
51                if not present:
52                    obj._rehydrate()
53            result = fn(*args, **kwargs)
54            return result

Returns a List of Models objects that have been created in this workspace.

def pipelines(*args, **kwargs):
41        def wrapper(*args, **kwargs):
42            obj = args[0]
43            if not getattr(obj, "_standalone", None):
44                present = getattr(obj, attr) != DehydratedValue()
45                # Uncomment to debug while testing
46                # print(
47                #    "rehydrate: {} -> {}".format(
48                #        attr, "present" if present else "not present"
49                #    )
50                # )
51                if not present:
52                    obj._rehydrate()
53            result = fn(*args, **kwargs)
54            return result
def users(*args, **kwargs):
41        def wrapper(*args, **kwargs):
42            obj = args[0]
43            if not getattr(obj, "_standalone", None):
44                present = getattr(obj, attr) != DehydratedValue()
45                # Uncomment to debug while testing
46                # print(
47                #    "rehydrate: {} -> {}".format(
48                #        attr, "present" if present else "not present"
49                #    )
50                # )
51                if not present:
52                    obj._rehydrate()
53            result = fn(*args, **kwargs)
54            return result
def add_user(self, user_email: str) -> wallaroo.workspace.Workspace:
184    def add_user(self, user_email: str) -> "Workspace":
185        """Add a user to workspace as participant"""
186        user = self.client.get_user_by_email(user_email)
187        if user is None:
188            raise EntityNotFoundError("User", {"email": user_email})
189        return self._add_user(user, UserType.COLLABORATOR)

Add a user to workspace as participant

def add_owner(self, user_email: str) -> wallaroo.workspace.Workspace:
191    def add_owner(self, user_email: str) -> "Workspace":
192        """Add a user to workspace as owner"""
193
194        user = self.client.get_user_by_email(user_email)
195        if user is None:
196            raise EntityNotFoundError("User", {"email": user_email})
197
198        return self._add_user(user, UserType.OWNER)

Add a user to workspace as owner

def remove_user(self, user_email: str):
223    def remove_user(self, user_email: str):
224
225        user = self.client.get_user_by_email(user_email)
226        assert user is not None
227        if user.id() == self.client.auth.user_id():
228            default_workspace = self._get_user_default_workspace(self.client)
229            if default_workspace.id() == self.id():
230                raise Exception("User cannot be removed from their default workspace")
231
232        if len(self.users()) <= 1:
233            raise Exception("Last user of a workspace cannot be removed")
234
235        res = self.client._gql_client.execute(
236            gql.gql(queries.named("DeleteWorkspaceUser")),
237            variable_values={"user_id": user.id(), "workspace_id": self.id()},
238        )
239        if user.id() == self.client.auth.user_id():
240            self._fill({"id": default_workspace.id()})
241        else:
242            self._rehydrate()
class Workspaces(typing.List[wallaroo.workspace.Workspace]):
245class Workspaces(List[Workspace]):
246    """Wraps a list of workspaces for display."""
247
248    @staticmethod
249    def _format_users(users: List[User]) -> List[str]:
250        return [u.email() for u in users]
251
252    def _repr_html_(self) -> str:
253        rows = [
254            f"""
255        <tr >
256            <td>{r.name()}</td>
257            <td>{r.created_at()[:19].replace("T", " ")}</td>
258            <td>{Workspaces._format_users(r.users())}</td>
259            <td>{len(r.models())}</td>
260            <td>{len(r.pipelines())}</td>
261        </tr>
262        """
263            for r in self
264        ]
265        table = """
266        <table>
267            <tr>
268                <th>Name</th>
269                <th>Created At</th>
270                <th>Users</th>
271                <th>Models</th>
272                <th>Pipelines</th>
273            </tr>
274            {0}
275        </table>
276        """.format(
277            "\n".join(rows)
278        )
279        return table

Wraps a list of workspaces for display.

Inherited Members
builtins.list
list
clear
copy
append
insert
extend
pop
remove
index
count
reverse
sort