From c09d5b86a28fde89396ced83d3c13b287ed8881c Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Tue, 21 Jan 2025 13:26:00 +0200 Subject: [PATCH] Add constructs to rename workspaces This adds the CRUD API endpoints and pseudo-CLI subcommands to rename workspaces. Note that `default` and `active` are reserved keywords and cannot be used. Closes: https://github.com/stacklok/codegate/issues/671 Signed-off-by: Juan Antonio Osorio --- src/codegate/api/v1.py | 36 ++++++++++++++++++++++- src/codegate/api/v1_models.py | 7 ++++- src/codegate/pipeline/cli/commands.py | 41 ++++++++++++++++++++++++++- src/codegate/workspaces/crud.py | 26 +++++++++++++++++ 4 files changed, 107 insertions(+), 3 deletions(-) diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index 885727a8..be588381 100644 --- a/src/codegate/api/v1.py +++ b/src/codegate/api/v1.py @@ -61,8 +61,18 @@ async def activate_workspace(request: v1_models.ActivateWorkspaceRequest, status @v1.post("/workspaces", tags=["Workspaces"], generate_unique_id_function=uniq_name, status_code=201) -async def create_workspace(request: v1_models.CreateWorkspaceRequest) -> v1_models.Workspace: +async def create_workspace( + request: v1_models.CreateOrRenameWorkspaceRequest, +) -> v1_models.Workspace: """Create a new workspace.""" + if request.rename_to is not None: + return await rename_workspace(request) + return await create_new_workspace(request) + + +async def create_new_workspace( + request: v1_models.CreateOrRenameWorkspaceRequest, +) -> v1_models.Workspace: # Input validation is done in the model try: _ = await wscrud.add_workspace(request.name) @@ -83,6 +93,30 @@ async def create_workspace(request: v1_models.CreateWorkspaceRequest) -> v1_mode return v1_models.Workspace(name=request.name, is_active=False) +async def rename_workspace( + request: v1_models.CreateOrRenameWorkspaceRequest, +) -> v1_models.Workspace: + try: + _ = await wscrud.rename_workspace(request.name, request.rename_to) + except crud.WorkspaceDoesNotExistError: + raise HTTPException(status_code=404, detail="Workspace does not exist") + except AlreadyExistsError: + raise HTTPException(status_code=409, detail="Workspace already exists") + except ValidationError: + raise HTTPException( + status_code=400, + detail=( + "Invalid workspace name. " "Please use only alphanumeric characters and dashes" + ), + ) + except crud.WorkspaceCrudError as e: + raise HTTPException(status_code=400, detail=str(e)) + except Exception: + raise HTTPException(status_code=500, detail="Internal server error") + + return v1_models.Workspace(name=request.rename_to, is_active=False) + + @v1.delete( "/workspaces/{workspace_name}", tags=["Workspaces"], diff --git a/src/codegate/api/v1_models.py b/src/codegate/api/v1_models.py index 6416b15c..2435c1e1 100644 --- a/src/codegate/api/v1_models.py +++ b/src/codegate/api/v1_models.py @@ -48,9 +48,14 @@ def from_db_workspaces( ) -class CreateWorkspaceRequest(pydantic.BaseModel): +class CreateOrRenameWorkspaceRequest(pydantic.BaseModel): name: str + # If set, rename the workspace to this name. Note that + # the 'name' field is still required and the workspace + # workspace must exist. + rename_to: Optional[str] = None + class ActivateWorkspaceRequest(pydantic.BaseModel): name: str diff --git a/src/codegate/pipeline/cli/commands.py b/src/codegate/pipeline/cli/commands.py index 2a97b3b1..08df1eca 100644 --- a/src/codegate/pipeline/cli/commands.py +++ b/src/codegate/pipeline/cli/commands.py @@ -154,6 +154,7 @@ def subcommands(self) -> Dict[str, Callable[[List[str]], Awaitable[str]]]: "add": self._add_workspace, "activate": self._activate_workspace, "remove": self._remove_workspace, + "rename": self._rename_workspace, } async def _list_workspaces(self, flags: Dict[str, str], args: List[str]) -> str: @@ -193,6 +194,37 @@ async def _add_workspace(self, flags: Dict[str, str], args: List[str]) -> str: return f"Workspace **{new_workspace_name}** has been added" + async def _rename_workspace(self, flags: Dict[str, str], args: List[str]) -> str: + """ + Rename a workspace + """ + if args is None or len(args) < 2: + return ( + "Please provide a name and a new name. " + "Use `codegate workspace rename workspace_name new_workspace_name`" + ) + + old_workspace_name = args[0] + new_workspace_name = args[1] + if not old_workspace_name or not new_workspace_name: + return ( + "Please provide a name and a new name. " + "Use `codegate workspace rename workspace_name new_workspace_name`" + ) + + try: + await self.workspace_crud.rename_workspace(old_workspace_name, new_workspace_name) + except crud.WorkspaceDoesNotExistError: + return f"Workspace **{old_workspace_name}** does not exist" + except AlreadyExistsError: + return f"Workspace **{new_workspace_name}** already exists" + except crud.WorkspaceCrudError: + return "An error occurred while renaming the workspace" + except Exception: + return "An error occurred while renaming the workspace" + + return f"Workspace **{old_workspace_name}** has been renamed to **{new_workspace_name}**" + async def _activate_workspace(self, flags: Dict[str, str], args: List[str]) -> str: """ Activate a workspace @@ -249,7 +281,14 @@ def help(self) -> str: " - `workspace_name`\n\n" "- `activate`: Activate a workspace\n\n" " - *args*:\n\n" - " - `workspace_name`" + " - `workspace_name`\n\n" + "- `remove`: Remove a workspace\n\n" + " - *args*:\n\n" + " - `workspace_name`\n\n" + "- `rename`: Rename a workspace\n\n" + " - *args*:\n\n" + " - `workspace_name`\n" + " - `new_workspace_name`\n\n" ) diff --git a/src/codegate/workspaces/crud.py b/src/codegate/workspaces/crud.py index 0af7b950..1352170f 100644 --- a/src/codegate/workspaces/crud.py +++ b/src/codegate/workspaces/crud.py @@ -43,6 +43,32 @@ async def add_workspace(self, new_workspace_name: str) -> Workspace: workspace_created = await db_recorder.add_workspace(new_workspace_name) return workspace_created + async def rename_workspace(self, old_workspace_name: str, new_workspace_name: str) -> Workspace: + """ + Rename a workspace + + Args: + old_name (str): The old name of the workspace + new_name (str): The new name of the workspace + """ + if new_workspace_name == "": + raise WorkspaceCrudError("Workspace name cannot be empty.") + if old_workspace_name == "": + raise WorkspaceCrudError("Workspace name cannot be empty.") + if old_workspace_name in DEFAULT_WORKSPACE_NAME: + raise WorkspaceCrudError("Cannot rename default workspace.") + if new_workspace_name in RESERVED_WORKSPACE_KEYWORDS: + raise WorkspaceCrudError(f"Workspace name {new_workspace_name} is reserved.") + if old_workspace_name == new_workspace_name: + raise WorkspaceCrudError("Old and new workspace names are the same.") + ws = await self._db_reader.get_workspace_by_name(old_workspace_name) + if not ws: + raise WorkspaceDoesNotExistError(f"Workspace {old_workspace_name} does not exist.") + db_recorder = DbRecorder() + new_ws = Workspace(id=ws.id, name=new_workspace_name, system_prompt=ws.system_prompt) + workspace_renamed = await db_recorder.update_workspace(new_ws) + return workspace_renamed + async def get_workspaces(self) -> List[WorkspaceActive]: """ Get all workspaces