diff --git a/database/repositories/__init__.py b/database/repositories/__init__.py index 24fcbf5..8bd8ba9 100644 --- a/database/repositories/__init__.py +++ b/database/repositories/__init__.py @@ -1,3 +1,4 @@ from .user_repository import UserRepository +from .robot_repository import RobotRepository -__all__ = ['UserRepository'] +__all__ = ['UserRepository', 'RobotRepository'] diff --git a/database/repositories/robot_repository.py b/database/repositories/robot_repository.py new file mode 100644 index 0000000..904224c --- /dev/null +++ b/database/repositories/robot_repository.py @@ -0,0 +1,134 @@ +# database/repositories/robot_repository.py +from typing import List, Optional +from database.connection import get_connection +from model.robot import Robot +import logging + +class RobotRepository: + def get_all(self) -> List[Robot]: + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute("SELECT * FROM robots ORDER BY id") + return [ + Robot( + id=row[0], + status=row[1], + battery_level=row[2], + last_update=row[3], + current_zone=row[4], + current_row=row[5], + current_shelf=row[6] + ) for row in cur.fetchall() + ] + except Exception as e: + print(f"Ошика получения всех роботов: {e}") + return [] + + def get_by_id(self, robot_id: str) -> Optional[Robot]: + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute("SELECT * FROM robots WHERE id = %s", (robot_id,)) + row = cur.fetchone() + return Robot(*row) if row else None + except Exception as e: + print(f"Ошибка получения роботов {robot_id}: {e}") + return None + + def update_robot(self, robot_id: str, status: str = None, battery_level: int = None, + current_zone: str = None, current_row: int = None, current_shelf: int = None) -> bool: + try: + with get_connection() as conn: + with conn.cursor() as cur: + updates = ["last_update = CURRENT_TIMESTAMP"] + params = [] + + if status is not None: + updates.append("status = %s") + params.append(status) + + if battery_level is not None: + updates.append("battery_level = %s") + params.append(battery_level) + + if current_zone is not None: + updates.append("current_zone = %s") + params.append(current_zone) + + if current_row is not None: + updates.append("current_row = %s") + params.append(current_row) + + if current_shelf is not None: + updates.append("current_shelf = %s") + params.append(current_shelf) + + params.append(robot_id) + query = f"UPDATE robots SET {', '.join(updates)} WHERE id = %s" + + cur.execute(query, params) + conn.commit() + return cur.rowcount > 0 + except Exception as e: + print(f"Ошибка обновления робота {robot_id}: {e}") + return False + + def get_robots_by_status(self, status: str) -> List[Robot]: + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute("SELECT * FROM robots WHERE status = %s ORDER BY id", (status,)) + return [Robot(*row) for row in cur.fetchall()] + except Exception as e: + print(f"Ошибка получения роботов по статусу {status}: {e}") + return [] + + def get_low_battery_robots(self, threshold: int = 20) -> List[Robot]: + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute(""" + SELECT * FROM robots + WHERE battery_level < %s AND status = 'active' + ORDER BY battery_level ASC + """, (threshold,)) + return [Robot(*row) for row in cur.fetchall()] + except Exception as e: + print(f"Ошибка получения роботов с низким зарядом: {e}") + return [] + + def get_robots_in_zone(self, zone: str) -> List[Robot]: + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute("SELECT * FROM robots WHERE current_zone = %s ORDER BY id", (zone,)) + return [Robot(*row) for row in cur.fetchall()] + except Exception as e: + print(f"Ошибка получения роботов в зоне {zone}: {e}") + return [] + + def create_robot(self, robot_id: str, status: str = 'active', battery_level: int = 100) -> bool: + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute(""" + INSERT INTO robots (id, status, battery_level, last_update) + VALUES (%s, %s, %s, CURRENT_TIMESTAMP) + """, (robot_id, status, battery_level)) + conn.commit() + return cur.rowcount > 0 + except Exception as e: + print(f"Ошибка создания робота {robot_id}: {e}") + return False + + def delete_robot(self, robot_id: str) -> bool: + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute("DELETE FROM robots WHERE id = %s", (robot_id,)) + conn.commit() + return cur.rowcount > 0 + except Exception as e: + print(f"Ошибка удаления робота {robot_id}: {e}") + return False \ No newline at end of file diff --git a/database/repositories/user_repository.py b/database/repositories/user_repository.py index 35f24a5..d624842 100644 --- a/database/repositories/user_repository.py +++ b/database/repositories/user_repository.py @@ -2,7 +2,6 @@ from typing import List, Optional from model.user import User from database.connection import get_connection - class UserRepository: def get_all(self) -> List[User]: with get_connection() as conn: