добавил закрытие бд для пресечения утечки данных

main.py

@@ -459,6 +459,10 @@ class Parserall(BaseModel):
     url: HttpUrl
     promt: str
 
+@app.post("/add_sources" , summary="Добавление парсинга любого источника")
+async def add_sources_all_ist(sources: Parserall):
+    return wp.add_sources(sources.url, sources.promt)
+
 @app.post("/parser_all" , summary="Запуск процеса парсинга любого источника")
 async def process_parser_all_ist(url: Parserall, background_tasks: BackgroundTasks):
     background_tasks.add_task(start_pars_all_istochnik(str(url.url), url.promt))

work_parser.py

@@ -13,150 +13,228 @@ conn = psycopg2.connect(
 )
 conn.autocommit = True
 
+def close_connection():
+    """Закрывает подключение к БД"""
+    global conn
+    if conn:
+        conn.close()
+        conn = None
+
 def create_table():
-    with conn.cursor() as cur:
-        cur.execute("""
-            CREATE TABLE IF NOT EXISTS work_parser (
-                id SERIAL PRIMARY KEY,
-                status VARCHAR(20) NOT NULL CHECK (status IN ('queued', 'in_progress', 'completed', 'failed')),
-                created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
-                started_at TIMESTAMPTZ,
-                finished_at TIMESTAMPTZ,
-                source_url TEXT,
-                error_message TEXT,
-                attempts INTEGER DEFAULT 0,
-                priority INTEGER DEFAULT 0
-            );
-        """)
-    print("Таблица work_parser создана или уже существует")
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                CREATE TABLE IF NOT EXISTS work_parser (
+                    id SERIAL PRIMARY KEY,
+                    status VARCHAR(20) NOT NULL CHECK (status IN ('queued', 'in_progress', 'completed', 'failed')),
+                    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                    started_at TIMESTAMPTZ,
+                    finished_at TIMESTAMPTZ,
+                    source_url TEXT,
+                    error_message TEXT,
+                    attempts INTEGER DEFAULT 0,
+                    priority INTEGER DEFAULT 0
+                );
+            """)
+        print("Таблица work_parser создана или уже существует")
+    finally:
+        if conn:
+            conn.close()
 
 def insert_task(status, source_url=None, source_id=None, priority=0):
-    with conn.cursor() as cur:
-        cur.execute("""
-            INSERT INTO work_parser (status, source_url, priority)
-            VALUES (%s, %s, %s)
-            RETURNING id;
-        """, (status, source_url, priority))
-        task_id = cur.fetchone()[0]
-        return task_id
-
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                INSERT INTO work_parser (status, source_url, priority)
+                VALUES (%s, %s, %s)
+                RETURNING id;
+            """, (status, source_url, priority))
+            task_id = cur.fetchone()[0]
+            return task_id
+    finally:
+        if conn:
+            conn.close()
 
 def get_tasks_offset(limit, offset):
-    with conn.cursor(cursor_factory=RealDictCursor) as cur:
-        cur.execute("""
-            SELECT * FROM work_parser
-            ORDER BY created_at DESC
-            LIMIT %s OFFSET %s
-            """, (limit, offset))
-        tasks = cur.fetchall()
-        return tasks
+    try:
+        with conn.cursor(cursor_factory=RealDictCursor) as cur:
+            cur.execute("""
+                SELECT * FROM work_parser
+                ORDER BY created_at DESC
+                LIMIT %s OFFSET %s
+                """, (limit, offset))
+            tasks = cur.fetchall()
+            return tasks
+    finally:
+        if conn:
+            conn.close()
 
 def delete_task(task_id: int):
