limit fix

Author: EEnota

main.py

@@ -54,12 +54,14 @@ def pipeline() -> None:
     num_of_news = 0
     page = 1
     while num_of_news < limit:
-        raw_file_name = make_extract(category, key_word, page, page_size)
-        clean_file_name = transform_article(raw_file_name, category, key_word, page)
-        result_num_of_news = load_news(clean_file_name)
-        if result_num_of_news == 0:
+        remaining = limit - num_of_news
+        raw_file_name,raw_articles_count = make_extract(category, key_word, page, page_size)
+        if raw_articles_count == 0:
             logger.warning("there is no more artical")
             break
+        clean_file_name = transform_article(raw_file_name, category, key_word, page)
+        result_num_of_news = load_news(clean_file_name, max_rows=remaining)
+        
         num_of_news += result_num_of_news
         page += 1
     logger.info(f"{num_of_news} news on category {category} already aploaded")

src/extract.py

@@ -15,14 +15,15 @@ def import_to_raw_json(data:dict[str, Any], category: str, key_word: str, page:
     timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
     create_data = f"{timestamp}_{category}_{key_word}_page_{page}.json"
     file_path = raw_dir / create_data
+
     with open(file_path, "w", encoding="utf-8") as f:
         json.dump(data, f, ensure_ascii=False, indent=2)
 
     return create_data
 
 
 
-def make_extract(category: str, key_word: str, page: int = 1, page_size: int = 100) -> str:
+def make_extract(category: str, key_word: str, page: int = 1, page_size: int = 100) -> tuple[str,int]:
     params = {
     "apiKey": settings.KEY_API,
     "country":settings.COUNTRY,
@@ -33,18 +34,18 @@ def make_extract(category: str, key_word: str, page: int = 1, page_size: int = 1
     }
     try:
         data = r.get(settings.NEWS_URL, params=params, timeout=15)
-        data.raise_for_status()
         payload = data.json()
         logger.info(f"raise of status: {data.raise_for_status()}")
         payload["fetched_at"] = datetime.now().isoformat()
         payload["country"] = settings.COUNTRY
         payload["category"] = category
         payload["key_word"] = key_word
-        if payload.get("totalResults", 0) == 0:
-            logger.info("There no more articles")
+        articles_count = len(payload.get("articles", []))
+        if articles_count == 0:
+            logger.info("There are no more articles")
 
         new_file_name = import_to_raw_json(payload, category, key_word, page)
-        return new_file_name
+        return new_file_name, articles_count
 
     except r.exceptions.Timeout:
         logger.error("Error: NewsAPI reauest time out")

src/load.py

@@ -2,9 +2,10 @@
 from pathlib import Path
 import json
 from config.config import settings
+from typing import Optional
 BASE_DIR = (Path(__file__).resolve().parent.parent)/"data"/"clean"
 
-def load_news(clean_news:str)-> int:
+def load_news(clean_news:str, max_rows: Optional[int] = None)-> int:
     LOAD_DIR = BASE_DIR/clean_news
     num_of_news = 0
     with open(LOAD_DIR, 'r', encoding='utf-8') as f:
@@ -30,6 +31,8 @@ def load_news(clean_news:str)-> int:
 
     with get_cursor(settings.db_news) as (conn, cur):
         for new in data:
+            if max_rows is not None and num_of_news >= max_rows:
+                break
             cur.execute(query, 
                         (new["country"],
                          new["category"],