Refactoring: Entities - Sinks (#4973)

Author: warrkan

application/backend/app/api/dependencies.py

@@ -11,17 +11,19 @@
 
 from app.core.jobs.control_plane import JobQueue
 from app.db import get_db_session
+from app.models import Sink
 from app.scheduler import Scheduler
 from app.schemas import ProjectView
 from app.services import (
-    ConfigurationService,
     DatasetService,
     MetricsService,
     ModelService,
     PipelineMetricsService,
     PipelineService,
     ProjectService,
     ResourceNotFoundError,
+    SinkService,
+    SourceService,
     SystemService,
 )
 from app.services.base_weights_service import BaseWeightsService
@@ -146,12 +148,18 @@ def get_metrics_service(scheduler: Annotated[Scheduler, Depends(get_scheduler)])
     return MetricsService(scheduler.shm_metrics.name, scheduler.shm_metrics_lock)
 
 
-def get_configuration_service(
-    event_bus: Annotated[EventBus, Depends(get_event_bus)],
-    db: Annotated[Session, Depends(get_db)],
-) -> ConfigurationService:
-    """Provides a ConfigurationService instance."""
-    return ConfigurationService(event_bus=event_bus, db_session=db)
+def get_sink_service(
+    event_bus: Annotated[EventBus, Depends(get_event_bus)], db: Annotated[Session, Depends(get_db)]
+) -> SinkService:
+    """Provides a SinkService instance."""
+    return SinkService(event_bus=event_bus, db_session=db)
+
+
+def get_source_service(
+    event_bus: Annotated[EventBus, Depends(get_event_bus)], db: Annotated[Session, Depends(get_db)]
+) -> SourceService:
+    """Provides a SourceService instance."""
+    return SourceService(event_bus=event_bus, db_session=db)
 
 
 def get_pipeline_service(
@@ -225,6 +233,17 @@ def get_project(
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
 
 
+def get_sink(
+    sink_id: Annotated[UUID, Depends(get_sink_id)],
+    sink_service: Annotated[SinkService, Depends(get_sink_service)],
+) -> Sink:
+    """Provides a Sink instance for request scoped sink."""
+    try:
+        return sink_service.get_by_id(sink_id)
+    except ResourceNotFoundError as e:
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
+
+
 def get_base_weights_service(data_dir: Annotated[Path, Depends(get_data_dir)]) -> BaseWeightsService:
     """Provides a BaseWeightsService instance for managing base weights."""
     return BaseWeightsService(data_dir)

application/backend/app/api/routers/sinks.py

@@ -5,7 +5,6 @@
 
 import logging
 from typing import Annotated
-from uuid import UUID
 
 import yaml
 from fastapi import APIRouter, Body, Depends, File, UploadFile, status
@@ -14,15 +13,15 @@
 from fastapi.responses import FileResponse, Response
 from pydantic import ValidationError
 
-from app.api.dependencies import get_configuration_service, get_sink_id
-from app.schemas import Sink, SinkCreate
-from app.schemas.sink import SinkCreateAdapter
+from app.api.dependencies import get_sink, get_sink_service
+from app.api.schemas.sink import SinkCreate, SinkCreateAdapter, SinkView, SinkViewAdapter
+from app.models import Sink
 from app.services import (
-    ConfigurationService,
     ResourceInUseError,
     ResourceNotFoundError,
     ResourceWithIdAlreadyExistsError,
     ResourceWithNameAlreadyExistsError,
+    SinkService,
 )
 
 logger = logging.getLogger(__name__)
@@ -83,7 +82,7 @@
 @router.post(
     "",
     status_code=status.HTTP_201_CREATED,
-    response_model=Sink,
+    response_model=SinkView,
     responses={
         status.HTTP_201_CREATED: {"description": "Sink created"},
         status.HTTP_400_BAD_REQUEST: {"description": "Invalid sink ID"},
@@ -94,72 +93,84 @@ def create_sink(
     sink_create: Annotated[
         SinkCreate, Body(description=CREATE_SINK_BODY_DESCRIPTION, openapi_examples=CREATE_SINK_BODY_EXAMPLES)
     ],
-    configuration_service: Annotated[ConfigurationService, Depends(get_configuration_service)],
-) -> Sink:
+    sink_service: Annotated[SinkService, Depends(get_sink_service)],
+) -> SinkView:
     """Create and configure a new sink"""
     try:
-        return configuration_service.create_sink(sink_create)
+        sink = sink_service.create_sink(
+            name=sink_create.name,
+            sink_type=sink_create.sink_type,
+            rate_limit=sink_create.rate_limit,
+            config_data=sink_create.config_data,
+            output_formats=sink_create.output_formats,
+            sink_id=sink_create.id,
+        )
+        return SinkViewAdapter.validate_python(sink, from_attributes=True)
     except (ResourceWithNameAlreadyExistsError, ResourceWithIdAlreadyExistsError) as e:
         raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(e))
 
 
 @router.get(
     "",
     responses={
-        status.HTTP_200_OK: {"description": "List of available sink configurations", "model": list[Sink]},
+        status.HTTP_200_OK: {"description": "List of available sink configurations", "model": list[SinkView]},
     },
 )
 def list_sinks(
-    configuration_service: Annotated[ConfigurationService, Depends(get_configuration_service)],
-) -> list[Sink]:
+    sink_service: Annotated[SinkService, Depends(get_sink_service)],
+) -> list[SinkView]:
     """List the available sinks"""
-    return configuration_service.list_sinks()
+    sinks = sink_service.list_all()
+    return [SinkViewAdapter.validate_python(sink, from_attributes=True) for sink in sinks]
 
 
 @router.get(
     "/{sink_id}",
     responses={
-        status.HTTP_200_OK: {"description": "Sink found", "model": Sink},
+        status.HTTP_200_OK: {"description": "Sink found", "model": SinkView},
         status.HTTP_400_BAD_REQUEST: {"description": "Invalid sink ID"},
         status.HTTP_404_NOT_FOUND: {"description": "Sink not found"},
     },
 )
-def get_sink(
-    sink_id: Annotated[UUID, Depends(get_sink_id)],
-    configuration_service: Annotated[ConfigurationService, Depends(get_configuration_service)],
-) -> Sink:
+def get_sink_view(sink: Annotated[Sink, Depends(get_sink)]) -> SinkView:
     """Get info about a sink"""
-    try:
-        return configuration_service.get_sink_by_id(sink_id)
-    except ResourceNotFoundError as e:
-        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
+    return SinkViewAdapter.validate_python(sink, from_attributes=True)
 
 
 @router.patch(
     "/{sink_id}",
     responses={
-        status.HTTP_200_OK: {"description": "Sink successfully updated", "model": Sink},
+        status.HTTP_200_OK: {"description": "Sink successfully updated", "model": SinkView},
         status.HTTP_400_BAD_REQUEST: {"description": "Invalid sink ID or request body"},
         status.HTTP_404_NOT_FOUND: {"description": "Sink not found"},
         status.HTTP_409_CONFLICT: {"description": "Sink already exists"},
     },
 )
 def update_sink(
-    sink_id: Annotated[UUID, Depends(get_sink_id)],
+    sink: Annotated[Sink, Depends(get_sink)],
     sink_config: Annotated[
         dict,
         Body(
             description=UPDATE_SINK_BODY_DESCRIPTION,
             openapi_examples=UPDATE_SINK_BODY_EXAMPLES,
         ),
     ],
-    configuration_service: Annotated[ConfigurationService, Depends(get_configuration_service)],
-) -> Sink:
+    sink_service: Annotated[SinkService, Depends(get_sink_service)],
+) -> SinkView:
     """Reconfigure an existing sink"""
     if "sink_type" in sink_config:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="The 'sink_type' field cannot be changed")
     try:
