diff --git a/app/connectors_service/connectors/es/management_client.py b/app/connectors_service/connectors/es/management_client.py index 0aa2917f0..75ff8753a 100644 --- a/app/connectors_service/connectors/es/management_client.py +++ b/app/connectors_service/connectors/es/management_client.py @@ -189,7 +189,7 @@ async def bulk_insert(self, operations, pipeline): ) async def yield_existing_documents_metadata(self, index): - """Returns an iterator on the `id` and `_timestamp` fields of all documents in an index. + """Returns an iterator on the `_id` and `_timestamp` fields of all documents in an index. WARNING @@ -203,10 +203,10 @@ async def yield_existing_documents_metadata(self, index): return async for doc in async_scan( - client=self.client, index=index, _source=["id", TIMESTAMP_FIELD] + client=self.client, index=index, _source=[TIMESTAMP_FIELD] ): source = doc["_source"] - doc_id = source.get("id", doc["_id"]) + doc_id = doc["_id"] timestamp = source.get(TIMESTAMP_FIELD) yield doc_id, timestamp diff --git a/app/connectors_service/tests/es/test_management_client.py b/app/connectors_service/tests/es/test_management_client.py index 18ac77858..c8cbca5f9 100644 --- a/app/connectors_service/tests/es/test_management_client.py +++ b/app/connectors_service/tests/es/test_management_client.py @@ -192,6 +192,30 @@ async def test_yield_existing_documents_metadata_when_index_exists( assert ids == ["1", "2"] + @pytest.mark.asyncio + async def test_yield_existing_documents_metadata_when_source_id_differs_from_id( + self, es_management_client, mock_responses + ): + es_management_client.index_exists = AsyncMock(return_value=True) + + records = [ + {"_id": "1", "_source": {"id": "a", "_timestamp": str(datetime.now())}}, + {"_id": "2", "_source": {"id": "b", "_timestamp": str(datetime.now())}}, + ] + + with mock.patch( + "connectors.es.management_client.async_scan", + return_value=AsyncIterator(records), + ): + ids = [] + async for ( + doc_id, + _, + ) in es_management_client.yield_existing_documents_metadata("something"): + ids.append(doc_id) + + assert ids == ["1", "2"] + @pytest.mark.asyncio async def test_get_connector_secret(self, es_management_client, mock_responses): secret_id = "secret-id" diff --git a/app/connectors_service/tests/test_sink.py b/app/connectors_service/tests/test_sink.py index 2d5f80449..256f94a14 100644 --- a/app/connectors_service/tests/test_sink.py +++ b/app/connectors_service/tests/test_sink.py @@ -766,6 +766,38 @@ async def test_get_docs( assert queue_called_with_operations(queue, expected_queue_operations) +@mock.patch("connectors.es.management_client.async_scan") +@mock.patch( + "connectors.es.management_client.ESManagementClient.index_exists", + return_value=True, +) +@pytest.mark.asyncio +async def test_get_docs_matches_and_deletes_existing_docs_by_es_id( + index_exists, async_scan +): + # e.g. an ingest pipeline rewrote `id`, so it no longer matches `_id` + async_scan.return_value = AsyncIterator( + [ + {"_id": "1", "_source": {"id": "rewritten-1", "_timestamp": TIMESTAMP}}, + {"_id": "2", "_source": {"id": "rewritten-2", "_timestamp": TIMESTAMP}}, + ] + ) + queue = await queue_mock() + extractor = await setup_extractor(queue) + + await extractor.run( + AsyncIterator([(deepcopy(DOC_ONE), None, "index")]), JobType.FULL + ) + + assert extractor.counters.get(UPDATES_QUEUED) == 1 + assert extractor.counters.get(CREATES_QUEUED) == 0 + assert extractor.counters.get(DELETES_QUEUED) == 1 + assert queue_called_with_operations( + queue, + [index_operation(DOC_ONE), delete_operation(DOC_TWO), end_docs_operation()], + ) + + def test_log_progress_includes_extraction_counters(): logger_mock = Mock() extractor = Extractor(None, Mock(), INDEX, logger_=logger_mock)