feat(db): Database initialization has been added
This commit is contained in:
Kita Trofimov
+2
-1
@@ -1,5 +1,6 @@
|
||||
from flask import Blueprint, request
|
||||
from model.user import user
|
||||
|
||||
from model import user
|
||||
|
||||
loginBP = Blueprint("loginapi", __name__)
|
||||
|
||||
|
||||
+31
-27
@@ -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())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_connection()
|
||||
+147
@@ -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()
|
||||
Reference in New Issue
Block a user