Initial CRUD API for workspaces

This adds a simple and unimplemented REST API for workspaces. Workspaces
will be the base for all other resources in terms of REST resource
mapping, so these go first.

These are initially left entirely unimplemented as
#600 needs to merge

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

Author: JAORMX

src/codegate/api/__init__.py

@@ -0,0 +1,77 @@
+from fastapi import APIRouter
+from fastapi.exceptions import HTTPException
+from fastapi.routing import APIRoute
+from pydantic import ValidationError
+
+from codegate.api import v1_models
+from codegate.db.connection import WorkspaceAlreadyExistsError
+from codegate.pipeline.workspace import commands as wscmd
+
+v1 = APIRouter()
+wscrud = wscmd.WorkspaceCrud()
+
+
+def uniq_name(route: APIRoute):
+    return f"v1_{route.name}"
+
+
+@v1.get("/workspaces", tags=["Workspaces"], generate_unique_id_function=uniq_name)
+async def list_workspaces() -> v1_models.ListWorkspacesResponse:
+    """List all workspaces."""
+    wslist = await wscrud.get_workspaces()
+
+    resp = v1_models.ListWorkspacesResponse.from_db_workspaces(wslist)
+
+    return resp
+
+
+@v1.get("/workspaces/active", tags=["Workspaces"], generate_unique_id_function=uniq_name)
+async def list_active_workspaces() -> v1_models.ListActiveWorkspacesResponse:
+    """List all active workspaces.
+
+    In it's current form, this function will only return one workspace. That is,
+    the globally active workspace."""
+    activews = await wscrud.get_active_workspace()
+
+    resp = v1_models.ListActiveWorkspacesResponse.from_db_workspaces(activews)
+
+    return resp
+
+
+@v1.post("/workspaces/active", tags=["Workspaces"], generate_unique_id_function=uniq_name)
+async def activate_workspace(request: v1_models.ActivateWorkspaceRequest, status_code=204):
+    """Activate a workspace by name."""
+    activated = await wscrud.activate_workspace(request.name)
+
+    # TODO: Refactor
+    if not activated:
+        return HTTPException(status_code=409, detail="Workspace already active")
+
+    return None
+
+
+@v1.post("/workspaces", tags=["Workspaces"], generate_unique_id_function=uniq_name, status_code=201)
+async def create_workspace(request: v1_models.CreateWorkspaceRequest):
+    """Create a new workspace."""
+    # Input validation is done in the model
+    try:
+        created = await wscrud.add_workspace(request.name, should_raise=True)
+    except ValidationError as e:
+        raise HTTPException(status_code=400,
+                            detail=str(e))
+    except WorkspaceAlreadyExistsError:
+        raise HTTPException(status_code=409,
+                            detail="Workspace already exists")
+
+    if not created:
+        raise HTTPException(status_code=500, detail="Failed to create workspace")
+
+    return v1_models.Workspace(name=request.name)
+
+
+
+@v1.delete("/workspaces/{workspace_name}", tags=["Workspaces"],
+           generate_unique_id_function=uniq_name, status_code=204)
+async def delete_workspace(workspace_name: str):
+    """Delete a workspace by name."""
+    raise NotImplementedError

src/codegate/api/v1.py

@@ -0,0 +1,44 @@
+from typing import Any, List, Optional
+
+import pydantic
+
+from codegate.db import models as db_models
+
+
+class Workspace(pydantic.BaseModel):
+    name: str
+    is_active: bool
+
+class ActiveWorkspace(Workspace):
+    # TODO: use a more specific type for last_updated
+    last_updated: Any
+
+class ListWorkspacesResponse(pydantic.BaseModel):
+    workspaces: list[Workspace]
+
+    @classmethod
+    def from_db_workspaces(
+        cls, db_workspaces: List[db_models.WorkspaceActive])-> "ListWorkspacesResponse":
+        return cls(workspaces=[
+            Workspace(name=ws.name, is_active=ws.active_workspace_id is not None)
+            for ws in db_workspaces])
+
+class ListActiveWorkspacesResponse(pydantic.BaseModel):
+    workspaces: list[ActiveWorkspace]
+
+    @classmethod
+    def from_db_workspaces(
+        cls, ws: Optional[db_models.ActiveWorkspace]) -> "ListActiveWorkspacesResponse":
+        if ws is None:
+            return cls(workspaces=[])
+        return cls(workspaces=[
+            ActiveWorkspace(name=ws.name,
+                            is_active=True,
+                            last_updated=ws.last_update)
+                            ])
+
+class CreateWorkspaceRequest(pydantic.BaseModel):
+    name: str
+
+class ActivateWorkspaceRequest(pydantic.BaseModel):
+    name: str

src/codegate/api/v1_models.py