-        return configuration_service.update_sink(sink_id, sink_config)
+        updated_sink: Sink = sink.model_copy(update=sink_config)
+        updated_sink.config_data = updated_sink.config_data.model_copy(update=sink_config)
+        sink = sink_service.update_sink(
+            sink=sink,
+            new_name=updated_sink.name,
+            new_rate_limit=updated_sink.rate_limit,
+            new_config_data=updated_sink.config_data,
+            new_output_formats=updated_sink.output_formats,
+        )
+        return SinkViewAdapter.validate_python(sink, from_attributes=True)
     except ResourceNotFoundError as e:
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
     except ResourceWithNameAlreadyExistsError as e:
@@ -180,47 +191,45 @@ def update_sink(
         status.HTTP_404_NOT_FOUND: {"description": "Sink not found"},
     },
 )
-def export_sink(
-    sink_id: Annotated[UUID, Depends(get_sink_id)],
-    configuration_service: Annotated[ConfigurationService, Depends(get_configuration_service)],
-) -> Response:
+def export_sink(sink: Annotated[Sink, Depends(get_sink)]) -> Response:
     """Export a sink to file"""
-    sink = configuration_service.get_sink_by_id(sink_id)
-    if not sink:
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND,
-            detail=f"Sink with ID {sink_id} not found",
-        )
-
     yaml_content = yaml.safe_dump(sink.model_dump(mode="json", exclude={"id"}))
 
     return Response(
         content=yaml_content.encode("utf-8"),
         media_type="application/x-yaml",
-        headers={"Content-Disposition": f"attachment; filename=sink_{sink_id}.yaml"},
+        headers={"Content-Disposition": f"attachment; filename=sink_{sink.id}.yaml"},
     )
 
 
 @router.post(
     ":import",
     status_code=status.HTTP_201_CREATED,
     responses={
-        status.HTTP_201_CREATED: {"description": "Sink imported successfully", "model": Sink},
+        status.HTTP_201_CREATED: {"description": "Sink imported successfully", "model": SinkView},
         status.HTTP_400_BAD_REQUEST: {"description": "Invalid YAML format"},
         status.HTTP_409_CONFLICT: {"description": "Sink already exists"},
         status.HTTP_422_UNPROCESSABLE_ENTITY: {"description": "Validation error(s)"},
     },
 )
 def import_sink(
     yaml_file: Annotated[UploadFile, File(description="YAML file containing the sink configuration")],
-    configuration_service: Annotated[ConfigurationService, Depends(get_configuration_service)],
-) -> Sink:
+    sink_service: Annotated[SinkService, Depends(get_sink_service)],
+) -> SinkView:
     """Import a sink from file"""
     try:
         yaml_content = yaml_file.file.read()
         sink_data = yaml.safe_load(yaml_content)
         sink_create = SinkCreateAdapter.validate_python(sink_data)
-        return configuration_service.create_sink(sink_create)
+        sink = sink_service.create_sink(
+            name=sink_create.name,
+            sink_type=sink_create.sink_type,
+            rate_limit=sink_create.rate_limit,
+            config_data=sink_create.config_data,
+            output_formats=sink_create.output_formats,
+            sink_id=sink_create.id,
+        )
+        return SinkViewAdapter.validate_python(sink, from_attributes=True)
     except yaml.YAMLError as e:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid YAML format: {str(e)}")
     except (ResourceWithNameAlreadyExistsError, ResourceWithIdAlreadyExistsError) as e:
@@ -242,13 +251,11 @@ def import_sink(
     },
 )
 def delete_sink(
-    sink_id: Annotated[UUID, Depends(get_sink_id)],
-    configuration_service: Annotated[ConfigurationService, Depends(get_configuration_service)],
+    sink: Annotated[Sink, Depends(get_sink)],
+    sink_service: Annotated[SinkService, Depends(get_sink_service)],
 ) -> None:
     """Remove a sink"""
     try:
-        configuration_service.delete_sink_by_id(sink_id)
-    except ResourceNotFoundError as e:
-        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
+        sink_service.delete_sink(sink)
     except ResourceInUseError as e:
         raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(e))