diff --git a/work_parser.py b/work_parser.py index 5538e76..9f96a0c 100644 --- a/work_parser.py +++ b/work_parser.py @@ -1,17 +1,26 @@ import psycopg2 from psycopg2.extras import RealDictCursor -# Подключение к БД (укажи свои параметры) -conn = psycopg2.connect( - dbname="parsed_url", - user="postgres", - password="qwertyqwerty123123", - host="45.129.78.228", - # host="127.0.0.1" - connect_timeout=10, - options="-c statement_timeout=30000" # таймаут запроса 30 сек -) -conn.autocommit = True +# Параметры подключения к БД +DB_CONFIG = { + "dbname": "parsed_url", + "user": "postgres", + "password": "qwertyqwerty123123", + "host": "45.129.78.228", + "connect_timeout": 10, + "options": "-c statement_timeout=30000" +} + +# Глобальное подключение к БД +conn = None + +def get_connection(): + """Получает подключение к БД, создавая новое при необходимости""" + global conn + if conn is None or conn.closed: + conn = psycopg2.connect(**DB_CONFIG) + conn.autocommit = True + return conn def close_connection(): """Закрывает подключение к БД""" @@ -21,6 +30,7 @@ def close_connection(): conn = None # работа с базой данных показывания задач work_parser def create_table(): + conn = get_connection() try: with conn.cursor() as cur: cur.execute(""" @@ -38,10 +48,10 @@ def create_table(): """) print("Таблица work_parser создана или уже существует") finally: - if conn: - conn.close() + pass # Не закрываем подключение def insert_task(status, source_url=None, source_id=None, priority=0): + conn = get_connection() try: with conn.cursor() as cur: cur.execute(""" @@ -52,10 +62,10 @@ def insert_task(status, source_url=None, source_id=None, priority=0): task_id = cur.fetchone()[0] return task_id finally: - if conn: - conn.close() + pass # Не закрываем подключение def get_tasks_offset(limit, offset): + conn = get_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute(""" @@ -66,10 +76,10 @@ def get_tasks_offset(limit, offset): tasks = cur.fetchall() return tasks finally: - if conn: - conn.close() + pass # Не закрываем подключение def delete_task(task_id: int): + conn = get_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute("DELETE FROM work_parser WHERE id = %s RETURNING *;", (task_id,)) @@ -79,10 +89,10 @@ def delete_task(task_id: int): else: return {"message": f"Задача с id {task_id} не найдена"} finally: - if conn: - conn.close() + pass # Не закрываем подключение def update_task(task_id, **fields): + conn = get_connection() try: allowed_fields = ['status', 'started_at', 'finished_at', 'source_url', 'error_message', 'attempts', 'priority'] set_parts = [] @@ -99,11 +109,11 @@ def update_task(task_id, **fields): cur.execute(f"UPDATE work_parser SET {set_sql} WHERE id = %s;", values) return True finally: - if conn: - conn.close() + pass # Не закрываем подключение # Создание и работа с таблицей по созданию и редактированию промтов def create_table_config_gpt(): + conn = get_connection() try: with conn.cursor() as cur: cur.execute(""" @@ -115,10 +125,10 @@ def create_table_config_gpt(): """) print("Таблица config_gpt создана или уже существует") finally: - if conn: - conn.close() + pass # Не закрываем подключение def update_promt(url: str, name: str, promt: str): + conn = get_connection() try: with conn.cursor() as cur: cur.execute(""" @@ -130,21 +140,21 @@ def update_promt(url: str, name: str, promt: str): """, (url, name, promt)) conn.commit() finally: - if conn: - conn.close() + pass # Не закрываем подключение def get_promt(promt_name_url): + conn = get_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute("SELECT promt FROM config_gpt WHERE url = %s", (promt_name_url,)) promt = cur.fetchone() return promt['promt'] finally: - if conn: - conn.close() + pass # Не закрываем подключение def get_all_promt(): # Возвращает список всех значений поля name из таблицы config_gpt + conn = get_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute("SELECT * FROM config_gpt") @@ -153,10 +163,10 @@ def get_all_promt(): sources = [{"url": row["url"], "name": row["name"], "promt": row["promt"]} for row in rows] return {"sources": sources} finally: - if conn: - conn.close() + pass # Не закрываем подключение def get_all_categories_promt(): + conn = get_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute("SELECT name FROM config_gpt") @@ -164,11 +174,11 @@ def get_all_categories_promt(): return [row["name"] for row in rows] finally: - if conn: - conn.close() + pass # Не закрываем подключение # Создание, сохранение и работа с таблицей ошибочных ссылок (error_url) def create_table_error_url(): + conn = get_connection() try: with conn.cursor() as cur: cur.execute(""" @@ -180,10 +190,10 @@ def create_table_error_url(): """) print("Таблица error_url создана или уже существует") finally: - if conn: - conn.close() + pass # Не закрываем подключение def add_error_url(source_url: str, error_sources_url: str): + conn = get_connection() try: with conn.cursor() as cur: cur.execute(""" @@ -193,10 +203,10 @@ def add_error_url(source_url: str, error_sources_url: str): """, (source_url, error_sources_url)) return cur.fetchone()[0] finally: - if conn: - conn.close() + pass # Не закрываем подключение def check_error_url(error_sources_url: str) -> bool: + conn = get_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute(""" @@ -208,11 +218,11 @@ def check_error_url(error_sources_url: str) -> bool: return row is None finally: - if conn: - conn.close() + pass # Не закрываем подключение # Создание и работа с таблицей источников add_sources def create_table_add_sourse(): + conn = get_connection() try: with conn.cursor() as cur: cur.execute(""" @@ -223,10 +233,10 @@ def create_table_add_sourse(): """) print("Таблица sourse создана или уже существует") finally: - if conn: - conn.close() + pass # Не закрываем подключение def add_sources(url: str, promt: str): + conn = get_connection() try: with conn.cursor() as cur: cur.execute(""" @@ -237,8 +247,7 @@ def add_sources(url: str, promt: str): """, (url, promt)) conn.commit() finally: - if conn: - conn.close() + pass # Не закрываем подключение # Пример использования # if __name__ == "__main__":