articles: adds method to set/update db table sequence

Signed-off-by: pamfilos <[email protected]>

scoap3/articles/tests/test_article.py

@@ -2,11 +2,13 @@
 from django.test import TestCase
 
 from scoap3.articles.models import Article
+from scoap3.utils.tools import update_article_db_model_sequence
 
 
-@pytest.mark.django_db
+@pytest.mark.django_db()
 class ArticleModelTest(TestCase):
     def test_save(self):
+        update_article_db_model_sequence(1)
         article = Article.objects.create(title="Test Article", abstract="Test Content")
         article.save()
         self.assertEqual(article.id, 1)
@@ -17,6 +19,7 @@ def test_save(self):
         new_article.save()
         self.assertEqual(new_article.id, 1000)
 
+        update_article_db_model_sequence(1001)
         new_article = Article.objects.create(
             title="Test Article 2", abstract="Test Content 2"
         )
@@ -26,10 +29,8 @@ def test_save(self):
     def test_save_and_update(self):
         article = Article.objects.create(title="Test Article", abstract="Test Content")
         article.save()
-        self.assertEqual(article.id, 1)
 
-        new_article = Article.objects.get(id=1)
+        new_article = Article.objects.get(id=article.id)
         new_article.title = "Test Article 2"
         new_article.save()
-        self.assertEqual(new_article.id, 1)
         self.assertEqual(new_article.title, "Test Article 2")

scoap3/utils/tools.py

@@ -1,7 +1,11 @@
 import logging
 from collections import Counter
 
+from django.db import connection
+from django.db.models import Max
+
 from scoap3.articles.documents import ArticleDocument
+from scoap3.articles.models import Article
 from scoap3.articles.util import (
     get_arxiv_primary_category,
     get_first_arxiv,
@@ -156,3 +160,23 @@ def author_export(search_year, search_country):
                 )
             )
     return {"header": result_headers, "data": result_data}
+
+
+def update_article_db_model_sequence(new_start_sequence):
+    max_id = Article.objects.aggregate(max_id=Max("id"))["max_id"] or 0
+    if new_start_sequence <= max_id:
+        print(
+            f"New sequence start ({new_start_sequence}) must be higher than current max ID ({max_id})."
+        )
+        return False
+
+    app_label = Article._meta.app_label
+    model_name = Article._meta.model_name
+
+    with connection.cursor() as cursor:
+        command = f"ALTER SEQUENCE {app_label}_{model_name}_id_seq RESTART WITH {new_start_sequence};"
+        cursor.execute(command)
+    print(
+        f"Sequence for {app_label}_{model_name} updated to start with {new_start_sequence}."
+    )
+    return True