-    with conn.cursor(cursor_factory=RealDictCursor) as cur:
-        cur.execute("DELETE FROM work_parser WHERE id = %s RETURNING *;", (task_id,))
-        deleted_task = cur.fetchone()
-        if deleted_task:
-            return {"message": f"Задача {task_id} удалена", "deleted_task": dict(deleted_task)}
-        else:
-            return {"message": f"Задача с id {task_id} не найдена"}
+    try:
+        with conn.cursor(cursor_factory=RealDictCursor) as cur:
+            cur.execute("DELETE FROM work_parser WHERE id = %s RETURNING *;", (task_id,))
+            deleted_task = cur.fetchone()
+            if deleted_task:
+                return {"message": f"Задача {task_id} удалена", "deleted_task": dict(deleted_task)}
+            else:
+                return {"message": f"Задача с id {task_id} не найдена"}
+    finally:
+        if conn:
+            conn.close()
 
 def update_task(task_id, **fields):
-    # dynamic update query generator
-    allowed_fields = ['status', 'started_at', 'finished_at', 'source_url', 'error_message', 'attempts', 'priority']
-    set_parts = []
-    values = []
-    for key, value in fields.items():
-        if key in allowed_fields:
-            set_parts.append(f"{key} = %s")
-            values.append(value)
-    if not set_parts:
-        return False
-    values.append(task_id)
-    set_sql = ", ".join(set_parts)
-    with conn.cursor() as cur:
-        cur.execute(f"UPDATE work_parser SET {set_sql} WHERE id = %s;", values)
-    return True
-
+    try:
+        allowed_fields = ['status', 'started_at', 'finished_at', 'source_url', 'error_message', 'attempts', 'priority']
+        set_parts = []
+        values = []
+        for key, value in fields.items():
+            if key in allowed_fields:
+                set_parts.append(f"{key} = %s")
+                values.append(value)
+        if not set_parts:
+            return False
+        values.append(task_id)
+        set_sql = ", ".join(set_parts)
+        with conn.cursor() as cur:
+            cur.execute(f"UPDATE work_parser SET {set_sql} WHERE id = %s;", values)
+        return True
+    finally:
+        if conn:
+            conn.close()
 
 def create_table_config_gpt():
-    with conn.cursor() as cur:
-        cur.execute("""
-            CREATE TABLE IF NOT EXISTS config_gpt (
-                url TEXT PRIMARY KEY,
-                name VARCHAR(20),
-                promt TEXT
-            );
-        """)
-    print("Таблица config_gpt создана или уже существует")
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                CREATE TABLE IF NOT EXISTS config_gpt (
+                    url TEXT PRIMARY KEY,
+                    name VARCHAR(20),
+                    promt TEXT
+                );
+            """)
+        print("Таблица config_gpt создана или уже существует")
+    finally:
+        if conn:
+            conn.close()
 
 def update_promt(url: str, name: str, promt: str):
-    with conn.cursor() as cur:
-        cur.execute("""
-            INSERT INTO config_gpt (url, name, promt)
-            VALUES (%s, %s, %s)
-            ON CONFLICT (url) DO UPDATE SET 
-                name = EXCLUDED.name,
-                promt = EXCLUDED.promt
-        """, (url, name, promt))
-    conn.commit()
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                INSERT INTO config_gpt (url, name, promt)
+                VALUES (%s, %s, %s)
+                ON CONFLICT (url) DO UPDATE SET 
+                    name = EXCLUDED.name,
+                    promt = EXCLUDED.promt
+            """, (url, name, promt))
+        conn.commit()
+    finally:
+        if conn:
+            conn.close()
 
 def get_promt(promt_name_url):
-    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'] 
+    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()
     
 def get_all_promt():
-    with conn.cursor(cursor_factory=RealDictCursor) as cur:
-        cur.execute("SELECT * FROM config_gpt")
-        rows = cur.fetchall()
+    try:
+        with conn.cursor(cursor_factory=RealDictCursor) as cur:
+            cur.execute("SELECT * FROM config_gpt")
+            rows = cur.fetchall()
         
