database.py - master - evgeny.kumpyakov/restman

database.py
376 строк16 KB
EvgenyKumpyakov
Добавлены лицензии и общее описание проекта
24 июн 2026, 14:44
24 июн 2026, 14:4490478e3
#!/usr/bin/env python3
"""
Модуль для работы с базой данных SQLite

Copyright (C) 2026 EvgenyKumpyakov

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import sqlite3
import json
from typing import List, Dict, Optional


class Database:
    """Класс для работы с базой данных SQLite"""
    
    def __init__(self, db_path: str):
        """
        Инициализирует подключение к базе данных
        
        Args:
            db_path: Путь к файлу базы данных
        """
        self.db_path = db_path
        self.conn = None
        self.init_database()
    
    def connect(self):
        """Устанавливает подключение к базе данных"""
        self.conn = sqlite3.connect(self.db_path)
        self.conn.row_factory = sqlite3.Row
    
    def close(self):
        """Закрывает подключение к базе данных"""
        if self.conn:
            self.conn.close()
            self.conn = None
    
    def init_database(self):
        """Инициализирует структуру базы данных"""
        self.connect()
        cursor = self.conn.cursor()
        
        # Таблица групп запросов
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS request_groups (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL UNIQUE
            )
        """)
        
        # Таблица запросов
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS requests (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                group_id INTEGER NOT NULL,
                name TEXT NOT NULL,
                method TEXT NOT NULL,
                url TEXT NOT NULL,
                headers TEXT,
                body TEXT,
                params TEXT,
                username TEXT,
                password TEXT,
                auth_type TEXT,
                auth_enabled INTEGER,
                bearer_token TEXT,
                bearer_prefix TEXT,
                basic_username TEXT,
                basic_password TEXT,
                FOREIGN KEY (group_id) REFERENCES request_groups(id) ON DELETE CASCADE
            )
        """)
        
        # Миграция: добавляем колонки, если их еще нет
        cursor.execute("PRAGMA table_info(requests)")
        columns = [row[1] for row in cursor.fetchall()]
        
        if 'params' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN params TEXT")
        if 'username' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN username TEXT")
        if 'password' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN password TEXT")
        if 'auth_type' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN auth_type TEXT")
        if 'auth_enabled' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN auth_enabled INTEGER")
        if 'bearer_token' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN bearer_token TEXT")
        if 'bearer_prefix' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN bearer_prefix TEXT")
        if 'basic_username' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN basic_username TEXT")
        if 'basic_password' not in columns:
            cursor.execute("ALTER TABLE requests ADD COLUMN basic_password TEXT")
        
        self.conn.commit()
        self.close()
    
    def save_groups_and_requests(self, groups_data: List[Dict]) -> bool:
        """
        Сохраняет группы запросов и сами запросы в базу данных
        
        Args:
            groups_data: Список словарей с данными групп и запросов
                Формат: [{
                    'name': 'Название группы',
                    'requests': [{
                        'name': 'Название запроса',
                        'method': 'GET',
                        'url': 'https://...',
                        'headers': '{}',
                        'body': ''
                    }]
                }]
        
        Returns:
            True если сохранение прошло успешно, False в противном случае
        """
        try:
            self.connect()
            cursor = self.conn.cursor()
            
            # Очищаем существующие данные
            cursor.execute("DELETE FROM requests")
            cursor.execute("DELETE FROM request_groups")
            
            # Сохраняем группы и запросы
            for group_data in groups_data:
                # Вставляем группу
                cursor.execute(
                    "INSERT INTO request_groups (name) VALUES (?)",
                    (group_data['name'],)
                )
                group_id = cursor.lastrowid
                
                # Вставляем запросы группы
                for request_data in group_data.get('requests', []):
                    cursor.execute(
                        """INSERT INTO requests 
                           (group_id, name, method, url, headers, body, params, username, password,
                            auth_type, auth_enabled, bearer_token, bearer_prefix, 
                            basic_username, basic_password) 
                           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                        (
                            group_id,
                            request_data['name'],
                            request_data['method'],
                            request_data['url'],
                            request_data.get('headers', '{}'),
                            request_data.get('body', ''),
                            request_data.get('params', '{}'),
                            request_data.get('username', ''),
                            request_data.get('password', ''),
                            request_data.get('auth_type', 'Bearer Token'),
                            1 if request_data.get('auth_enabled', False) else 0,
                            request_data.get('bearer_token', ''),
                            request_data.get('bearer_prefix', 'Bearer'),
                            request_data.get('basic_username', ''),
                            request_data.get('basic_password', '')
                        )
                    )
            
            self.conn.commit()
            self.close()
            return True
            
        except Exception as e:
            print(f"Ошибка при сохранении данных: {e}")
            if self.conn:
                self.conn.rollback()
                self.close()
            return False
    
    def load_groups_and_requests(self) -> List[Dict]:
        """
        Загружает группы запросов и сами запросы из базы данных
        
        Returns:
            Список словарей с данными групп и запросов
        """
        try:
            self.connect()
            cursor = self.conn.cursor()
            
            # Проверяем, какие колонки есть в таблице requests
            cursor.execute("PRAGMA table_info(requests)")
            columns_info = cursor.fetchall()
            available_columns = [col[1] for col in columns_info]
            
            # Формируем список колонок для SELECT
            base_columns = ['name', 'method', 'url', 'headers', 'body']
            optional_columns = {
                'params': 'params',
                'username': 'username',
                'password': 'password',
                'auth_type': 'auth_type',
                'auth_enabled': 'auth_enabled',
                'bearer_token': 'bearer_token',
                'bearer_prefix': 'bearer_prefix',
                'basic_username': 'basic_username',
                'basic_password': 'basic_password'
            }
            
            # Добавляем только существующие колонки
            select_columns = base_columns.copy()
            for col_name, col_alias in optional_columns.items():
                if col_name in available_columns:
                    select_columns.append(col_name)
            
            # Получаем все группы
            cursor.execute("SELECT id, name FROM request_groups ORDER BY name")
            groups = cursor.fetchall()
            
            result = []
            for group in groups:
                group_id = group['id']
                group_name = group['name']
                
                # Формируем и выполняем запрос с доступными колонками
                columns_str = ', '.join(select_columns)
                query = f"""SELECT {columns_str}
                           FROM requests 
                           WHERE group_id = ? 
                           ORDER BY name"""
                
                cursor.execute(query, (group_id,))
                requests = cursor.fetchall()
                
                requests_data = []
                for req in requests:
                    # Поддержка обратной совместимости
                    # sqlite3.Row не имеет метода get(), используем индексацию с проверкой
                    auth_type = None
                    if 'auth_type' in available_columns:
                        try:
                            auth_type = req['auth_type']
                        except (KeyError, IndexError):
                            auth_type = None
                    
                    if not auth_type or auth_type == '':
                        # Если есть старые данные username/password, используем Basic
                        username = ''
                        password = ''
                        if 'username' in available_columns:
                            try:
                                username = req['username'] or ''
                            except (KeyError, IndexError):
                                username = ''
                        if 'password' in available_columns:
                            try:
                                password = req['password'] or ''
                            except (KeyError, IndexError):
                                password = ''
                        
                        if username or password:
                            auth_type = 'Basic'
                        else:
                            auth_type = 'Bearer Token'
                    
                    # Получаем auth_enabled
                    auth_enabled = False
                    if 'auth_enabled' in available_columns:
                        try:
                            auth_enabled_val = req['auth_enabled']
                            auth_enabled = bool(auth_enabled_val) if auth_enabled_val is not None else False
                        except (KeyError, IndexError):
                            auth_enabled = False
                    
                    # Получаем данные для Bearer Token
                    bearer_token = ''
                    bearer_prefix = 'Bearer'
                    if 'bearer_token' in available_columns:
                        try:
                            bearer_token = req['bearer_token'] or ''
                        except (KeyError, IndexError):
                            bearer_token = ''
                    if 'bearer_prefix' in available_columns:
                        try:
                            bearer_prefix = req['bearer_prefix'] or 'Bearer'
                        except (KeyError, IndexError):
                            bearer_prefix = 'Bearer'
                    
                    # Получаем данные для Basic
                    basic_username = ''
                    basic_password = ''
                    if 'basic_username' in available_columns:
                        try:
                            basic_username = req['basic_username'] or ''
                        except (KeyError, IndexError):
                            basic_username = ''
                    if 'basic_password' in available_columns:
                        try:
                            basic_password = req['basic_password'] or ''
                        except (KeyError, IndexError):
                            basic_password = ''
                    
                    # Если нет новых полей, используем старые
                    if not basic_username and 'username' in available_columns:
                        try:
                            basic_username = req['username'] or ''
                        except (KeyError, IndexError):
                            basic_username = ''
                    if not basic_password and 'password' in available_columns:
                        try:
                            basic_password = req['password'] or ''
                        except (KeyError, IndexError):
                            basic_password = ''
                    
                    # Получаем базовые поля
                    username_val = ''
                    password_val = ''
                    if 'username' in available_columns:
                        try:
                            username_val = req['username'] or ''
                        except (KeyError, IndexError):
                            username_val = ''
                    if 'password' in available_columns:
                        try:
                            password_val = req['password'] or ''
                        except (KeyError, IndexError):
                            password_val = ''
                    
                    # Получаем параметры
                    params_val = '{}'
                    if 'params' in available_columns:
                        try:
                            params_val = req['params'] or '{}'
                        except (KeyError, IndexError):
                            params_val = '{}'
                    
                    requests_data.append({
                        'name': req['name'],
                        'method': req['method'],
                        'url': req['url'],
                        'headers': req['headers'] or '{}',
                        'body': req['body'] or '',
                        'params': params_val,
                        'username': username_val,
                        'password': password_val,
                        'auth_type': auth_type,
                        'auth_enabled': auth_enabled,
                        'bearer_token': bearer_token,
                        'bearer_prefix': bearer_prefix,
                        'basic_username': basic_username,
                        'basic_password': basic_password
                    })
                
                result.append({
                    'name': group_name,
                    'requests': requests_data
                })
            
            self.close()
            return result
            
        except Exception as e:
            print(f"Ошибка при загрузке данных: {e}")
            import traceback
            traceback.print_exc()
            if self.conn:
                self.close()
            return []