From 7e4b963ce117b0d2b8c1da80a6502c717b195220 Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Fri, 17 Jan 2025 11:29:24 +0200 Subject: [PATCH] Initial CRUD API for workspaces (#620) 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 https://github.com/stacklok/codegate/pull/600 needs to merge Signed-off-by: Juan Antonio Osorio --- src/codegate/api/__init__.py | 0 src/codegate/api/v1.py | 69 +++++++++++++++++++++ src/codegate/api/v1_models.py | 44 +++++++++++++ src/codegate/db/connection.py | 9 ++- src/codegate/pipeline/workspace/commands.py | 15 +++-- src/codegate/server.py | 4 ++ tests/pipeline/workspace/test_workspace.py | 5 +- 7 files changed, 137 insertions(+), 9 deletions(-) create mode 100644 src/codegate/api/__init__.py create mode 100644 src/codegate/api/v1.py create mode 100644 src/codegate/api/v1_models.py diff --git a/src/codegate/api/__init__.py b/src/codegate/api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py new file mode 100644 index 00000000..f1dba385 --- /dev/null +++ b/src/codegate/api/v1.py @@ -0,0 +1,69 @@ +from fastapi import APIRouter, Response +from fastapi.exceptions import HTTPException +from fastapi.routing import APIRoute + +from codegate.api import v1_models +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 Response(status_code=204) + + +@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 + created = await wscrud.add_workspace(request.name) + + # TODO: refactor to use a more specific exception + if not created: + raise HTTPException(status_code=400, 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 diff --git a/src/codegate/api/v1_models.py b/src/codegate/api/v1_models.py new file mode 100644 index 00000000..55418e2b --- /dev/null +++ b/src/codegate/api/v1_models.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 diff --git a/src/codegate/db/connection.py b/src/codegate/db/connection.py index c2185d8f..006616ba 100644 --- a/src/codegate/db/connection.py +++ b/src/codegate/db/connection.py @@ -30,7 +30,6 @@ alert_queue = asyncio.Queue() fim_cache = FimCache() - class DbCodeGate: _instance = None @@ -256,7 +255,13 @@ 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) + except Exception as e: + logger.error(f"Failed to add workspace: {workspace_name}.", error=str(e)) + return None + return added_workspace async def update_session(self, session: Session) -> Optional[Session]: diff --git a/src/codegate/pipeline/workspace/commands.py b/src/codegate/pipeline/workspace/commands.py index 9651db8f..812d6c85 100644 --- a/src/codegate/pipeline/workspace/commands.py +++ b/src/codegate/pipeline/workspace/commands.py @@ -1,8 +1,8 @@ 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: @@ -18,15 +18,22 @@ async def add_workspace(self, new_workspace_name: str) -> bool: 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) 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]]: diff --git a/src/codegate/server.py b/src/codegate/server.py index b995fdd7..d1da668e 100644 --- a/src/codegate/server.py +++ b/src/codegate/server.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 diff --git a/tests/pipeline/workspace/test_workspace.py b/tests/pipeline/workspace/test_workspace.py index 85f10edc..d67c1f6c 100644 --- a/tests/pipeline/workspace/test_workspace.py +++ b/tests/pipeline/workspace/test_workspace.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