helpers/firebase.py

import firebase_admin
from firebase_admin import credentials
from firebase_admin import firestore


class FirebaseClient:
    def __init__(self, path_to_certificate):
        # Initialize Firebase Admin SDK
        cred = credentials.Certificate(path_to_certificate)  # Path to your service account key JSON file
        firebase_admin.initialize_app(cred)

        # Initialize Firestore database
        self.db = firestore.client()

    def add_task(self, task_data):
        """
        Add a new task to Firestore.
        
        Args:
            task_data (dict): Dictionary containing task data.
                Example: {'title': 'Task Title', 'description': 'Task Description', 'status': 'pending'}
        """
        # Add task data to Firestore
        doc_ref = self.db.collection('tasks').document()
        doc_ref.set(task_data)
        return doc_ref.id
    
    def get_task_by_status(self, status):
        # Reference to the tasks collection
        tasks_ref = self.db.collection('tasks')

        # Query tasks with status 'pending'
        query = tasks_ref.where('status', '==', status)

        # Get documents that match the query
        pending_tasks = query.stream()

        # Convert documents to dictionaries
        pending_tasks_data = []
        for doc in pending_tasks:
            task_data = doc.to_dict() 
            task_data['id'] = doc.id
            pending_tasks_data.append(task_data)

        return pending_tasks_data

    def get_all_tasks(self):
        """
        Retrieve all tasks from Firestore.
        
        Returns:
            list: A list containing dictionaries, each representing a task.
        """
        # Reference to the 'tasks' collection
        tasks_ref = self.db.collection('tasks')
        
        # Get all documents in the collection
        docs = tasks_ref.stream()
        
        # Initialize an empty list to store tasks
        tasks = []
        
        # Iterate over each document and add it to the tasks list
        for doc in docs:
            doc_dict = doc.to_dict()
            doc_dict['id'] = doc.id
            tasks.append(doc_dict)
        
        return tasks

    def update(self, task_id, data):
        """
        Reserve a task by a worker.
        
        Args:
            task_id (str): ID of the task to be reserved.
            worker_id (str): ID of the worker reserving the task.
        """
        # Reference to the task document
        task_ref = self.db.collection('tasks').document(task_id)
        
        # Update the task document to indicate it has been reserved by the worker
        task_ref.update(data)
    
    def delete_task(self, task_id):
        """
        Delete a task from Firestore by its ID.
        
        Args:
            task_id (str): ID of the task to be deleted.
        """
        # Reference to the task document
        task_ref = self.db.collection('tasks').document(task_id)
        
        # Delete the task document
        task_ref.delete()
    
    def get_task_by_id(self, task_id):
        """
        Retrieve a task from Firestore by its ID.
        
        Args:
            task_id (str): ID of the task to be retrieved.
        
        Returns:
            dict or None: Dictionary containing the task data if found, None otherwise.
        """
        # Reference to the task document
        task_ref = self.db.collection('tasks').document(task_id)
        
        # Retrieve the task document
        task_doc = task_ref.get()
        
        # Check if the task document exists
        if task_doc.exists:
            return task_doc.to_dict()
        else:
            return None
        
    def find_tasks_by_status(self, status):
        """
        Find all tasks in Firestore with the specified status.
        
        Args:
            status (str): Status value to filter tasks by.
        
        Returns:
            list: List of dictionaries containing task data.
        """
        # Reference to the 'tasks' collection
        tasks_ref = self.db.collection('tasks')
        
        # Query tasks with the specified status
        query = tasks_ref.where('status', '==', status)
        
        # Get documents that match the query
        docs = query.stream()
        
        # Initialize an empty list to store tasks
        tasks = []
        
        # Iterate over each document and add it to the tasks list
        for doc in docs:
            task = doc.to_dict()
            task['id'] = doc.id
            tasks.append(task)
        
        return tasks