Added a separate get_total_count() function

Author: Your Name

sde_collections/sinequa_api.py

@@ -257,7 +257,7 @@ def get_full_texts(
                 if total_count is None:
                     total_count = response.get("TotalRowCount", 0)
 
-                yield (self._process_rows_to_records(rows), total_count)
+                yield (self._process_rows_to_records(rows))
 
                 current_offset += len(rows)
 
@@ -275,6 +275,35 @@ def get_full_texts(
                 print(f"Reducing batch size to {current_batch_size} and retrying...")
                 continue
 
+    def get_total_count(self, collection_config_folder: str, source: str = None) -> int:
+        """
+        Retrieves the total count of records for a given collection using Sinequa's TotalRowCount metadata.
+
+        Args:
+            collection_config_folder (str): The collection folder to query (e.g., "EARTHDATA", "CASEI").
+            source (str, optional): The source to query. If None, defaults to "scrapers" for dev servers
+                or "SDE" for other servers.
+
+        Returns:
+            int: The total number of records in the collection.
+        """
+        if not source:
+            source = self._get_source_name()
+
+        if (index := self.config.get("index")) is None:
+            raise ValueError(
+                f"Configuration error: Index not defined for server '{self.server_name}'. "
+                "Please update server configuration with the required index."
+            )
+
+        # Minimal query to get only metadata, no data retrieval
+        sql = f"SELECT * FROM {index} WHERE collection = '/{source}/{collection_config_folder}/' SKIP 0 COUNT 0"
+
+        response = self._execute_sql_query(sql)
+
+        # Extract TotalRowCount from metadata
+        return response.get("TotalRowCount", 0)
+
     @staticmethod
     def _process_full_text_response(batch_data: dict):
         if "Rows" not in batch_data or not isinstance(batch_data["Rows"], list):

sde_collections/tasks.py

@@ -174,13 +174,13 @@ def fetch_and_replace_full_text(collection_id, server_name):
     # Step 1: Delete existing DumpUrl entries
     deleted_count, _ = DumpUrl.objects.filter(collection=collection).delete()
     print(f"Deleted {deleted_count} old records.")
-    total_server_count = 0
     try:
+        total_server_count = api.get_total_count(collection.config_folder)
+        print(f"Total records on the server: {total_server_count}")
+
         # Step 2: Process data in batches
         total_processed = 0
-        for batch, total_count in api.get_full_texts(collection.config_folder):
-            if total_server_count == 0:
-                total_server_count = total_count
+        for batch in api.get_full_texts(collection.config_folder):
             with transaction.atomic():
                 DumpUrl.objects.bulk_create(
                     [

sde_collections/tests/test_import_fulltexts.py

@@ -30,7 +30,7 @@ def test_fetch_and_replace_full_text(disconnect_signals):
     ]
 
     def mock_generator():
-        yield (mock_batch, len(mock_batch))
+        yield (mock_batch)
 
     with patch("sde_collections.sinequa_api.Api.get_full_texts") as mock_get_full_texts, patch(
         "sde_collections.utils.slack_utils.send_detailed_import_notification"
@@ -61,7 +61,7 @@ def mock_batch_generator():
         total_records = 20000
 
         for start in range(0, total_records, batch_size):
-            yield (create_batch(start, min(batch_size, total_records - start)), total_records)
+            yield (create_batch(start, min(batch_size, total_records - start)))
 
     with patch("sde_collections.sinequa_api.Api.get_full_texts") as mock_get_full_texts, patch(
         "sde_collections.utils.slack_utils.send_detailed_import_notification"

sde_collections/tests/test_sinequa_api.py

@@ -170,17 +170,16 @@ def test_get_full_texts_pagination(self, mock_execute_sql, api_instance):
 
         # Collect all batches from the iterator
         batches = list(api_instance.get_full_texts("test_folder"))
-        records_batches = [batch[0] for batch in batches]
-        assert len(records_batches) == 2  # Should have two batches
-        assert len(records_batches[0]) == 2  # First batch has 2 records
-        assert len(records_batches[1]) == 1  # Second batch has 1 record
+        assert len(batches) == 2  # Should have two batches
+        assert len(batches[0]) == 2  # First batch has 2 records
+        assert len(batches[1]) == 1  # Second batch has 1 record
 
         # Verify content of batches
-        assert records_batches[0] == [
+        assert batches[0] == [
             {"url": "http://example.com/1", "full_text": "Text 1", "title": "Title 1"},
             {"url": "http://example.com/2", "full_text": "Text 2", "title": "Title 2"},
         ]
-        assert records_batches[1] == [{"url": "http://example.com/3", "full_text": "Text 3", "title": "Title 3"}]
+        assert batches[1] == [{"url": "http://example.com/3", "full_text": "Text 3", "title": "Title 3"}]
 
     def test_get_full_texts_missing_index(self, api_instance):
         """
@@ -249,12 +248,11 @@ def test_get_full_texts_batch_size_reduction(self, mock_execute_sql, api_instanc
         ]
 
         batches = list(api_instance.get_full_texts("test_folder", batch_size=100, min_batch_size=1))
-        records_batches = [batch[0] for batch in batches]
 
         # Verify the batches were processed correctly after size reduction
         assert len(batches) == 1
-        assert len(records_batches[0]) == 1
-        assert records_batches[0][0]["url"] == "http://example.com/1"
+        assert len(batches[0]) == 1
+        assert batches[0][0]["url"] == "http://example.com/1"
 
         # Verify batch size reduction logic
         assert mock_execute_sql.call_count == 2

sde_collections/tests/test_workflow_status_triggers.py

@@ -149,8 +149,7 @@ def test_full_text_import_workflow(self, MockGitHub, MockApi, MockSlackNotificat
 
         # Setup mock API
         mock_api = Mock()
-        total_count = len(self.api_response)
-        mock_api.get_full_texts.return_value = iter([(self.api_response, total_count)])
+        mock_api.get_full_texts.return_value = iter([self.api_response])
         MockApi.return_value = mock_api
 
         # Setup initial workflow state