-    sources = [{"url": row["url"], "name": row["name"], "promt": row["promt"]} for row in rows]
-    return {"sources": sources}
+        sources = [{"url": row["url"], "name": row["name"], "promt": row["promt"]} for row in rows]
+        return {"sources": sources}
+    finally:
+        if conn:
+            conn.close()
 
 def get_all_categories_promt():
-    """Возвращает список всех значений поля name из таблицы config_gpt"""
-    with conn.cursor(cursor_factory=RealDictCursor) as cur:
-        cur.execute("SELECT name FROM config_gpt")
-        rows = cur.fetchall()
+    try:
+        with conn.cursor(cursor_factory=RealDictCursor) as cur:
+            cur.execute("SELECT name FROM config_gpt")
+            rows = cur.fetchall()
         
-    return [row["name"] for row in rows]
+        return [row["name"] for row in rows]
+    finally:
+        if conn:
+            conn.close()
 
 def create_table_error_url():
-    with conn.cursor() as cur:
-        cur.execute("""
-            CREATE TABLE IF NOT EXISTS error_url (
-                id SERIAL PRIMARY KEY,
-                source_url TEXT NOT NULL,
-                error_sources_url TEXT NOT NULL
-            );
-        """)
-    print("Таблица error_url создана или уже существует")
-
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                CREATE TABLE IF NOT EXISTS error_url (
+                    id SERIAL PRIMARY KEY,
+                    source_url TEXT NOT NULL,
+                    error_sources_url TEXT NOT NULL
+                );
+            """)
+        print("Таблица error_url создана или уже существует")
+    finally:
+        if conn:
+            conn.close()
 
 def add_error_url(source_url: str, error_sources_url: str):
-    """Добавляет запись об ошибке URL"""
-    with conn.cursor() as cur:
-        cur.execute("""
-            INSERT INTO error_url (source_url, error_sources_url)
-            VALUES (%s, %s)
-            RETURNING id;
-        """, (source_url, error_sources_url))
-        return cur.fetchone()[0]
-
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                INSERT INTO error_url (source_url, error_sources_url)
+                VALUES (%s, %s)
+                RETURNING id;
+            """, (source_url, error_sources_url))
+            return cur.fetchone()[0]
+    finally:
+        if conn:
+            conn.close()
 
 def check_error_url(error_sources_url: str) -> bool:
-    """Проверяет, есть ли запись в таблице error_url с таким URL"""
-    with conn.cursor(cursor_factory=RealDictCursor) as cur:
-        cur.execute("""
-            SELECT 1 FROM error_url 
-            WHERE error_sources_url = %s
-            LIMIT 1;
-        """, (error_sources_url,))
-        row = cur.fetchone()
+    try:
+        with conn.cursor(cursor_factory=RealDictCursor) as cur:
+            cur.execute("""
+                SELECT 1 FROM error_url 
+                WHERE error_sources_url = %s
+                LIMIT 1;
+            """, (error_sources_url,))
+            row = cur.fetchone()
         
-    return row is None
+        return row is None
+    finally:
+        if conn:
+            conn.close()
 
+def create_table_add_sourse():
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                CREATE TABLE IF NOT EXISTS sourse (
+                    url TEXT PRIMARY KEY,
+                    promt TEXT
+                );
+                """)
+        print("Таблица sourse создана или уже существует")
+    finally:
+        if conn:
+            conn.close()
+
+def add_sources(url: str, promt: str):
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                INSERT INTO config_gpt (url, promt)
+                VALUES (%s, %s)
+                ON CONFLICT (url) DO UPDATE SET 
+                    promt = EXCLUDED.promt
+            """, (url, promt))
+        conn.commit()
+    finally:
+        if conn:
+            conn.close()
 
 # Пример использования
 # if __name__ == "__main__":
@@ -168,3 +246,4 @@ def check_error_url(error_sources_url: str) -> bool:
 #     })
 #     # print(get_promt("japan"))
 #     # create_table_error_url()
+    # create_table_add_sourse()