diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index be588381..6d9d2e13 100644 --- a/src/codegate/api/v1.py +++ b/src/codegate/api/v1.py @@ -136,6 +136,25 @@ async def delete_workspace(workspace_name: str): return Response(status_code=204) +@v1.delete( + "/workspaces/del/{workspace_name}", + tags=["Workspaces"], + generate_unique_id_function=uniq_name, +) +async def delete_workspace_permanent(workspace_name: str): + """Delete a workspace by name.""" + try: + _ = await wscrud.hard_delete_workspace(workspace_name) + except crud.WorkspaceDoesNotExistError: + raise HTTPException(status_code=404, detail="Workspace does not exist") + 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 Response(status_code=204) + + @v1.get( "/workspaces/{workspace_name}/alerts", tags=["Workspaces"], diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index 103a8eec..a8d25ef9 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -8,7 +8,8 @@ from alembic import command as alembic_command from alembic.config import Config as AlembicConfig from pydantic import BaseModel -from sqlalchemy import CursorResult, TextClause, text +from sqlalchemy import CursorResult, TextClause, event, text +from sqlalchemy.engine import Engine from sqlalchemy.exc import IntegrityError, OperationalError from sqlalchemy.ext.asyncio import create_async_engine @@ -35,6 +36,21 @@ class AlreadyExistsError(Exception): pass +@event.listens_for(Engine, "connect") +def set_sqlite_pragma(dbapi_connection, connection_record): + """ + Ensures that foreign keys are enabled for the SQLite database at every connection. + + SQLite does not enforce foreign keys by default, so we need to enable them manually. + [SQLAlchemy docs](https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#foreign-key-support) + [SQLite docs](https://www.sqlite.org/foreignkeys.html) + [SO](https://stackoverflow.com/questions/2614984/sqlite-sqlalchemy-how-to-enforce-foreign-keys) + """ + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + + class DbCodeGate: _instance = None @@ -318,6 +334,19 @@ async def soft_delete_workspace(self, workspace: Workspace) -> Optional[Workspac ) return deleted_workspace + async def hard_delete_workspace(self, workspace: Workspace) -> Optional[Workspace]: + sql = text( + """ + DELETE FROM workspaces + WHERE id = :id + RETURNING * + """ + ) + deleted_workspace = await self._execute_update_pydantic_model( + workspace, sql, should_raise=True + ) + return deleted_workspace + class DbReader(DbCodeGate): @@ -431,7 +460,7 @@ async def get_workspaces(self) -> List[WorkspaceActive]: workspaces = await self._execute_select_pydantic_model(WorkspaceActive, sql) return workspaces - async def get_workspace_by_name(self, name: str) -> Optional[Workspace]: + async def get_non_deleted_workspace_by_name(self, name: str) -> Optional[Workspace]: sql = text( """ SELECT @@ -446,6 +475,21 @@ async def get_workspace_by_name(self, name: str) -> Optional[Workspace]: ) return workspaces[0] if workspaces else None + async def get_workspace_by_name(self, name: str) -> Optional[Workspace]: + sql = text( + """ + SELECT + id, name, system_prompt, deleted_at + FROM workspaces + WHERE name = :name + """ + ) + conditions = {"name": name} + workspaces = await self._exec_select_conditions_to_pydantic( + Workspace, sql, conditions, should_raise=True + ) + return workspaces[0] if workspaces else None + async def get_sessions(self) -> List[Session]: sql = text( """ diff --git a/src/codegate/db/models.py b/src/codegate/db/models.py index 6120ea1f..d606e0c6 100644 --- a/src/codegate/db/models.py +++ b/src/codegate/db/models.py @@ -44,6 +44,7 @@ class Workspace(BaseModel): id: str name: str system_prompt: Optional[str] + deleted_at: Optional[datetime.datetime] = None @field_validator("name", mode="plain") @classmethod diff --git a/src/codegate/pipeline/cli/commands.py b/src/codegate/pipeline/cli/commands.py index 08df1eca..010311c4 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, + "delete": self._delete_workspace, "rename": self._rename_workspace, } @@ -267,6 +268,27 @@ async def _remove_workspace(self, flags: Dict[str, str], args: List[str]) -> str return "An error occurred while removing the workspace" return f"Workspace **{workspace_name}** has been removed" + async def _delete_workspace(self, flags: Dict[str, str], args: List[str]) -> str: + """ + Remove a workspace + """ + if args is None or len(args) == 0: + return "Please provide a name. Use `codegate workspace delete workspace_name`" + + workspace_name = args[0] + if not workspace_name: + return "Please provide a name. Use `codegate workspace delete workspace_name`" + + try: + await self.workspace_crud.hard_delete_workspace(workspace_name) + except crud.WorkspaceDoesNotExistError: + return f"Workspace **{workspace_name}** does not exist" + except crud.WorkspaceCrudError as e: + return str(e) + except Exception: + return "An error occurred while deleting the workspace" + return f"Workspace **{workspace_name}** has been deleted permanently" + @property def help(self) -> str: return ( @@ -282,13 +304,17 @@ def help(self) -> str: "- `activate`: Activate a workspace\n\n" " - *args*:\n\n" " - `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" + "- `remove`: Remove a workspace. It can be recovered later.\n\n" + " - *args*:\n\n" + " - `workspace_name`" + "- `delete`: Delete permanently a workspace and its' associated info. The workspace " + "first needs to be `remove`.\n\n" + " - *args*:\n\n" + " - `workspace_name`" ) diff --git a/src/codegate/workspaces/crud.py b/src/codegate/workspaces/crud.py index 1352170f..148f6f55 100644 --- a/src/codegate/workspaces/crud.py +++ b/src/codegate/workspaces/crud.py @@ -89,7 +89,7 @@ async def _is_workspace_active( """ # TODO: All of this should be done within a transaction. - selected_workspace = await self._db_reader.get_workspace_by_name(workspace_name) + selected_workspace = await self._db_reader.get_non_deleted_workspace_by_name(workspace_name) if not selected_workspace: raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.") @@ -118,7 +118,7 @@ async def activate_workspace(self, workspace_name: str): async def update_workspace_system_prompt( self, workspace_name: str, sys_prompt_lst: List[str] ) -> Workspace: - selected_workspace = await self._db_reader.get_workspace_by_name(workspace_name) + selected_workspace = await self._db_reader.get_non_deleted_workspace_by_name(workspace_name) if not selected_workspace: raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.") @@ -132,7 +132,7 @@ async def update_workspace_system_prompt( updated_workspace = await db_recorder.update_workspace(workspace_update) return updated_workspace - async def soft_delete_workspace(self, workspace_name: str): + async def soft_delete_workspace(self, workspace_name: str) -> None: """ Soft delete a workspace """ @@ -141,11 +141,11 @@ async def soft_delete_workspace(self, workspace_name: str): if workspace_name == DEFAULT_WORKSPACE_NAME: raise WorkspaceCrudError("Cannot delete default workspace.") - selected_workspace = await self._db_reader.get_workspace_by_name(workspace_name) + selected_workspace = await self._db_reader.get_non_deleted_workspace_by_name(workspace_name) if not selected_workspace: raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.") - # Check if workspace is active, if it is, make the default workspace active + # Check if workspace is active, if it is, avoid deleting it active_workspace = await self._db_reader.get_active_workspace() if active_workspace and active_workspace.id == selected_workspace.id: raise WorkspaceCrudError("Cannot delete active workspace.") @@ -157,8 +157,30 @@ async def soft_delete_workspace(self, workspace_name: str): raise WorkspaceCrudError(f"Error deleting workspace {workspace_name}") return + async def hard_delete_workspace(self, workspace_name: str) -> None: + """ + Soft delete a workspace + """ + if workspace_name == "": + raise WorkspaceCrudError("Workspace name cannot be empty.") + + selected_workspace = await self._db_reader.get_workspace_by_name(workspace_name) + if not selected_workspace: + raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.") + + # Check if workspace is soft deleted, if it is not, don't delete it + if not selected_workspace.deleted_at: + raise WorkspaceCrudError("Cannot delete workspace that is not soft-deleted.") + + db_recorder = DbRecorder() + try: + _ = await db_recorder.hard_delete_workspace(selected_workspace) + except Exception: + raise WorkspaceCrudError(f"Error deleting workspace {workspace_name}") + return + async def get_workspace_by_name(self, workspace_name: str) -> Workspace: - workspace = await self._db_reader.get_workspace_by_name(workspace_name) + workspace = await self._db_reader.get_non_deleted_workspace_by_name(workspace_name) if not workspace: raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.") return workspace