From a6275232a5d28b9d3d634f8a2e171312fb067efb Mon Sep 17 00:00:00 2001
From: axif <muhamadasif570@gmail.com>
Date: Sat, 26 Oct 2024 01:40:25 +0600
Subject: [PATCH] retry mechanism for total and fix overwrite

---
 src/scribe_data/cli/total.py           | 20 +++++++++++++++++++-
 src/scribe_data/wikidata/query_data.py | 12 ++++++------
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/scribe_data/cli/total.py b/src/scribe_data/cli/total.py
index 61df9b4e9..d7538e9bf 100644
--- a/src/scribe_data/cli/total.py
+++ b/src/scribe_data/cli/total.py
@@ -21,6 +21,8 @@
 """
 
 from SPARQLWrapper import JSON
+from urllib.error import HTTPError
+from http.client import IncompleteRead
 
 from scribe_data.utils import (
     LANGUAGE_DATA_EXTRACTION_DIR,
@@ -244,7 +246,23 @@ def get_total_lexemes(language, data_type, doPrint=True):
 
     sparql.setQuery(query)
     sparql.setReturnFormat(JSON)
-    results = sparql.query().convert()
+    try_count = 0
+    max_retries = 2
+    results = None
+
+    while try_count <= max_retries and results is None:
+        try:
+            results = sparql.query().convert()
+        except HTTPError as http_err:
+            print(f"HTTPError occurred: {http_err}")
+        except IncompleteRead as read_err:
+            print(f"Incomplete read error occurred: {read_err}")
+        try_count += 1
+
+        if results is None and try_count <= max_retries:
+            print("The query will be retried..")
+        elif results is None:
+            print("Query failed after retries.")
 
     # Check if the query returned any results.
     if (
diff --git a/src/scribe_data/wikidata/query_data.py b/src/scribe_data/wikidata/query_data.py
index c1f70ab99..f54ccce32 100644
--- a/src/scribe_data/wikidata/query_data.py
+++ b/src/scribe_data/wikidata/query_data.py
@@ -181,13 +181,13 @@ def query_data(
                         for file in existing_files:
                             file.unlink()
 
-                # elif choice in ["k", "K"]:
-                #     timestamp = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
-                #     file_name = f"{target_type}_{timestamp}.json"
+                    # elif choice in ["k", "K"]:
+                    #     timestamp = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
+                    #     file_name = f"{target_type}_{timestamp}.json"
 
-                else:
-                    print(f"Skipping update for {lang} {target_type}.")
-                    continue
+                    else:
+                        print(f"Skipping update for {lang} {target_type}.")
+                        break
 
         print(f"Querying and formatting {lang} {target_type}")