<commit_before>import logging
from banal import is_mapping, ensure_list
from followthemoney import model
from followthemoney.exc import InvalidData
from followthemoney.pragma import remove_checksums

from aleph.model import Entity, Document
from aleph.queues import ingest_entity
from aleph.analysis import analyze_entity
from aleph.queues import queue_task, OP_INDEX
from aleph.index.entities import index_bulk
from aleph.logic.entities import refresh_entity_id
from aleph.logic.collections import refresh_collection
from aleph.logic.aggregator import get_aggregator

log = logging.getLogger(__name__)


def _collection_proxies(collection):
    for entity in Entity.by_collection(collection.id).yield_per(5000):
        yield entity.to_proxy()
    for document in Document.by_collection(collection.id).yield_per(5000):
        yield document.to_proxy()


def process_collection(stage, collection, ingest=True, sync=False):
    """Trigger a full re-parse of all documents and re-build the
    search index from the aggregator."""
    aggregator = get_aggregator(collection)
    for proxy in _collection_proxies(collection):
        if ingest and proxy.schema.is_a(Document.SCHEMA):
            ingest_entity(collection, proxy,
                          job_id=stage.job.id,
                          sync=sync)
        else:
            aggregator.put(proxy, fragment='db')
            queue_task(collection, OP_INDEX,
                       job_id=stage.job.id,
                       payload={'entity_id': proxy.id},
                       context={'sync': sync})
    aggregator.close()


def _process_entity(entity, sync=False):
    """Perform pre-index processing on an entity, includes running the
    NLP pipeline."""
    analyze_entity(entity)
    if sync:
        refresh_entity_id(entity.id)
    # log.debug("Index: %r", entity)
    return entity


def _fetch_entities(stage, collection, entity_ids=None, batch=100):
    aggregator = get_aggregator(collection)
    if entity_ids is not None:
        entity_ids = ensure_list(entity_ids)
        # WEIRD: Instead of indexing a single entity, this will try
        # pull a whole batch of them off the queue and do it at once.
        tasks = stage.get_tasks(limit=max(1, batch - len(entity_ids)))
        for task in tasks:
            entity_ids.extend(ensure_list(task.payload.get('entity_ids')))
        # FIXME: this doesn't retain mapping_id properly.
        stage.mark_done(len(tasks))

    yield from aggregator.iterate(entity_id=entity_ids)
    aggregator.close()


def index_aggregate(stage, collection, sync=False, entity_ids=None,
                    mapping_id=None):
    """Project the contents of the collections aggregator into the index."""
    entities = _fetch_entities(stage, collection, entity_ids=entity_ids)
    entities = (_process_entity(e, sync=sync) for e in entities)
    extra = {'job_id': stage.job.id, 'mapping_id': mapping_id}
    index_bulk(collection, entities, extra)
    refresh_collection(collection.id)<commit_msg>Support really sync indexing<commit_after>import logging
from banal import is_mapping, ensure_list
from followthemoney import model
from followthemoney.exc import InvalidData
from followthemoney.pragma import remove_checksums

from aleph.model import Entity, Document
from aleph.queues import ingest_entity
from aleph.analysis import analyze_entity
from aleph.queues import queue_task, OP_INDEX
from aleph.index.entities import index_bulk
from aleph.logic.entities import refresh_entity_id
from aleph.logic.collections import refresh_collection
from aleph.logic.aggregator import get_aggregator

log = logging.getLogger(__name__)


def _collection_proxies(collection):
    for entity in Entity.by_collection(collection.id).yield_per(5000):
        yield entity.to_proxy()
    for document in Document.by_collection(collection.id).yield_per(5000):
        yield document.to_proxy()


def process_collection(stage, collection, ingest=True, sync=False):
    """Trigger a full re-parse of all documents and re-build the
    search index from the aggregator."""
    aggregator = get_aggregator(collection)
    for proxy in _collection_proxies(collection):
        if ingest and proxy.schema.is_a(Document.SCHEMA):
            ingest_entity(collection, proxy,
                          job_id=stage.job.id,
                          sync=sync)
        else:
            aggregator.put(proxy, fragment='db')
            queue_task(collection, OP_INDEX,
                       job_id=stage.job.id,
                       payload={'entity_ids': [proxy.id]},
                       context={'sync': sync})
    aggregator.close()


def _process_entity(entity, sync=False):
    """Perform pre-index processing on an entity, includes running the
    NLP pipeline."""
    analyze_entity(entity)
    if sync:
        refresh_entity_id(entity.id)
    # log.debug("Index: %r", entity)
    return entity


def _fetch_entities(stage, collection, entity_ids=None, batch=100):
    aggregator = get_aggregator(collection)
    if entity_ids is not None:
        entity_ids = ensure_list(entity_ids)
        # WEIRD: Instead of indexing a single entity, this will try
        # pull a whole batch of them off the queue and do it at once.
        tasks = stage.get_tasks(limit=max(1, batch - len(entity_ids)))
        for task in tasks:
            entity_ids.extend(ensure_list(task.payload.get('entity_ids')))
        # FIXME: this doesn't retain mapping_id properly.
        stage.mark_done(len(tasks))

    yield from aggregator.iterate(entity_id=entity_ids)
    aggregator.close()


def index_aggregate(stage, collection, sync=False, entity_ids=None,
                    mapping_id=None):
    """Project the contents of the collections aggregator into the index."""
    entities = _fetch_entities(stage, collection, entity_ids=entity_ids)
    entities = (_process_entity(e, sync=sync) for e in entities)
    extra = {'job_id': stage.job.id, 'mapping_id': mapping_id}
    index_bulk(collection, entities, extra, sync=sync)
    refresh_collection(collection.id, sync=sync)
{'fraction_non_alphanumeric': 0.0858085808580858, 'fraction_numerical': 0.0039603960396039604, 'mean_word_length': 3.3920289855072463, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3482647', 'n_tokens_mistral': 1757, 'n_tokens_neox': 1702, 'n_words': 466}