@@ -30,6 +30,9 @@
 alert_queue = asyncio.Queue()
 fim_cache = FimCache()
 
+class WorkspaceAlreadyExistsError(Exception):
+    pass
+
 
 class DbCodeGate:
     _instance = None
@@ -70,7 +73,7 @@ def __init__(self, sqlite_path: Optional[str] = None):
         super().__init__(sqlite_path)
 
     async def _execute_update_pydantic_model(
-        self, model: BaseModel, sql_command: TextClause
+        self, model: BaseModel, sql_command: TextClause, should_raise: bool = False
     ) -> Optional[BaseModel]:
         """Execute an update or insert command for a Pydantic model."""
         async with self._async_db_engine.begin() as conn:
@@ -85,6 +88,8 @@ async def _execute_update_pydantic_model(
                 return model_class(**row._asdict())
             except Exception as e:
                 logger.error(f"Failed to update model: {model}.", error=str(e))
+                if should_raise:
+                    raise e
                 return None
 
     async def record_request(self, prompt_params: Optional[Prompt] = None) -> Optional[Prompt]:
@@ -242,11 +247,14 @@ 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, should_raise: bool = False) -> Optional[Workspace]:
         try:
             workspace = Workspace(id=str(uuid.uuid4()), name=workspace_name)
         except ValidationError as e:
             logger.error(f"Failed to create workspace with name: {workspace_name}: {str(e)}")
+            if should_raise:
+                raise e
             return None
 
         sql = text(
@@ -256,7 +264,17 @@ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
             RETURNING *
             """
         )
-        added_workspace = await self._execute_update_pydantic_model(workspace, sql)
+        try:
+            added_workspace = await self._execute_update_pydantic_model(
+                workspace, sql, should_raise=should_raise)
+        except Exception as e:
+            logger.error(f"Failed to add workspace: {workspace_name}.", error=str(e))
+            if should_raise:
+                # Overwrite the exception; we know that the workspace already exists
+                raise WorkspaceAlreadyExistsError(
+                    f"Workspace with name {workspace_name} already exists.")
+            return None
+
         return added_workspace
 
     async def update_session(self, session: Session) -> Optional[Session]:

src/codegate/db/connection.py

@@ -1,32 +1,39 @@
 import datetime
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 from codegate.db.connection import DbReader, DbRecorder
-from codegate.db.models import Session, Workspace
+from codegate.db.models import ActiveWorkspace, Session, Workspace, WorkspaceActive
 
 
 class WorkspaceCrud:
 
     def __init__(self):
         self._db_reader = DbReader()
 
-    async def add_workspace(self, new_workspace_name: str) -> bool:
+    async def add_workspace(self, new_workspace_name: str, should_raise: bool = True) -> bool:
         """
         Add a workspace
 
         Args:
             name (str): The name of the workspace
         """
         db_recorder = DbRecorder()
-        workspace_created = await db_recorder.add_workspace(new_workspace_name)
+        workspace_created = await db_recorder.add_workspace(
+            new_workspace_name, should_raise=should_raise)
         return bool(workspace_created)
 
-    async def get_workspaces(self):
+    async def get_workspaces(self) -> List[WorkspaceActive]:
         """
         Get all workspaces
         """
         return await self._db_reader.get_workspaces()
 
+    async def get_active_workspace(self) -> Optional[ActiveWorkspace]:
+        """
+        Get the active workspace
+        """
+        return await self._db_reader.get_active_workspace()
+
     async def _is_workspace_active_or_not_exist(
         self, workspace_name: str
     ) -> Tuple[bool, Optional[Session], Optional[Workspace]]:

src/codegate/pipeline/workspace/commands.py

@@ -7,6 +7,7 @@
 from starlette.middleware.errors import ServerErrorMiddleware
 
 from codegate import __description__, __version__
+from codegate.api.v1 import v1
 from codegate.dashboard.dashboard import dashboard_router
 from codegate.pipeline.factory import PipelineFactory
 from codegate.providers.anthropic.provider import AnthropicProvider
@@ -97,4 +98,7 @@ async def health_check():
     app.include_router(system_router)
     app.include_router(dashboard_router)
 
+    # CodeGate API
+    app.include_router(v1, prefix="/api/v1", tags=["CodeGate API"])
+
     return app

src/codegate/server.py

@@ -1,10 +1,9 @@
-import datetime
 from unittest.mock import AsyncMock, patch
 
 import pytest
 
-from codegate.db.models import Session, Workspace, WorkspaceActive
-from codegate.pipeline.workspace.commands import WorkspaceCommands, WorkspaceCrud
+from codegate.db.models import WorkspaceActive
+from codegate.pipeline.workspace.commands import WorkspaceCommands
 
 
 @pytest.mark.asyncio