Spaces:
Sleeping
Sleeping
| """Manages images table in the database.""" | |
| from datetime import datetime | |
| from sqlalchemy import text | |
| from data_models.sql_connection import get_db_connection | |
| class ImageManager: | |
| """ImageManager class.""" | |
| def __init__(self): | |
| """Initialise connection and session.""" | |
| self.engine, self.session = get_db_connection() | |
| def add_image(self, name: str, created_at: datetime, park_id=None) -> int: | |
| """ | |
| Add an image to the `images` table. | |
| Args: | |
| name (str): Image name. | |
| created_at (datetime): Image creation date. | |
| park_id (int):id of the park where the image was taken. | |
| Returns: | |
| dict: Information of the added image. | |
| """ | |
| query = text( | |
| """ | |
| INSERT INTO images (name, created_at, park_id) | |
| VALUES (:name, :created_at, :park_id) | |
| RETURNING id | |
| """ | |
| ) | |
| try: | |
| response = self.session.execute( | |
| query, {"name": name, "created_at": created_at, "park_id": park_id} | |
| ).fetchone() | |
| self.session.commit() | |
| return response[0] | |
| except Exception as e: | |
| self.session.rollback() | |
| raise Exception(f"An error occurred while adding the image: {e}") | |
| def get_image_id(self, image_name): | |
| """ | |
| Get the image ID from the image name. | |
| Args: | |
| image_name (str): Name of the image. | |
| Returns: | |
| int: ID of the image. | |
| """ | |
| query = text("SELECT id FROM images WHERE name = :image_name") | |
| try: | |
| result = self.session.execute(query, {"image_name": image_name}) | |
| return result[0] if result else None | |
| except Exception as e: | |
| raise Exception(f"An error occurred while getting the image ID: {e}") | |
| def get_images_by_park(self, park_id): | |
| """ | |
| Get all images from the `images` table. | |
| Args: | |
| park_id (int): ID of the park to filter images. | |
| Returns: | |
| list[dict]: List of images. | |
| """ | |
| query = text("SELECT * FROM images") | |
| if park_id: | |
| query = text("SELECT * FROM images WHERE park_id = :park_id") | |
| try: | |
| result = self.session.execute( | |
| query, {"park_id": park_id} if park_id else {} | |
| ) | |
| return [row._asdict() for row in result] | |
| except Exception as e: | |
| raise Exception(f"Erreur lors de la récupération des images : {e}") | |
| def delete_image(self, image_id): | |
| """ | |
| Delete an image by its ID. | |
| Args: | |
| image_id (int): Image ID. | |
| Returns: | |
| bool: True if the deletion was successful, False otherwise. | |
| """ | |
| query = text("DELETE FROM images WHERE id = :image_id") | |
| try: | |
| result = self.session.execute(query, {"image_id": image_id}) | |
| self.session.commit() | |
| return result.rowcount > 0 | |
| except Exception as e: | |
| self.session.rollback() | |
| raise Exception(f"Erreur lors de la suppression de l'image : {e}") | |
| def close_connection(self): | |
| """Close the connection.""" | |
| self.session.close() | |
| def get_image_count(self): | |
| """ | |
| Get the number of images in the `images` table. | |
| Returns: | |
| int: Number of images. | |
| """ | |
| query = text("SELECT COUNT(*) FROM images") | |
| try: | |
| result = self.session.execute(query).fetchone() | |
| return result[0] | |
| except Exception as e: | |
| raise Exception(f"An error occurred while getting the image count: {e}") |