Kita Trofimov
133 lines
5.6 KiB
Python
133 lines
5.6 KiB
Python
# database/repositories/robot_repository.py
|
|
from typing import List, Optional
|
|
from database.connection import get_connection
|
|
from model.robot import Robot
|
|
|
|
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 |