Raise WorkspaceDoesNotExistError on update system-prompt

Author: aponcedeleonch

src/codegate/db/connection.py

@@ -248,7 +248,7 @@ async def record_context(self, context: Optional[PipelineContext]) -> None:
         except Exception as e:
             logger.error(f"Failed to record context: {context}.", error=str(e))
 
-    async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
+    async def add_workspace(self, workspace_name: str) -> Workspace:
         """Add a new workspace to the DB.
 
         This handles validation and insertion of a new workspace.
@@ -274,7 +274,7 @@ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
             raise AlreadyExistsError(f"Workspace {workspace_name} already exists.")
         return added_workspace
 
-    async def update_workspace(self, workspace: Workspace) -> Optional[Workspace]:
+    async def update_workspace(self, workspace: Workspace) -> Workspace:
         sql = text(
             """
             UPDATE workspaces SET
@@ -284,8 +284,7 @@ async def update_workspace(self, workspace: Workspace) -> Optional[Workspace]:
             RETURNING *
             """
         )
-        # We only pass an object to respect the signature of the function
-        updated_workspace = await self._execute_update_pydantic_model(workspace, sql)
+        updated_workspace = await self._execute_update_pydantic_model(workspace, sql, should_raise=True)
         return updated_workspace
 
     async def update_session(self, session: Session) -> Optional[Session]:

src/codegate/pipeline/cli/commands.py

@@ -244,14 +244,15 @@ async def _set_system_prompt(self, flags: Dict[str, str], args: List[str]) -> st
             active_workspace = await self.workspace_crud.get_active_workspace()
             workspace_name = active_workspace.name
 
-        updated_worksapce = await self.workspace_crud.update_workspace_system_prompt(
-            workspace_name, args
-        )
-        if not updated_worksapce:
+        try:
+            updated_worksapce = await self.workspace_crud.update_workspace_system_prompt(
+                workspace_name, args
+            )
+        except crud.WorkspaceDoesNotExistError:
             return (
-                f"Workspace system prompt not updated. "
-                f"Check if the workspace `{workspace_name}` exists"
+                f"Workspace system prompt not updated. Workspace `{workspace_name}` doesn't exist"
             )
+
         return (
             f"Workspace `{updated_worksapce.name}` system prompt "
             f"updated to:\n```\n{updated_worksapce.system_prompt}\n```"

src/codegate/workspaces/crud.py

@@ -85,10 +85,10 @@ async def activate_workspace(self, workspace_name: str):
 
     async def update_workspace_system_prompt(
         self, workspace_name: str, sys_prompt_lst: List[str]
-    ) -> Optional[Workspace]:
+    ) -> Workspace:
         selected_workspace = await self._db_reader.get_workspace_by_name(workspace_name)
         if not selected_workspace:
-            return None
+            raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.")
 
         system_prompt = " ".join(sys_prompt_lst)
         workspace_update = Workspace(