Merge pull request #4 from ia-labo/feat/sql_db

src/models/bbox_manager.py

@@ -0,0 +1,64 @@
+"""Manage bounding boxes in the database."""
+
+from sqlalchemy import text
+
+from models.sql_connection import get_db_connection
+
+
+class BoundingBoxManager:
+    def __init__(self):
+        """Initialise connection and session."""
+        self.engine, self.session = get_db_connection()
+
+    def add_bbox(self, confidence, class_id, img_id, x_min, y_min, x_max, y_max):
+        """
+        Add a bounding box to the `bboxes` table.
+
+        Args:
+            confidence (float): Confidence of the detection.
+            class_id (str): Class ID of the detection.
+            img_id (int): ID of the image where the bounding box was detected.
+            x_min (float): Minimum X coordinate.
+            y_min (float): Minimum Y coordinate.
+            x_max (float): Maximum X coordinate.
+            y_max (float): Maximum Y coordinate.
+
+        Returns:
+            dict: Information of the added bounding box.
+        """
+        query = text(
+            """
+            INSERT INTO bboxes (confidence, class_id, img_id, x_min, y_min, x_max, y_max)
+            VALUES (:confidence, :class_id, :img_id, :x_min, :y_min, :x_max, :y_max)
+        """
+        )
+        try:
+            self.session.execute(
+                query,
+                {
+                    "confidence": confidence,
+                    "class_id": class_id,
+                    "img_id": img_id,
+                    "x_min": x_min,
+                    "y_min": y_min,
+                    "x_max": x_max,
+                    "y_max": y_max,
+                },
+            )
+            self.session.commit()
+            return {
+                "confidence": confidence,
+                "class_id": class_id,
+                "img_id": img_id,
+                "x_min": x_min,
+                "y_min": y_min,
+                "x_max": x_max,
+                "y_max": y_max,
+            }
+        except Exception as e:
+            self.session.rollback()
+            raise Exception(f"An error occurred while adding the bounding box: {e}")
+
+    def close_connection(self):
+        """Close the connection."""
+        self.session.close()

src/models/image_manager.py

@@ -0,0 +1,99 @@
+"""Mnages images table in the database."""
+
+from sqlalchemy import text
+
+from 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, created_at, park_id=None):
+        """
+        Add an image to the `images` table.
+
+        Args:
+            name (str): Image name.
+            created_at (str): Image creation date. Format: "YYYY-MM-DD HH:MM:SS".
+            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)
+        """
+        )
+        try:
+            self.session.execute(
+                query, {"name": name, "created_at": created_at, "park_id": park_id}
+            )
+            self.session.commit()
+            return {
+                "name": name,
+                "created_at": created_at,
+                "park_id": park_id,
+            }
+        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}).fetchone()
+            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(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 [dict(row) 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()

src/models/park_manager.py

@@ -0,0 +1,76 @@
+"""Manage parks table in the database."""
+
+from sqlalchemy import text
+
+from models.sql_connection import get_db_connection
+
+
+class ParkManager:
+    def __init__(self):
+        """Initialise connection and session."""
+        self.engine, self.session = get_db_connection()
+
+    def add_park(self, name):
+        query = text(
+            """
+            INSERT INTO parks (name)
+            VALUES (:name)
+        """
+        )
+        try:
+            self.session.execute(query, {"name": name})
+            self.session.commit()
+            return {"name": name}
+        except Exception as e:
+            self.session.rollback()
+            raise Exception(f"An error occurred while adding the park: {e}")
+
+    def get_parks(self):
+        """
+        get all parks from the `parks` table.
+
+        Returns:
+            list[dict]: list of parks.
+        """
+        query = text("SELECT * FROM parks")
+        try:
+            result = self.session.execute(query)
+            return [dict(row) for row in result]
+        except Exception as e:
+            raise Exception(f"An error occurred while adding the park: {e}")
+
+    def get_park_id(self, park_name):
+        """Get the park ID from the park name.
+
+        Args:
+            park_name (str): Name of the park.
+        """
+        query = text("SELECT id FROM parks WHERE name = :park_name")
+        try:
+            result = self.session.execute(query, {"park_name": park_name}).fetchone()
+            return result[0] if result else None
+        except Exception as e:
+            raise Exception(f"An error occurred while getting the park ID: {e}")
+
+    def delete_park(self, park_id):
+        """
+        Delete a park from the `parks` table.
+
+        Args:
+            park_id (int): ID of the park to delete.
+
+        Returns:
+            bool: True if the park was deleted, False otherwise.
+        """
+        query = text("DELETE FROM parks WHERE id = :park_id")
+        try:
+            result = self.session.execute(query, {"park_id": park_id})
+            self.session.commit()
+            return result.rowcount > 0
+        except Exception as e:
+            self.session.rollback()
+            raise Exception(f"An error occurred while adding the park: {e}")
+
+    def close_connection(self):
+        """Close the connection."""
+        self.session.close()

src/models/sql_connection.py

@@ -0,0 +1,30 @@
+""""Create a connection to the database and return the engine and session objects."""
+
+import os
+
+from dotenv import load_dotenv
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+load_dotenv()
+HOST = os.getenv("DB_HOST")
+PORT = os.getenv("DB_PORT")
+DBNAME = os.getenv("DB_DATABASE")
+USERNAME = os.getenv("DB_USERNAME")
+PASSWORD = os.getenv("DB_PASSWORD")
+
+
+def get_db_connection() -> tuple:
+    """Create a connection to the database and return the engine and session objects.
+
+    Returns:
+        engine: connection to the database
+        session: session to interact with the database
+    """
+    database_url = f"postgresql+psycopg2://{USERNAME}:{PASSWORD}@{HOST}:{PORT}/{DBNAME}"
+
+    engine = create_engine(database_url)
+    Session = sessionmaker(bind=engine)
+    session = Session()
+
+    return engine, session