Merge branch 'main' into run-on-main

Author: rdimitrov

api/openapi.json

@@ -306,6 +306,108 @@
         }
       }
     },
+    "/api/v1/workspaces/archive": {
+      "get": {
+        "tags": [
+          "CodeGate API",
+          "Workspaces"
+        ],
+        "summary": "List Archived Workspaces",
+        "description": "List all archived workspaces.",
+        "operationId": "v1_list_archived_workspaces",
+        "responses": {
+          "200": {
+            "description": "Successful Response",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/ListWorkspacesResponse"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/api/v1/workspaces/archive/{workspace_name}/recover": {
+      "post": {
+        "tags": [
+          "CodeGate API",
+          "Workspaces"
+        ],
+        "summary": "Recover Workspace",
+        "description": "Recover an archived workspace by name.",
+        "operationId": "v1_recover_workspace",
+        "parameters": [
+          {
+            "name": "workspace_name",
+            "in": "path",
+            "required": true,
+            "schema": {
+              "type": "string",
+              "title": "Workspace Name"
+            }
+          }
+        ],
+        "responses": {
+          "204": {
+            "description": "Successful Response"
+          },
+          "422": {
+            "description": "Validation Error",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/HTTPValidationError"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/api/v1/workspaces/archive/{workspace_name}": {
+      "delete": {
+        "tags": [
+          "CodeGate API",
+          "Workspaces"
+        ],
+        "summary": "Hard Delete Workspace",
+        "description": "Hard delete an archived workspace by name.",
+        "operationId": "v1_hard_delete_workspace",
+        "parameters": [
+          {
+            "name": "workspace_name",
+            "in": "path",
+            "required": true,
+            "schema": {
+              "type": "string",
+              "title": "Workspace Name"
+            }
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Successful Response",
+            "content": {
+              "application/json": {
+                "schema": {}
+              }
+            }
+          },
+          "422": {
+            "description": "Validation Error",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/HTTPValidationError"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
     "/api/v1/workspaces/{workspace_name}/alerts": {
       "get": {
         "tags": [

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"
         )