добавил закрытие бд для пресечения утечки данных
All checks were successful
continuous-integration/drone/push Build is passing

This commit is contained in:
2026-04-11 12:14:01 +10:00
parent 8c790b018c
commit a813a3805e

View File

@@ -12,8 +12,16 @@ conn = psycopg2.connect(
options="-c statement_timeout=30000" # таймаут запроса 30 сек options="-c statement_timeout=30000" # таймаут запроса 30 сек
) )
conn.autocommit = True conn.autocommit = True
# работа с базой данных показывания задач work_parser
def close_connection():
"""Закрывает подключение к БД"""
global conn
if conn:
conn.close()
conn = None
def create_table(): def create_table():
try:
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
CREATE TABLE IF NOT EXISTS work_parser ( CREATE TABLE IF NOT EXISTS work_parser (
@@ -29,8 +37,12 @@ def create_table():
); );
""") """)
print("Таблица work_parser создана или уже существует") print("Таблица work_parser создана или уже существует")
finally:
if conn:
conn.close()
def insert_task(status, source_url=None, source_id=None, priority=0): def insert_task(status, source_url=None, source_id=None, priority=0):
try:
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
INSERT INTO work_parser (status, source_url, priority) INSERT INTO work_parser (status, source_url, priority)
@@ -39,8 +51,12 @@ def insert_task(status, source_url=None, source_id=None, priority=0):
""", (status, source_url, priority)) """, (status, source_url, priority))
task_id = cur.fetchone()[0] task_id = cur.fetchone()[0]
return task_id return task_id
finally:
if conn:
conn.close()
def get_tasks_offset(limit, offset): def get_tasks_offset(limit, offset):
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur: with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(""" cur.execute("""
SELECT * FROM work_parser SELECT * FROM work_parser
@@ -49,8 +65,12 @@ def get_tasks_offset(limit, offset):
""", (limit, offset)) """, (limit, offset))
tasks = cur.fetchall() tasks = cur.fetchall()
return tasks return tasks
finally:
if conn:
conn.close()
def delete_task(task_id: int): def delete_task(task_id: int):
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur: with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute("DELETE FROM work_parser WHERE id = %s RETURNING *;", (task_id,)) cur.execute("DELETE FROM work_parser WHERE id = %s RETURNING *;", (task_id,))
deleted_task = cur.fetchone() deleted_task = cur.fetchone()
@@ -58,9 +78,12 @@ def delete_task(task_id: int):
return {"message": f"Задача {task_id} удалена", "deleted_task": dict(deleted_task)} return {"message": f"Задача {task_id} удалена", "deleted_task": dict(deleted_task)}
else: else:
return {"message": f"Задача с id {task_id} не найдена"} return {"message": f"Задача с id {task_id} не найдена"}
finally:
if conn:
conn.close()
def update_task(task_id, **fields): def update_task(task_id, **fields):
# dynamic update query generator try:
allowed_fields = ['status', 'started_at', 'finished_at', 'source_url', 'error_message', 'attempts', 'priority'] allowed_fields = ['status', 'started_at', 'finished_at', 'source_url', 'error_message', 'attempts', 'priority']
set_parts = [] set_parts = []
values = [] values = []
@@ -75,9 +98,12 @@ def update_task(task_id, **fields):
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(f"UPDATE work_parser SET {set_sql} WHERE id = %s;", values) cur.execute(f"UPDATE work_parser SET {set_sql} WHERE id = %s;", values)
return True return True
finally:
if conn:
conn.close()
# Создание и работа с таблицей по созданию и редактированию промтов
def create_table_config_gpt(): def create_table_config_gpt():
try:
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
CREATE TABLE IF NOT EXISTS config_gpt ( CREATE TABLE IF NOT EXISTS config_gpt (
@@ -87,8 +113,12 @@ def create_table_config_gpt():
); );
""") """)
print("Таблица config_gpt создана или уже существует") print("Таблица config_gpt создана или уже существует")
finally:
if conn:
conn.close()
def update_promt(url: str, name: str, promt: str): def update_promt(url: str, name: str, promt: str):
try:
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
INSERT INTO config_gpt (url, name, promt) INSERT INTO config_gpt (url, name, promt)
@@ -98,31 +128,45 @@ def update_promt(url: str, name: str, promt: str):
promt = EXCLUDED.promt promt = EXCLUDED.promt
""", (url, name, promt)) """, (url, name, promt))
conn.commit() conn.commit()
finally:
if conn:
conn.close()
def get_promt(promt_name_url): def get_promt(promt_name_url):
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur: with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute("SELECT promt FROM config_gpt WHERE url = %s", (promt_name_url,)) cur.execute("SELECT promt FROM config_gpt WHERE url = %s", (promt_name_url,))
promt = cur.fetchone() promt = cur.fetchone()
return promt['promt'] return promt['promt']
finally:
if conn:
conn.close()
def get_all_promt(): def get_all_promt():
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur: with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute("SELECT * FROM config_gpt") cur.execute("SELECT * FROM config_gpt")
rows = cur.fetchall() rows = cur.fetchall()
sources = [{"url": row["url"], "name": row["name"], "promt": row["promt"]} for row in rows] sources = [{"url": row["url"], "name": row["name"], "promt": row["promt"]} for row in rows]
return {"sources": sources} return {"sources": sources}
finally:
if conn:
conn.close()
# Возвращает список всех значений поля name из таблицы config_gpt
def get_all_categories_promt(): def get_all_categories_promt():
try:
with conn.cursor(cursor_factory=RealDictCursor) as cur: with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute("SELECT name FROM config_gpt") cur.execute("SELECT name FROM config_gpt")
rows = cur.fetchall() rows = cur.fetchall()
return [row["name"] for row in rows] return [row["name"] for row in rows]
finally:
if conn:
conn.close()
# Создание, сохранение и работа с таблицей ошибочных ссылок (error_url)
def create_table_error_url(): def create_table_error_url():
try:
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
CREATE TABLE IF NOT EXISTS error_url ( CREATE TABLE IF NOT EXISTS error_url (
@@ -132,9 +176,12 @@ def create_table_error_url():
); );
""") """)
print("Таблица error_url создана или уже существует") print("Таблица error_url создана или уже существует")
finally:
if conn:
conn.close()
def add_error_url(source_url: str, error_sources_url: str): def add_error_url(source_url: str, error_sources_url: str):
"""Добавляет запись об ошибке URL""" try:
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
INSERT INTO error_url (source_url, error_sources_url) INSERT INTO error_url (source_url, error_sources_url)
@@ -142,9 +189,12 @@ def add_error_url(source_url: str, error_sources_url: str):
RETURNING id; RETURNING id;
""", (source_url, error_sources_url)) """, (source_url, error_sources_url))
return cur.fetchone()[0] return cur.fetchone()[0]
finally:
if conn:
conn.close()
def check_error_url(error_sources_url: str) -> bool: def check_error_url(error_sources_url: str) -> bool:
"""Проверяет, есть ли запись в таблице error_url с таким URL""" try:
with conn.cursor(cursor_factory=RealDictCursor) as cur: with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(""" cur.execute("""
SELECT 1 FROM error_url SELECT 1 FROM error_url
@@ -154,9 +204,12 @@ def check_error_url(error_sources_url: str) -> bool:
row = cur.fetchone() row = cur.fetchone()
return row is None return row is None
finally:
if conn:
conn.close()
# Создание и работа с таблицей источников add_sources
def create_table_add_sourse(): def create_table_add_sourse():
try:
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
CREATE TABLE IF NOT EXISTS sourse ( CREATE TABLE IF NOT EXISTS sourse (
@@ -165,8 +218,12 @@ def create_table_add_sourse():
); );
""") """)
print("Таблица sourse создана или уже существует") print("Таблица sourse создана или уже существует")
finally:
if conn:
conn.close()
def add_sources(url: str, promt: str): def add_sources(url: str, promt: str):
try:
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
INSERT INTO config_gpt (url, promt) INSERT INTO config_gpt (url, promt)
@@ -175,6 +232,9 @@ def add_sources(url: str, promt: str):
promt = EXCLUDED.promt promt = EXCLUDED.promt
""", (url, promt)) """, (url, promt))
conn.commit() conn.commit()
finally:
if conn:
conn.close()
# Пример использования # Пример использования
# if __name__ == "__main__": # if __name__ == "__main__":