From 052c93928c2fd2f2f415f1e5791400ace925c035 Mon Sep 17 00:00:00 2001 From: Kita Trofimov <144846094+NiTro005@users.noreply.github.com> Date: Mon, 27 Oct 2025 23:02:19 +0300 Subject: [PATCH] feat(db): Database initialization has been added --- api/loginapi.py | 3 +- db/connection.py | 58 ++++++++++--------- db/init.py | 147 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 180 insertions(+), 28 deletions(-) create mode 100644 db/init.py diff --git a/api/loginapi.py b/api/loginapi.py index 3ca2e88..299d109 100644 --- a/api/loginapi.py +++ b/api/loginapi.py @@ -1,5 +1,6 @@ from flask import Blueprint, request -from model.user import user + +from model import user loginBP = Blueprint("loginapi", __name__) diff --git a/db/connection.py b/db/connection.py index 3950913..20e1425 100644 --- a/db/connection.py +++ b/db/connection.py @@ -1,21 +1,29 @@ import psycopg2 import os -import logging from contextlib import contextmanager from typing import Generator +from loguru import logger from utils.loadDotEnv import initializeENV initializeENV() -logger = logging.getLogger(__name__) def PSQLConnect(): - conn = psycopg2.connect(os.getenv('POSTDRESS_CONNECTION')) + conn_str = os.getenv('POSTGRES_CONNECTION') + + if not conn_str: + logger.error("POSTGRES_CONNECTION не найден в .env файле") + raise ValueError("POSTGRES_CONNECTION не найден в .env файле") + + conn = psycopg2.connect(conn_str) + logger.debug("Подключение к БД установлено") return conn + def PSQLCursor(conn): - cur = conn.cursor() + cur = conn.cursor() + logger.debug("Курсор БД создан") return cur @@ -24,38 +32,34 @@ def get_connection() -> Generator[psycopg2.extensions.connection, None, None]: conn = None try: conn = PSQLConnect() - logger.debug("Подключение к БД установлено") + logger.debug("Контекст подключения к БД открыт") yield conn - except psycopg2.OperationalError as e: - logger.error(f"Ошибка подключения к БД: {e}") - raise - except psycopg2.Error as e: - logger.error(f"Ошибка PostgreSQL: {e}") - raise except Exception as e: - logger.error(f"Неожиданная ошибка при работе с БД: {e}") + logger.error(f"Ошибка в контексте подключения: {e}") + if conn: + conn.rollback() + logger.debug("Откат транзакции выполнен") raise finally: if conn: - try: - conn.close() - logger.debug("Соединение с БД закрыто") - except Exception as e: - logger.warning(f"Ошибка при закрытии соединения: {e}") + conn.close() + logger.debug("Подключение к БД закрыто") + def test_connection() -> bool: try: with get_connection() as conn: cur = PSQLCursor(conn) - try: - cur.execute("SELECT version();") - version = cur.fetchone() - logger.debug(f"Версия PostgreSQL: {version[0]}") - return True - finally: - cur.close() - logger.debug("Курсор закрыт") + cur.execute("SELECT version();") + version = cur.fetchone() + logger.info(f"Подключение к БД успешно: {version[0]}") + cur.close() + logger.debug("Курсор БД закрыт") + return True except Exception as e: - logger.error(f"Тест подключения к БД провален: {e}") + logger.error(f"Ошибка подключения к БД: {e}") return False -print(test_connection()) \ No newline at end of file + + +if __name__ == "__main__": + test_connection() \ No newline at end of file diff --git a/db/init.py b/db/init.py new file mode 100644 index 0000000..a7f6556 --- /dev/null +++ b/db/init.py @@ -0,0 +1,147 @@ +from db.connection import get_connection +from loguru import logger + +def create_tables(): + try: + with get_connection() as conn: + with conn.cursor() as cur: + # Пользователи + cur.execute(""" + CREATE TABLE IF NOT EXISTS users ( + id SERIAL PRIMARY KEY, + email VARCHAR(255) UNIQUE NOT NULL, + password_hash VARCHAR(255) NOT NULL, + name VARCHAR(255) NOT NULL, + role VARCHAR(50) NOT NULL DEFAULT 'viewer', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) + + # Роботы + cur.execute(""" + CREATE TABLE IF NOT EXISTS robots ( + id VARCHAR(50) PRIMARY KEY, + status VARCHAR(50) DEFAULT 'active', + battery_level INTEGER DEFAULT 100, + last_update TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + current_zone VARCHAR(10), + current_row INTEGER, + current_shelf INTEGER + ) + """) + + # Товары + cur.execute(""" + CREATE TABLE IF NOT EXISTS products ( + id VARCHAR(50) PRIMARY KEY, + name VARCHAR(255) NOT NULL, + category VARCHAR(100), + min_stock INTEGER DEFAULT 10, + optimal_stock INTEGER DEFAULT 100 + ) + """) + + # История инвентаризации + cur.execute(""" + CREATE TABLE IF NOT EXISTS inventory_history ( + id SERIAL PRIMARY KEY, + robot_id VARCHAR(50) REFERENCES robots(id), + product_id VARCHAR(50) REFERENCES products(id), + quantity INTEGER NOT NULL, + zone VARCHAR(10) NOT NULL, + row_number INTEGER, + shelf_number INTEGER, + status VARCHAR(50), + scanned_at TIMESTAMP NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) + + # Прогнозы ИИ + cur.execute(""" + CREATE TABLE IF NOT EXISTS ai_predictions ( + id SERIAL PRIMARY KEY, + product_id VARCHAR(50) REFERENCES products(id), + prediction_date DATE NOT NULL, + days_until_stockout INTEGER, + recommended_order INTEGER, + confidence_score DECIMAL(3,2), + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) + + conn.commit() + logger.debug("Все таблицы успешно созданы") + + except Exception as e: + logger.error(f"Ошибка создания таблиц: {e}") + raise + + +def create_indexes(): + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute("CREATE INDEX IF NOT EXISTS idx_inventory_scanned ON inventory_history(scanned_at DESC)") + cur.execute("CREATE INDEX IF NOT EXISTS idx_inventory_product ON inventory_history(product_id)") + cur.execute("CREATE INDEX IF NOT EXISTS idx_inventory_zone ON inventory_history(zone)") + + conn.commit() + logger.debug("Индексы созданы") + + except Exception as e: + logger.error(f"Ошибка создания индексов: {e}") + raise + + +def insert_sample_data(): + try: + with get_connection() as conn: + with conn.cursor() as cur: + cur.execute(""" + INSERT INTO users (email, password_hash, name, role) + VALUES + ('admin@warehouse.com', 'hash1', 'Администратор', 'admin'), + ('operator@warehouse.com', 'hash2', 'Оператор Иванов', 'operator'), + ('viewer@warehouse.com', 'hash3', 'Наблюдатель Петров', 'viewer') + ON CONFLICT (email) DO NOTHING + """) + + cur.execute(""" + INSERT INTO robots (id, status, battery_level, current_zone) + VALUES + ('RB-001', 'active', 85, 'A'), + ('RB-002', 'active', 45, 'B'), + ('RB-003', 'maintenance', 100, NULL) + ON CONFLICT (id) DO NOTHING + """) + + cur.execute(""" + INSERT INTO products (id, name, category, min_stock, optimal_stock) + VALUES + ('TEL-1234', 'Смартфон X', 'Электроника', 5, 50), + ('NOTE-567', 'Ноутбук Pro', 'Электроника', 3, 20), + ('ACC-999', 'Чехол для телефона', 'Аксессуары', 10, 100) + ON CONFLICT (id) DO NOTHING + """) + + conn.commit() + logger.debug("Тестовые данные добавлены") + + except Exception as e: + logger.error(f"Ошибка добавления тестовых данных: {e}") + raise + + +def initialize_database(): + logger.info("Начинаем инициализацию базы данных...") + + create_tables() + create_indexes() + insert_sample_data() + + logger.debug("База данных успешно инициализирована!") + + +if __name__ == "__main__": + initialize_database() \ No newline at end of file