Implement CRUD actions for dealing with archived workspaces (#686)

* Implement CRUD actions for dealing with archived workspaces

This implements listing, recovering and hard-deleting archived
workspaces.

Closes: #668

Signed-off-by: Juan Antonio Osorio <[email protected]>

* Add pragma to enable foreign key support in SQLite

Co-Authored-By: Alejandro Ponce de Leon <[email protected]>
Signed-off-by: Juan Antonio Osorio <[email protected]>

---------

Signed-off-by: Juan Antonio Osorio <[email protected]>
Co-authored-by: Alejandro Ponce de Leon <[email protected]>

Author: JAORMX

src/codegate/api/v1.py

@@ -27,7 +27,7 @@ async def list_workspaces() -> v1_models.ListWorkspacesResponse:
     """List all workspaces."""
     wslist = await wscrud.get_workspaces()
 
-    resp = v1_models.ListWorkspacesResponse.from_db_workspaces(wslist)
+    resp = v1_models.ListWorkspacesResponse.from_db_workspaces_active(wslist)
 
     return resp
 
@@ -136,6 +136,55 @@ async def delete_workspace(workspace_name: str):
     return Response(status_code=204)
 
 
+@v1.get("/workspaces/archive", tags=["Workspaces"], generate_unique_id_function=uniq_name)
+async def list_archived_workspaces() -> v1_models.ListWorkspacesResponse:
+    """List all archived workspaces."""
+    wslist = await wscrud.get_archived_workspaces()
+
+    resp = v1_models.ListWorkspacesResponse.from_db_workspaces(wslist)
+
+    return resp
+
+
+@v1.post(
+    "/workspaces/archive/{workspace_name}/recover",
+    tags=["Workspaces"],
+    generate_unique_id_function=uniq_name,
+    status_code=204,
+)
+async def recover_workspace(workspace_name: str):
+    """Recover an archived workspace by name."""
+    try:
+        _ = await wscrud.recover_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.delete(
+    "/workspaces/archive/{workspace_name}",
+    tags=["Workspaces"],
+    generate_unique_id_function=uniq_name,
+)
+async def hard_delete_workspace(workspace_name: str):
+    """Hard delete an archived 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"],

src/codegate/api/v1_models.py

@@ -23,7 +23,7 @@ class ListWorkspacesResponse(pydantic.BaseModel):
     workspaces: list[Workspace]
 
     @classmethod
-    def from_db_workspaces(
+    def from_db_workspaces_active(
         cls, db_workspaces: List[db_models.WorkspaceActive]
     ) -> "ListWorkspacesResponse":
         return cls(
@@ -33,6 +33,12 @@ def from_db_workspaces(
             ]
         )
 
+    @classmethod
+    def from_db_workspaces(
+        cls, db_workspaces: List[db_models.Workspace]
+    ) -> "ListWorkspacesResponse":
+        return cls(workspaces=[Workspace(name=ws.name, is_active=False) for ws in db_workspaces])
+
 
 class ListActiveWorkspacesResponse(pydantic.BaseModel):
     workspaces: list[ActiveWorkspace]

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,20 @@ 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 +333,33 @@ 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
+
+    async def recover_workspace(self, workspace: Workspace) -> Optional[Workspace]:
+        sql = text(
+            """
+            UPDATE workspaces
+            SET deleted_at = NULL
+            WHERE id = :id
+            RETURNING *
+            """
+        )
+        recovered_workspace = await self._execute_update_pydantic_model(
+            workspace, sql, should_raise=True
+        )
+        return recovered_workspace
+
 
 class DbReader(DbCodeGate):
 
@@ -431,6 +473,19 @@ async def get_workspaces(self) -> List[WorkspaceActive]:
         workspaces = await self._execute_select_pydantic_model(WorkspaceActive, sql)
         return workspaces
 
+    async def get_archived_workspaces(self) -> List[Workspace]:
+        sql = text(
+            """
+            SELECT
+                id, name, system_prompt
+            FROM workspaces
+            WHERE deleted_at IS NOT NULL
+            ORDER BY deleted_at DESC
+            """
+        )
+        workspaces = await self._execute_select_pydantic_model(Workspace, sql)
+        return workspaces
+
     async def get_workspace_by_name(self, name: str) -> Optional[Workspace]:
         sql = text(
             """
@@ -446,6 +501,21 @@ async def get_workspace_by_name(self, name: str) -> Optional[Workspace]:
         )
         return workspaces[0] if workspaces else None
 
