setup_database.py

import asyncio
import os
import aiomysql
from dotenv import load_dotenv

CREATE_TABLES_SQL = [
    """
    CREATE TABLE IF NOT EXISTS entities (
        id INT AUTO_INCREMENT PRIMARY KEY,
        name VARCHAR(255) UNIQUE NOT NULL,
        type VARCHAR(100) NOT NULL COMMENT 'e.g., shaxs, futbol_klubi, texnologiya',
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    ) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
    """,
    """
    CREATE TABLE IF NOT EXISTS facts (
        id INT AUTO_INCREMENT PRIMARY KEY,
        entity_id INT NOT NULL,
        attribute VARCHAR(255) NOT NULL,
        value TEXT NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
        UNIQUE KEY (entity_id, attribute),
        FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE CASCADE
    ) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
    """,
    """
    CREATE TABLE IF NOT EXISTS user_opinions (
        id INT AUTO_INCREMENT PRIMARY KEY,
        user_telegram_id BIGINT NOT NULL,
        entity_id INT NOT NULL,
        sentiment ENUM('ijobiy', 'salbiy', 'neytral') NOT NULL,
        strength FLOAT NOT NULL COMMENT 'Confidence or intensity from 0.0 to 1.0',
        source_message TEXT,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
        UNIQUE KEY (user_telegram_id, entity_id),
        FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE CASCADE
    ) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
    """,
    """
    CREATE TABLE IF NOT EXISTS conversation_history (
        id INT AUTO_INCREMENT PRIMARY KEY,
        chat_id BIGINT NOT NULL,
        user_id BIGINT NOT NULL,
        username VARCHAR(255) NOT NULL,
        message_text TEXT NOT NULL,
        timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        INDEX idx_chat_timestamp (chat_id, timestamp)
    ) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
    """
]

async def setup_database():
    """Connects to the database and creates/updates the necessary tables."""
    load_dotenv()
    conn = None
    try:
        conn = await aiomysql.connect(
            host=os.getenv("DB_HOST"),
            port=int(os.getenv("DB_PORT", 3306)),
            user=os.getenv("DB_USER"),
            password=os.getenv("DB_PASS"),
            db=os.getenv("DB_NAME"),
            charset='utf8mb4'
        )
        async with conn.cursor() as cur:
            print("INFO: Database connected. Ensuring all tables exist...")
            for statement in CREATE_TABLES_SQL:
                await cur.execute(statement)
            print("SUCCESS: All tables are created and ready.")
    except Exception as e:
        print(f"ERROR: Could not set up the database: {e}")
    finally:
        if conn:
            conn.close()

if __name__ == '__main__':
    print("Running database setup...")
    asyncio.run(setup_database())