270 lines
11 KiB
Python
270 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Автономный скрипт для инициализации тарифных планов в PostgreSQL.
|
||
Использует данные подключения из docker-compose.yml
|
||
"""
|
||
|
||
import asyncio
|
||
import argparse
|
||
import sys
|
||
from typing import List, Dict
|
||
from sqlalchemy import Column, Integer, String, Numeric, Text, delete, insert
|
||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||
from sqlalchemy.orm import declarative_base
|
||
from decimal import Decimal
|
||
|
||
Base = declarative_base()
|
||
|
||
|
||
class Plan(Base):
|
||
"""Модель тарифного плана"""
|
||
__tablename__ = 'plans'
|
||
|
||
id = Column(Integer, primary_key=True)
|
||
name = Column(String(100), nullable=False)
|
||
price = Column(Numeric(10, 2), nullable=False)
|
||
duration_days = Column(Integer, nullable=False)
|
||
description = Column(Text, nullable=True)
|
||
|
||
|
||
# Данные из вашего docker-compose.yml
|
||
DEFAULT_CONFIG = {
|
||
'host': 'localhost',
|
||
'port': 5432,
|
||
'database': 'postgresql',
|
||
'user': 'AH3J9GSPBYOP',
|
||
'password': 'uPS9?y~mcu2',
|
||
'url': 'postgresql+asyncpg://AH3J9GSPBYOP:uPS9?y~mcu2@localhost:5432/postgresql'
|
||
}
|
||
|
||
|
||
PLANS_DATA = [
|
||
{'name': 'Lark_Standart_1', 'price': Decimal('200.00'), 'duration_days': 30},
|
||
{'name': 'Lark_Pro_1', 'price': Decimal('400.00'), 'duration_days': 30},
|
||
{'name': 'Lark_Family_1', 'price': Decimal('700.00'), 'duration_days': 30},
|
||
{'name': 'Lark_Standart_6', 'price': Decimal('1200.00'), 'duration_days': 180},
|
||
{'name': 'Lark_Standart_12', 'price': Decimal('2400.00'), 'duration_days': 360},
|
||
{'name': 'Lark_Pro_6', 'price': Decimal('2000.00'), 'duration_days': 180},
|
||
{'name': 'Lark_Pro_12', 'price': Decimal('4800.00'), 'duration_days': 360},
|
||
{'name': 'Lark_Family_6', 'price': Decimal('4200.00'), 'duration_days': 180},
|
||
{'name': 'Lark_Family_12', 'price': Decimal('8400.00'), 'duration_days': 360},
|
||
]
|
||
|
||
|
||
def print_banner():
|
||
"""Печатает баннер скрипта"""
|
||
print("=" * 60)
|
||
print("🚀 ИНИЦИАЛИЗАЦИЯ ТАРИФНЫХ ПЛАНОВ В БАЗЕ ДАННЫХ")
|
||
print("=" * 60)
|
||
print()
|
||
|
||
|
||
def create_db_url(config: dict) -> str:
|
||
"""Создает URL для подключения к базе данных"""
|
||
if config.get('url'):
|
||
return config['url']
|
||
|
||
return f"postgresql+asyncpg://{config['user']}:{config['password']}@{config['host']}:{config['port']}/{config['database']}"
|
||
|
||
|
||
async def check_connection(engine) -> bool:
|
||
"""Проверяет подключение к базе данных"""
|
||
try:
|
||
async with engine.connect() as conn:
|
||
result = await conn.execute("SELECT version()")
|
||
version = result.scalar()
|
||
print(f"✅ Подключено к PostgreSQL: {version.split(',')[0]}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ Ошибка подключения к базе данных: {e}")
|
||
return False
|
||
|
||
|
||
async def get_existing_plans(session) -> List:
|
||
"""Получает существующие тарифные планы"""
|
||
result = await session.execute(
|
||
"SELECT id, name, price, duration_days FROM plans ORDER BY price"
|
||
)
|
||
return result.fetchall()
|
||
|
||
|
||
async def clear_table(session, table_name: str = 'plans') -> bool:
|
||
"""Очищает указанную таблицу"""
|
||
try:
|
||
await session.execute(delete(Plan))
|
||
await session.commit()
|
||
print(f"✅ Таблица '{table_name}' очищена")
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ Ошибка при очистке таблицы: {e}")
|
||
await session.rollback()
|
||
return False
|
||
|
||
|
||
async def add_plans_to_db(session, plans_data: List[Dict]) -> int:
|
||
"""Добавляет тарифные планы в базу данных"""
|
||
try:
|
||
added_count = 0
|
||
for plan in plans_data:
|
||
await session.execute(
|
||
insert(Plan).values(**plan)
|
||
)
|
||
added_count += 1
|
||
|
||
await session.commit()
|
||
return added_count
|
||
except Exception as e:
|
||
await session.rollback()
|
||
raise e
|
||
|
||
|
||
async def print_plans_table(plans: List) -> None:
|
||
"""Выводит таблицу с тарифными планами"""
|
||
if not plans:
|
||
print("📭 Таблица 'plans' пуста")
|
||
return
|
||
|
||
print(f"\n📊 Текущие тарифные планы ({len(plans)} шт.):")
|
||
print("-" * 70)
|
||
print(f"{'ID':<5} {'Название':<25} {'Цена (руб.)':<15} {'Дней':<10}")
|
||
print("-" * 70)
|
||
|
||
for plan in plans:
|
||
print(f"{plan[0]:<5} {plan[1]:<25} {plan[2]:<15} {plan[3]:<10}")
|
||
|
||
print("-" * 70)
|
||
|
||
# Подсчет статистики
|
||
total_price = sum(float(p[2]) for p in plans)
|
||
avg_price = total_price / len(plans) if plans else 0
|
||
|
||
print(f"💰 Общая сумма всех тарифов: {total_price:.2f} руб.")
|
||
print(f"📈 Средняя цена тарифа: {avg_price:.2f} руб.")
|
||
print(f"📅 Всего предложений: {len(plans)}")
|
||
|
||
|
||
async def main(config: dict, clear_existing: bool = True, dry_run: bool = False):
|
||
"""Основная функция скрипта"""
|
||
|
||
print_banner()
|
||
|
||
# Создаем URL для подключения
|
||
db_url = create_db_url(config)
|
||
print(f"📡 Параметры подключения:")
|
||
print(f" Хост: {config['host']}:{config['port']}")
|
||
print(f" База данных: {config['database']}")
|
||
print(f" Пользователь: {config['user']}")
|
||
print(f" {'🚨 РЕЖИМ ТЕСТА (dry-run)' if dry_run else ''}")
|
||
print()
|
||
|
||
try:
|
||
# Подключаемся к базе данных
|
||
print("🔄 Подключение к базе данных...")
|
||
engine = create_async_engine(db_url, echo=False)
|
||
|
||
# Проверяем подключение
|
||
if not await check_connection(engine):
|
||
print("\n❌ Не удалось подключиться к базе данных")
|
||
return False
|
||
|
||
# Создаем фабрику сессий
|
||
AsyncSessionLocal = async_sessionmaker(
|
||
engine, class_=AsyncSession, expire_on_commit=False
|
||
)
|
||
|
||
async with AsyncSessionLocal() as session:
|
||
# Получаем текущие тарифы
|
||
print("\n🔍 Проверяем существующие тарифы...")
|
||
existing_plans = await get_existing_plans(session)
|
||
|
||
if existing_plans:
|
||
await print_plans_table(existing_plans)
|
||
|
||
if clear_existing and not dry_run:
|
||
print("\n⚠️ ВНИМАНИЕ: Будут удалены все существующие тарифы!")
|
||
confirm = input("Продолжить? (y/N): ")
|
||
if confirm.lower() != 'y':
|
||
print("❌ Операция отменена пользователем")
|
||
return False
|
||
|
||
# Очищаем таблицу
|
||
await clear_table(session)
|
||
elif dry_run:
|
||
print("\n⚠️ DRY-RUN: Существующие тарифы НЕ будут удалены")
|
||
else:
|
||
print("📭 Таблица 'plans' пуста, создаем новые тарифы...")
|
||
|
||
# Добавляем новые тарифы
|
||
if not dry_run:
|
||
print(f"\n➕ Добавляем {len(PLANS_DATA)} тарифных планов...")
|
||
added_count = await add_plans_to_db(session, PLANS_DATA)
|
||
print(f"✅ Успешно добавлено {added_count} тарифов")
|
||
else:
|
||
print(f"\n⚠️ DRY-RUN: Планируется добавить {len(PLANS_DATA)} тарифов:")
|
||
for i, plan in enumerate(PLANS_DATA, 1):
|
||
print(f" {i}. {plan['name']} - {plan['price']} руб. ({plan['duration_days']} дней)")
|
||
|
||
# Показываем финальный результат
|
||
print("\n🎯 ФИНАЛЬНЫЙ РЕЗУЛЬТАТ:")
|
||
final_plans = await get_existing_plans(session)
|
||
await print_plans_table(final_plans)
|
||
|
||
await engine.dispose()
|
||
print("\n✅ Скрипт успешно выполнен!")
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ Критическая ошибка: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return False
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(
|
||
description='Инициализация тарифных планов в базе данных',
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
Примеры использования:
|
||
%(prog)s # Использует настройки по умолчанию
|
||
%(prog)s --no-clear # Не очищать существующие тарифы
|
||
%(prog)s --dry-run # Только показать что будет сделано
|
||
%(prog)s --host 192.168.1.100 # Указать другой хост
|
||
%(prog)s --url "postgresql://..." # Указать полный URL
|
||
"""
|
||
)
|
||
|
||
parser.add_argument('--host', help='Хост базы данных', default=DEFAULT_CONFIG['host'])
|
||
parser.add_argument('--port', type=int, help='Порт базы данных', default=DEFAULT_CONFIG['port'])
|
||
parser.add_argument('--database', help='Имя базы данных', default=DEFAULT_CONFIG['database'])
|
||
parser.add_argument('--user', help='Имя пользователя', default=DEFAULT_CONFIG['user'])
|
||
parser.add_argument('--password', help='Пароль', default=DEFAULT_CONFIG['password'])
|
||
parser.add_argument('--url', help='Полный URL подключения (игнорирует остальные параметры)')
|
||
parser.add_argument('--no-clear', action='store_true', help='Не очищать существующие тарифы')
|
||
parser.add_argument('--dry-run', action='store_true', help='Только показать что будет сделано')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Формируем конфигурацию
|
||
config = DEFAULT_CONFIG.copy()
|
||
|
||
if args.url:
|
||
config['url'] = args.url
|
||
else:
|
||
config.update({
|
||
'host': args.host,
|
||
'port': args.port,
|
||
'database': args.database,
|
||
'user': args.user,
|
||
'password': args.password,
|
||
'url': None # Будет сгенерирован автоматически
|
||
})
|
||
|
||
# Запускаем скрипт
|
||
success = asyncio.run(main(
|
||
config=config,
|
||
clear_existing=not args.no_clear,
|
||
dry_run=args.dry_run
|
||
))
|
||
|
||
sys.exit(0 if success else 1) |