theatre-system/db.py

import psycopg2
from config import DB_HOST, DB_PORT, DB_NAME

class Database:
    def __init__(self):
        self.conn = None

    def connect(self, username, password):
        self.conn = psycopg2.connect(
            host=DB_HOST,
            port=DB_PORT,
            dbname=DB_NAME,
            user=username,
            password=password
        )
        return self.conn

    def execute(self, query, params=None, fetch=False):
        try:
            with self.conn.cursor() as cur:
                cur.execute(query, params)
                if fetch:
                    return cur.fetchall()
                self.conn.commit()
                return True
        except Exception as e:
            self.conn.rollback()
            raise e

    #  Отделы 
    def get_departments(self):
        return self.execute(
            "SELECT department_id, head_last_name, head_first_name, head_middle_name, email, details "
            "FROM accounting_department ORDER BY department_id", fetch=True
        )

    def get_departments_list(self):
        rows = self.execute(
            "SELECT department_id, head_last_name, head_first_name FROM accounting_department ORDER BY department_id",
            fetch=True
        )
        return [(row[0], f"{row[0]}: {row[1]} {row[2]}") for row in rows]

    #  Пользователи 
    def _get_next_user_id(self):
        rows = self.execute("SELECT user_id FROM users ORDER BY user_id", fetch=True)
        existing = {row[0] for row in rows}
        next_id = 1
        while next_id in existing:
            next_id += 1
        return next_id

    def get_users(self):
        return self.execute(
            "SELECT user_id, last_name, first_name, middle_name, birth_date, phone "
            "FROM users ORDER BY user_id", fetch=True
        )

    def get_users_list(self):
        rows = self.execute(
            "SELECT user_id, last_name, first_name FROM users ORDER BY user_id",
            fetch=True
        )
        return [(row[0], f"{row[0]}: {row[1]} {row[2]}") for row in rows]

    def add_user(self, last_name, first_name, middle_name, birth_date, phone):
        new_id = self._get_next_user_id()
        self.execute(
            "INSERT INTO users (user_id, last_name, first_name, middle_name, birth_date, phone) "
            "VALUES (%s, %s, %s, %s, %s, %s)",
            (new_id, last_name, first_name, middle_name, birth_date, phone)
        )

    def update_user(self, user_id, last_name, first_name, middle_name, birth_date, phone):
        self.execute(
            "UPDATE users SET last_name=%s, first_name=%s, middle_name=%s, birth_date=%s, phone=%s "
            "WHERE user_id=%s",
            (last_name, first_name, middle_name, birth_date, phone, user_id)
        )

    def delete_user(self, user_id):
        self.execute("DELETE FROM users WHERE user_id=%s", (user_id,))

    #  Билеты (представление tickets_full_price) 
    def _get_next_ticket_id(self):
        rows = self.execute("SELECT ticket_id FROM tickets ORDER BY ticket_id", fetch=True)
        existing = {row[0] for row in rows}
        next_id = 1
        while next_id in existing:
            next_id += 1
        return next_id

    def get_tickets(self): 
        return self.execute(
            "SELECT ticket_id, department_id, user_id, performance_name, ticket_price, hall_seat, event_datetime, full_price "
            "FROM tickets_full_price ORDER BY ticket_id", fetch=True
        )

    def add_ticket(self, department_id, user_id, performance_name, ticket_price, hall_seat, event_datetime):
        new_id = self._get_next_ticket_id()
        self.execute(
            "INSERT INTO tickets (ticket_id, department_id, user_id, performance_name, ticket_price, hall_seat, event_datetime) "
            "VALUES (%s, %s, %s, %s, %s, %s, %s)",
            (new_id, department_id, user_id, performance_name, ticket_price, hall_seat, event_datetime)
        )

    def update_ticket(self, ticket_id, department_id, user_id, performance_name, ticket_price, hall_seat, event_datetime):
        self.execute(
            "UPDATE tickets SET department_id=%s, user_id=%s, performance_name=%s, ticket_price=%s, hall_seat=%s, event_datetime=%s "
            "WHERE ticket_id=%s",
            (department_id, user_id, performance_name, ticket_price, hall_seat, event_datetime, ticket_id)
        )

    def delete_ticket(self, ticket_id):
        self.execute("DELETE FROM tickets WHERE ticket_id=%s", (ticket_id,))

    #  Поиск 
    def search_tickets(self, last_name=None, performance=None):
        sql = """
            SELECT t.ticket_id, u.last_name, u.first_name, t.performance_name,
                   t.ticket_price, t.hall_seat, t.event_datetime,
                   ROUND(t.ticket_price * 1.10, 2) AS full_price
            FROM tickets t
            JOIN users u ON t.user_id = u.user_id
            WHERE 1=1
        """
        params = []
        if last_name:
            sql += " AND u.last_name ILIKE %s"
            params.append(f"%{last_name}%")
        if performance:
            sql += " AND t.performance_name ILIKE %s"
            params.append(f"%{performance}%")
        sql += " ORDER BY t.event_datetime"
        return self.execute(sql, params, fetch=True)

    #  Отчёт 
    def revenue_by_performance(self, performance_name):
        result = self.execute("SELECT get_revenue(%s)", (performance_name,), fetch=True)
        return result[0][0] if result else 0