Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions src/backend/common/database/database_base.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@

"""Database base class for managing database operations."""

# pylint: disable=unnecessary-pass

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Type

Expand All @@ -21,25 +24,30 @@ class DatabaseBase(ABC):
@abstractmethod
async def initialize(self) -> None:
"""Initialize the database client and create containers if needed."""
pass

@abstractmethod
async def close(self) -> None:
"""Close database connection."""
pass

# Core CRUD Operations
@abstractmethod
async def add_item(self, item: BaseDataModel) -> None:
"""Add an item to the database."""
pass

@abstractmethod
async def update_item(self, item: BaseDataModel) -> None:
"""Update an item in the database."""
pass

@abstractmethod
async def get_item_by_id(
self, item_id: str, partition_key: str, model_class: Type[BaseDataModel]
) -> Optional[BaseDataModel]:
"""Retrieve an item by its ID and partition key."""
pass

@abstractmethod
async def query_items(
Expand All @@ -49,92 +57,113 @@ async def query_items(
model_class: Type[BaseDataModel],
) -> List[BaseDataModel]:
"""Query items from the database and return a list of model instances."""
pass

@abstractmethod
async def delete_item(self, item_id: str, partition_key: str) -> None:
"""Delete an item from the database."""
pass

# Plan Operations
@abstractmethod
async def add_plan(self, plan: Plan) -> None:
"""Add a plan to the database."""
pass

@abstractmethod
async def update_plan(self, plan: Plan) -> None:
"""Update a plan in the database."""
pass

@abstractmethod
async def get_plan_by_plan_id(self, plan_id: str) -> Optional[Plan]:
"""Retrieve a plan by plan_id."""
pass

@abstractmethod
async def get_plan(self, plan_id: str) -> Optional[Plan]:
"""Retrieve a plan by plan_id."""
pass

@abstractmethod
async def get_all_plans(self) -> List[Plan]:
"""Retrieve all plans for the user."""
pass

@abstractmethod
async def get_all_plans_by_team_id(self, team_id: str) -> List[Plan]:
"""Retrieve all plans for a specific team."""
pass

@abstractmethod
async def get_all_plans_by_team_id_status(
self, user_id: str, team_id: str, status: str
) -> List[Plan]:
"""Retrieve all plans for a specific team."""
pass

# Step Operations
@abstractmethod
async def add_step(self, step: Step) -> None:
"""Add a step to the database."""
pass

@abstractmethod
async def update_step(self, step: Step) -> None:
"""Update a step in the database."""
pass

@abstractmethod
async def get_steps_by_plan(self, plan_id: str) -> List[Step]:
"""Retrieve all steps for a plan."""
pass

@abstractmethod
async def get_step(self, step_id: str, session_id: str) -> Optional[Step]:
"""Retrieve a step by step_id and session_id."""
pass

# Team Operations
@abstractmethod
async def add_team(self, team: TeamConfiguration) -> None:
"""Add a team configuration to the database."""
pass

@abstractmethod
async def update_team(self, team: TeamConfiguration) -> None:
"""Update a team configuration in the database."""
pass

@abstractmethod
async def get_team(self, team_id: str) -> Optional[TeamConfiguration]:
"""Retrieve a team configuration by team_id."""
pass

@abstractmethod
async def get_team_by_id(self, team_id: str) -> Optional[TeamConfiguration]:
"""Retrieve a team configuration by internal id."""
pass

@abstractmethod
async def get_all_teams(self) -> List[TeamConfiguration]:
"""Retrieve all team configurations for the given user."""
pass

@abstractmethod
async def delete_team(self, team_id: str) -> bool:
"""Delete a team configuration by team_id and return True if deleted."""
pass

# Data Management Operations
@abstractmethod
async def get_data_by_type(self, data_type: str) -> List[BaseDataModel]:
"""Retrieve all data of a specific type."""
pass

@abstractmethod
async def get_all_items(self) -> List[Dict[str, Any]]:
"""Retrieve all items as dictionaries."""
pass

# Context Manager Support
async def __aenter__(self):
Expand All @@ -149,14 +178,17 @@ async def __aexit__(self, exc_type, exc, tb):
@abstractmethod
async def get_steps_for_plan(self, plan_id: str) -> List[Step]:
"""Convenience method aliasing get_steps_by_plan for compatibility."""
pass

@abstractmethod
async def get_current_team(self, user_id: str) -> Optional[UserCurrentTeam]:
"""Retrieve the current team for a user."""
pass

@abstractmethod
async def delete_current_team(self, user_id: str) -> Optional[UserCurrentTeam]:
"""Retrieve the current team for a user."""
pass

@abstractmethod
async def set_current_team(self, current_team: UserCurrentTeam) -> None:
Expand All @@ -165,22 +197,27 @@ async def set_current_team(self, current_team: UserCurrentTeam) -> None:
@abstractmethod
async def update_current_team(self, current_team: UserCurrentTeam) -> None:
"""Update the current team for a user."""
pass

@abstractmethod
async def delete_plan_by_plan_id(self, plan_id: str) -> bool:
"""Retrieve the current team for a user."""
pass

@abstractmethod
async def add_mplan(self, mplan: messages.MPlan) -> None:
"""Add a team configuration to the database."""
pass

@abstractmethod
async def update_mplan(self, mplan: messages.MPlan) -> None:
"""Update a team configuration in the database."""
pass

@abstractmethod
async def get_mplan(self, plan_id: str) -> Optional[messages.MPlan]:
"""Retrieve a mplan configuration by plan_id."""
pass

@abstractmethod
async def add_agent_message(self, message: AgentMessageData) -> None:
Expand All @@ -189,7 +226,9 @@ async def add_agent_message(self, message: AgentMessageData) -> None:
@abstractmethod
async def update_agent_message(self, message: AgentMessageData) -> None:
"""Update an agent message in the database."""
pass

@abstractmethod
async def get_agent_messages(self, plan_id: str) -> Optional[AgentMessageData]:
"""Retrieve an agent message by message_id."""
pass
Loading