+    async def get_archived_workspace_by_name(self, name: str) -> Optional[Workspace]:
+        sql = text(
+            """
+            SELECT
+                id, name, system_prompt
+            FROM workspaces
+            WHERE name = :name AND deleted_at IS NOT NULL
+            """
+        )
+        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(
             """

src/codegate/pipeline/cli/commands.py

@@ -155,6 +155,9 @@ def subcommands(self) -> Dict[str, Callable[[List[str]], Awaitable[str]]]:
             "activate": self._activate_workspace,
             "remove": self._remove_workspace,
             "rename": self._rename_workspace,
+            "list-archived": self._list_archived_workspaces,
+            "restore": self._restore_workspace,
+            "delete-archived": self._delete_archived_workspace,
         }
 
     async def _list_workspaces(self, flags: Dict[str, str], args: List[str]) -> str:
@@ -267,6 +270,58 @@ 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 _list_archived_workspaces(self, flags: Dict[str, str], args: List[str]) -> str:
+        """
+        List all archived workspaces
+        """
+        workspaces = await self.workspace_crud.get_archived_workspaces()
+        respond_str = ""
+        for workspace in workspaces:
+            respond_str += f"- {workspace.name}\n"
+        return respond_str
+
+    async def _restore_workspace(self, flags: Dict[str, str], args: List[str]) -> str:
+        """
+        Restore an archived workspace
+        """
+        if args is None or len(args) == 0:
+            return "Please provide a name. Use `codegate workspace restore workspace_name`"
+
+        workspace_name = args[0]
+        if not workspace_name:
+            return "Please provide a name. Use `codegate workspace restore workspace_name`"
+
+        try:
+            await self.workspace_crud.recover_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 restoring the workspace"
+        return f"Workspace **{workspace_name}** has been restored"
+
+    async def _delete_archived_workspace(self, flags: Dict[str, str], args: List[str]) -> str:
+        """
+        Hard delete an archived workspace
+        """
+        if args is None or len(args) == 0:
+            return "Please provide a name. Use `codegate workspace delete-archived workspace_name`"
+
+        workspace_name = args[0]
+        if not workspace_name:
+            return "Please provide a name. Use `codegate workspace delete-archived 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"
+
     @property
     def help(self) -> str:
         return (
@@ -289,6 +344,14 @@ def help(self) -> str:
             "  - *args*:\n\n"
             "    - `workspace_name`\n"
             "    - `new_workspace_name`\n\n"
+            "- `list-archived`: List all archived workspaces\n\n"
+            "  - *args*: None\n\n"
+            "- `restore`: Restore an archived workspace\n\n"
+            "  - *args*:\n\n"
+            "    - `workspace_name`\n\n"
+            "- `delete-archived`: Hard delete an archived workspace\n\n"
+            "  - *args*:\n\n"
+            "    - `workspace_name`\n\n"
         )
 
 

src/codegate/workspaces/crud.py

@@ -20,7 +20,7 @@ class WorkspaceAlreadyActiveError(WorkspaceCrudError):
 DEFAULT_WORKSPACE_NAME = "default"
 
 # These are reserved keywords that cannot be used for workspaces
-RESERVED_WORKSPACE_KEYWORDS = [DEFAULT_WORKSPACE_NAME, "active"]
+RESERVED_WORKSPACE_KEYWORDS = [DEFAULT_WORKSPACE_NAME, "active", "archived"]
 
 
 class WorkspaceCrud:
@@ -75,6 +75,12 @@ async def get_workspaces(self) -> List[WorkspaceActive]:
         """
         return await self._db_reader.get_workspaces()
 
+    async def get_archived_workspaces(self) -> List[Workspace]:
+        """
+        Get all archived workspaces
+        """
+        return await self._db_reader.get_archived_workspaces()
+
     async def get_active_workspace(self) -> Optional[ActiveWorkspace]:
         """
         Get the active workspace
@@ -115,6 +121,18 @@ async def activate_workspace(self, workspace_name: str):
         await db_recorder.update_session(session)
         return
 
+    async def recover_workspace(self, workspace_name: str):
+        """
+        Recover an archived workspace
+        """
+        selected_workspace = await self._db_reader.get_archived_workspace_by_name(workspace_name)
+        if not selected_workspace:
+            raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.")
+
+        db_recorder = DbRecorder()
+        await db_recorder.recover_workspace(selected_workspace)
+        return
+
     async def update_workspace_system_prompt(
         self, workspace_name: str, sys_prompt_lst: List[str]
     ) -> Workspace:
@@ -157,6 +175,24 @@ 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):
+        """
+        Hard delete a workspace
+        """
+        if workspace_name == "":
+            raise WorkspaceCrudError("Workspace name cannot be empty.")
+
+        selected_workspace = await self._db_reader.get_archived_workspace_by_name(workspace_name)
+        if not selected_workspace:
+            raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.")
+
+        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)
         if not workspace: