db.py - master - Kirby/EnglishWork

db.py
131 строка3 KB
Alex
Исправлены замечания
13 май 2026, 20:53
13 май 2026, 20:534f2ad7d
import psycopg2
import random
import os
from dotenv import load_dotenv
load_dotenv()

DB_NAME = os.getenv("DB_NAME")
DB_USER = os.getenv("DB_USER")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_HOST = os.getenv("DB_HOST")
DB_PORT = os.getenv("DB_PORT")


def get_connection():
    return psycopg2.connect(
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD,
        host=DB_HOST,
        port=DB_PORT
    )


def get_or_create_user(username):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO users (username) VALUES (%s) ON CONFLICT (username) DO NOTHING",
        (username,)
    )
    conn.commit()
    cur.execute("SELECT id FROM users WHERE username = %s", (username,))
    user_id = cur.fetchone()[0]
    cur.close()
    conn.close()
    return user_id


def get_random_words(user_id, limit=4):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute("""
        SELECT id, russian, english FROM (
            SELECT id, russian, english FROM words
            UNION ALL
            SELECT id, russian, english FROM user_words WHERE user_id = %s
        ) AS all_words
        ORDER BY RANDOM()
        LIMIT %s
    """, (user_id, limit))
    result = cur.fetchall()
    cur.close()
    conn.close()
    return result


def get_user_words(user_id):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        "SELECT id, russian, english FROM user_words WHERE user_id = %s",
        (user_id,)
    )
    result = cur.fetchall()
    cur.close()
    conn.close()
    return result


def add_word(user_id, russian, english):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO user_words (user_id, russian, english) VALUES (%s, %s, %s)",
        (user_id, russian.strip(), english.strip())
    )
    conn.commit()
    cur.close()
    conn.close()


def delete_word(user_id, word_id):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute(
        "DELETE FROM user_words WHERE id = %s AND user_id = %s",
        (word_id, user_id)
    )
    conn.commit()
    cur.close()
    conn.close()


def get_choices(correct_answer, user_id):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute("""
        SELECT english FROM (
            SELECT english FROM words
            UNION ALL
            SELECT english FROM user_words WHERE user_id = %s
        ) AS all_english
        WHERE english != %s
        ORDER BY RANDOM()
        LIMIT 3
    """, (user_id, correct_answer))
    others = [row[0] for row in cur.fetchall()]
    cur.close()
    conn.close()

    while len(others) < 3:
        others.append(others[0] if others else "N/A")

    variants = others + [correct_answer]
    random.shuffle(variants)
    return variants


def get_word_count(user_id):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute("SELECT COUNT(*) FROM words")
    common_count = cur.fetchone()[0]
    cur.execute(
        "SELECT COUNT(*) FROM user_words WHERE user_id = %s",
        (user_id,)
    )
    personal_count = cur.fetchone()[0]
    cur.close()
    conn.close()
    return common_count + personal_count