generate_data.py - master - sonne/Task-RatsAD-31.group

generate_data.py
101 строка4 KB
arseon21
Add SQL
24 июн 2025, 20:24
24 июн 2025, 20:244416e94
import psycopg2
import random
import string
from faker import Faker
import uuid

# НАСТРОЙКИ ПОДКЛЮЧЕНИЯ К БАЗЕ ДАННЫХ

DB_PARAMS = {
    'dbname': 'fuzzy_search_lab',  # Имя базы данных
    'user': 'postgres',           # Имя пользователя
    'password': '1963',             # Ваш пароль PostgreSQL
    'host': 'localhost',
    'port': '5432'
}

fake = Faker()

categories = ['Электроника', 'Одежда', 'Продукты', 'Книги', 'Спорт', 'Дом', 'Игрушки']
brands = ['ТехноКорп', 'СтильБренд', 'ЕдаМастер', 'ЧитайБольше', 'СпортПро', 'ДомКомфорт', 'ИгрушкиМир']
test_keywords = ['computer', 'monitor', 'keyboard', 'software', 'processor', 'adapter', 'mouse', 'windows']


def introduce_typo(text: str) -> str:
    if not text or len(text) < 2:
        return text

    typo_type = random.choice(['swap', 'delete', 'insert', 'replace'])
    pos = random.randint(0, len(text) - 1)

    if typo_type == 'swap':
        pos = random.randint(0, len(text) - 2)
        return text[:pos] + text[pos + 1] + text[pos] + text[pos + 2:]
    elif typo_type == 'delete':
        return text[:pos] + text[pos + 1:]
    elif typo_type == 'insert':
        random_char = random.choice(string.ascii_lowercase)
        return text[:pos] + random_char + text[pos:]
    elif typo_type == 'replace':
        random_char = random.choice(string.ascii_lowercase)
        return text[:pos] + random_char + text[pos + 1:]
    return text


def generate_product_data(count: int) -> list[dict]:
    products = []
    num_keywords = len(test_keywords)

    for i in range(count):
        keyword = test_keywords[i % num_keywords]

        name = f"{random.choice(brands)} {keyword.capitalize()} Model {fake.bban()[:4]}"

        if random.random() < 0.15:
            keyword_with_typo = introduce_typo(keyword)
            name = name.lower().replace(keyword, keyword_with_typo)

        products.append({
            'name': name,
            'description': f"Description for {name}",
            'category': random.choice(categories),
            'brand': random.choice(brands),
            'sku': f"SKU-{uuid.uuid4().hex[:8].upper()}"
        })
    return products


def insert_data_to_db(connection, products: list[dict]):
    with connection.cursor() as cursor:
        from psycopg2.extras import execute_values
        sql_insert_query = """INSERT INTO products (name, description, category, brand, sku) \
                              VALUES %s"""
        values = [(p['name'], p['description'], p['category'], p['brand'], p['sku']) for p in products]
        execute_values(cursor, sql_insert_query, values)
        connection.commit()
        print(f"Успешно вставлено {cursor.rowcount} записей.")


if __name__ == "__main__":
    DATASET_SIZE = 1000

    try:
        conn = psycopg2.connect(**DB_PARAMS)

        with conn.cursor() as cur:
            print("Очистка таблицы 'products'...")
            cur.execute("TRUNCATE TABLE products RESTART IDENTITY CASCADE;")
            conn.commit()

        print(f"Генерация {DATASET_SIZE} записей...")
        product_list = generate_product_data(DATASET_SIZE)

        print("Вставка данных в таблицу 'products'...")
        insert_data_to_db(conn, product_list)

    except Exception as e:
        print(f"Произошла ошибка: {e}")
    finally:
        if 'conn' in locals() and conn is not None:
            conn.close()
            print("Соединение с БД закрыто.")