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

work_parser.py

@@ -26,6 +26,7 @@ class ParsedData(BaseModel):
     svodka: Optional[bool] = False
     donesenie: Optional[bool] = False
     bilutene: Optional[bool] = False
+    download: Optional[bool] = False
     other: str
     category: str
     translation_text: str
@@ -38,8 +39,8 @@ def save_parsed_data_to_db(data: ParsedData):
         conn = get_connection()
         with conn.cursor() as cur:
             cur.execute("""
-                INSERT INTO url (url, parsed_at, title, original_text, article_date, status, viewed, tematik, svodka, donesenie, bilutene, other, category, translation_text, short_text)
-                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+                INSERT INTO url (url, parsed_at, title, original_text, article_date, status, viewed, tematik, svodka, donesenie, download, bilutene, other, category, translation_text, short_text)
+                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                 ON CONFLICT (url) DO UPDATE SET
                     parsed_at = EXCLUDED.parsed_at,
                     title = EXCLUDED.title,
@@ -48,6 +49,7 @@ def save_parsed_data_to_db(data: ParsedData):
                     status = EXCLUDED.status,
                     viewed = EXCLUDED.viewed,
                     tematik = EXCLUDED.tematik,
+                    download = EXCLUDED.download,
                     svodka = EXCLUDED.svodka,
                     donesenie = EXCLUDED.donesenie,
                     bilutene = EXCLUDED.bilutene,
@@ -55,7 +57,7 @@ def save_parsed_data_to_db(data: ParsedData):
                     category = EXCLUDED.category,
                     translation_text = EXCLUDED.translation_text,
                     short_text = EXCLUDED.short_text;
-            """, (data.url, data.parsed_at, data.title, data.original_text, data.article_date, data.status, data.viewed, data.tematik, data.svodka, data.donesenie, data.bilutene, data.other, data.category, data.translation_text, data.short_text))
+            """, (data.url, data.parsed_at, data.title, data.original_text, data.article_date, data.status, data.viewed, data.tematik, data.svodka, data.donesenie, data.download, data.bilutene, data.other, data.category, data.translation_text, data.short_text))
             conn.commit()
         return {"status": "success", "message": "Данные успешно сохранены"}
     except Exception as e:
@@ -68,15 +70,7 @@ def save_parsed_data_to_db(data: ParsedData):
 
 def get_articles_by_filter(field_name: str, start_date: str, finish_date: str):
     """
-    Возвращает список заголовков статей по полю и диапазону дат
-    
-    Args:
-        field_name: имя поля (tematik, svodka, donesenie, bilutene, status)
-        start_date: дата начала в формате YYYY-MM-DD
-        finish_date: дата окончания в формате YYYY-MM-DD
-    
-    Returns:
-        List[str]: список заголовков (title)
+    Возвращает список заголовков статей по полю и диапазону дат для выгрузки
     """
     conn = get_connection()
     try:
@@ -100,6 +94,54 @@ def get_articles_by_filter(field_name: str, start_date: str, finish_date: str):
         raise
 
 
+def get_download_counts():
+    """
+    Возвращает количество статей для каждого поля, где поле = TRUE и download = FALSE
+    """
+    conn = get_connection()
+    try:
+        allowed_fields = ['tematik', 'svodka', 'donesenie', 'bilutene']
+        
+        with conn.cursor(cursor_factory=RealDictCursor) as cur:
+            counts = {}
+            for field in allowed_fields:
+                cur.execute(f"""
+                    SELECT COUNT(*) as count FROM url 
+                    WHERE {field} = TRUE 
+                    AND download = FALSE;
+                """)
+                row = cur.fetchone()
+                counts[field] = row['count']
+            
+            return counts
+    except Exception as e:
+        print(f"Ошибка в get_download_counts: {e}")
+        raise
+
+
+def mark_articles_as_downloaded(titles: list):
+    """
+    Обновляет download = TRUE для списка заголовков
+    """
+    if not titles:
+        return {"message": "Список заголовков пуст", "updated_rows": 0}
+    
+    conn = get_connection()
+    try:
+        with conn.cursor() as cur:
+            cur.execute("""
+                UPDATE url 
+                SET download = TRUE 
+                WHERE title = ANY(%s);
+            """, (titles,))
+            updated_rows = cur.rowcount
+        conn.commit()
+        return {"message": f"Статус download обновлён для {updated_rows} статей", "updated_rows": updated_rows}
+    except Exception as e:
+        print(f"Ошибка в mark_articles_as_downloaded: {e}")
+        raise
+
+
 # Глобальное подключение к БД
 conn = None
 
@@ -367,7 +409,6 @@ def create_table_add_sourse():
     except Exception as e:
         print(f"Ошибка при создании таблицы sourse: {e}")
 
-
 def add_sources(url: str, promt: str, status: bool = False):
     conn = get_connection()
     try:
@@ -384,7 +425,6 @@ def add_sources(url: str, promt: str, status: bool = False):
         print(f"Ошибка при добавлении источника: {e}")
         raise
 
-
 def get_all_sources(category: str):
     """Возвращает все записи из таблицы sourse. Сначала показываются записи со status=false"""
     conn = get_connection()
@@ -408,7 +448,6 @@ def get_all_sources(category: str):
         print(f"Ошибка при получении источников: {e}")
         return {"error": str(e), "sources": []}
 
-
 def get_true_sources():
     """Возвращает все записи из таблицы sourse. Сначала показываются записи со status=true"""
     conn = get_connection()
@@ -427,7 +466,6 @@ def get_true_sources():
         print(f"Ошибка при получении источников: {e}")
         return {"error": str(e), "sources": []}
 
-
 def update_source_status(url: str, status: bool = True):
     """Обновляет статус источника по URL"""
     conn = get_connection()
@@ -443,7 +481,6 @@ def update_source_status(url: str, status: bool = True):
         print(f"Ошибка при обновлении статуса: {e}")
         return {"error": str(e), "updated_rows": 0}
 
-
 def delete_sources(url: str):
     """Удаляет источник по URL из таблицы sourse"""
     conn = get_connection()