backend/app/database/travels.py

from typing import List, Dict, Set
from tinydb import Query
from backend.app.database.base import BaseDB

class TravelDB(BaseDB):
    def __init__(self):
        super().__init__()
        self.table = self.db.table('travels')
        self.query = Query()

    def create_travels(self, items: List[Dict]) -> List[Dict]:
        if items:
            inserted_ids = self.table.insert_multiple(items)
            return [self.table.get(doc_id=id) for id in inserted_ids]
        return []


    def update_travels(self , items: List[Dict]) -> List[Dict]:
        updated = []
        for item in items:
            self.table.update(item, self.query.id == item['id'])
            updated.append(self.table.get(self.query.id == item['id']))
        return updated


    def get_existing_travel_ids( self , input_ids: Set[str]) -> Set[str]:
        existing = self.table.search(self.query.id.one_of(list(input_ids)))
        return {item['id'] for item in existing}



    def delete_travels(self ,item_ids: List[str]) -> List[str]:
        self.table.remove(self.query.id.one_of(item_ids))
        return item_ids

    def get_travels(self , item_ids: List[str]) -> List[Dict]:
        return self.table.search(self.query.id.one_of(item_ids))

    def get_existing_items(self, item_ids: List[str]) -> List[Dict]:
        return self.table.search(self.query.id.one_of(item_ids))
    

    def get_items_by_embedding_state(self, states: List[str]):
        return [item for item in self.table if item.get("embedded") in states]