<commit_before>import os
import json
import elasticsearch
import pymongo
from celery import Celery
from flask import current_app
from bhs_api import create_app
from bhs_api.utils import uuids_to_str

MIGRATE_MODE = os.environ.get('MIGRATE_MODE')
MIGRATE_ES = os.environ.get('MIGRATE_MODE', '1')

def make_celery():
    app, conf = create_app()
    if app.config['REDIS_PASSWORD']:
        redis_broker='redis://:{}@{}:{}/0'.format(
            app.config['REDIS_PASSWORD'],
            app.config['REDIS_HOST'],
            app.config['REDIS_PORT'],
        )
    else:
        redis_broker='redis://{}:{}/0'.format(
            app.config['REDIS_HOST'],
            app.config['REDIS_PORT'],
        )
    app.logger.info('Broker at {}'.format(app.config['REDIS_HOST']))
    celery = Celery(app.import_name, broker=redis_broker)
    celery.conf.update(app.config)
    celery.data_db = app.data_db
    # boiler plate to get our tasks running in the app context
    TaskBase = celery.Task
    class ContextTask(TaskBase):
        abstract = True
        def __call__(self, *args, **kwargs):
            with app.app_context():
                return TaskBase.__call__(self, *args, **kwargs)
    celery.Task = ContextTask
    return celery
celery = make_celery()


def update_es(collection, doc, id):
    if MIGRATE_ES != '1':
        return

    index_name = current_app.data_db.name
    try:
        current_app.es.index(index=index_name,
                             doc_type=collection,
                             id=id,
                             body=doc)
    except elasticsearch.exceptions.SerializationError:
        # UUID fields are causing es to crash, turn them to strings
        uuids_to_str(doc)
        try:
            current_app.es.index(index=index_name,
                                 doc_type=collection,
                                 id=id,
                                 body=doc)
        except elasticsearch.exceptions.SerializationError as e:
            current_app.logger.error("Elastic search index failed for {}:{} with {}"
                                     .format(collection, id, e))


def reslugify(collection, document):
    ''' append the document id to the slug to ensure uniquness '''
    for lang, val in document['Slug'].items():
        if val:
            doc_id = get_collection_id_field(collection.name)
            document['Slug'][lang] += '-' + str(document[doc_id])


def get_collection_id_field(collection_name):
    doc_id = 'UnitId'
    if collection_name == 'photos':
        doc_id = 'PictureId'
    elif collection_name == 'genTreeIndividuals':
        doc_id = 'ID'
    elif collection_name == 'persons':
        doc_id = 'id'
    elif collection_name == 'synonyms':
        doc_id = '_id'
    elif collection_name == 'trees':
        doc_id = 'num'
    return doc_id


@celery.task
def update_tree(data, collection=None):
    # from celery.contrib import rdb; rdb.set_trace()
    if not collection:
        collection = celery.data_db['trees']
    num = data['num']
    tree = collection.find_one({'num': num})
    doc = data.copy()
    file_id = doc.pop('file_id')
    date = doc.pop('date')
    persons = doc.pop('persons')
    current_ver = {'file_id':file_id,
                   'update_date': date,
                   'persons': persons}
    if tree:
        # don't add the same version twice
        for i in tree['versions']:
            if i['file_id'] == file_id:
                return
        doc['versions'] = tree['versions']
        doc['versions'].append(current_ver)
        collection.update_one({'num': num},
                              {'$set': doc},
                              upsert=True)

    else:
        doc['versions'] = [current_ver]
        collection.insert_one(doc)
    current_app.redis.set('tree_vers_'+str(num),
                          json.dumps(doc['versions']),
                          300)


def find_version(tree_vers, file_id):
    ''' finding the version index based on an array of version and a file id '''
    for i, ver in enumerate(tree_vers):
        if ver['file_id'] == file_id:
            return i
    # let's assume it's the next version
    return i+1


@celery.task
def update_row(doc, collection_name):
    collection = celery.data_db[collection_name]
    # from celery.contrib import rdb; rdb.set_trace()
    update_doc(collection, doc)

def update_collection(collection, query, doc):
    if MIGRATE_MODE  == 'i':
        return collection.insert(doc)
    else:<commit_msg>the migration worker now ensures we have the indexes<commit_after>import os
import json
import elasticsearch
import pymongo
from celery import Celery
from flask import current_app
from bhs_api import create_app
from bhs_api.utils import uuids_to_str

MIGRATE_MODE = os.environ.get('MIGRATE_MODE')
MIGRATE_ES = os.environ.get('MIGRATE_MODE', '1')
INDICES = {
    'places' : ['UnitId', 'DisplayStatusDesc', 'RightsDesc', 'StatusDesc', 'Header.En', 'Header.He'],
    'familyNames' : ['UnitId', 'DisplayStatusDesc', 'RightsDesc', 'StatusDesc', 'Header.En', 'Header.He'],
    'lexicon' : ['UnitId'],
    'photoUnits' : ['UnitId', 'DisplayStatusDesc', 'RightsDesc', 'StatusDesc', 'Header.En', 'Header.He'],
    'photos' : ['PictureId', 'PictureFileName', 'PicturePath'],
    'persons' : ['name_lc.0', 'name_lc.1', 'sex', 'BIRT_PLAC_lc', 'MARR_PLAC_lc', 'tree_num', 'DEAT_PLAC_lc'],
    'synonyms': ['s_group', 'str_lc'],
    'personalities' : ['UnitId', 'DisplayStatusDesc', 'RightsDesc', 'StatusDesc', 'Header.En', 'Header.He']
}

def make_celery():
    app, conf = create_app()
    if app.config['REDIS_PASSWORD']:
        redis_broker='redis://:{}@{}:{}/0'.format(
            app.config['REDIS_PASSWORD'],
            app.config['REDIS_HOST'],
            app.config['REDIS_PORT'],
        )
    else:
        redis_broker='redis://{}:{}/0'.format(
            app.config['REDIS_HOST'],
            app.config['REDIS_PORT'],
        )
    app.logger.info('Broker at {}'.format(app.config['REDIS_HOST']))
    celery = Celery(app.import_name, broker=redis_broker)
    celery.conf.update(app.config)
    celery.data_db = app.data_db
    # boiler plate to get our tasks running in the app context
    TaskBase = celery.Task
    class ContextTask(TaskBase):
        abstract = True
        def __call__(self, *args, **kwargs):
            with app.app_context():
                return TaskBase.__call__(self, *args, **kwargs)
    celery.Task = ContextTask
    return celery
celery = make_celery()


def ensure_indices(collection):
    indices = INDICES.get(collection.name, set())
    for index in indices:
        collection.ensure_index(index)

def update_es(collection, doc, id):
    if MIGRATE_ES != '1':
        return

    index_name = current_app.data_db.name
    try:
        current_app.es.index(index=index_name,
                             doc_type=collection,
                             id=id,
                             body=doc)
    except elasticsearch.exceptions.SerializationError:
        # UUID fields are causing es to crash, turn them to strings
        uuids_to_str(doc)
        try:
            current_app.es.index(index=index_name,
                                 doc_type=collection,
                                 id=id,
                                 body=doc)
        except elasticsearch.exceptions.SerializationError as e:
            current_app.logger.error("Elastic search index failed for {}:{} with {}"
                                     .format(collection, id, e))


def reslugify(collection, document):
    ''' append the document id to the slug to ensure uniquness '''
    for lang, val in document['Slug'].items():
        if val:
            doc_id = get_collection_id_field(collection.name)
            document['Slug'][lang] += '-' + str(document[doc_id])


def get_collection_id_field(collection_name):
    doc_id = 'UnitId'
    if collection_name == 'photos':
        doc_id = 'PictureId'
    elif collection_name == 'genTreeIndividuals':
        doc_id = 'ID'
    elif collection_name == 'persons':
        doc_id = 'id'
    elif collection_name == 'synonyms':
        doc_id = '_id'
    elif collection_name == 'trees':
        doc_id = 'num'
    return doc_id


@celery.task
def update_tree(data, collection=None):
    # from celery.contrib import rdb; rdb.set_trace()
    if not collection:
        collection = celery.data_db['trees']
    num = data['num']
    tree = collection.find_one({'num': num})
    doc = data.copy()
    file_id = doc.pop('file_id')
    date = doc.pop('date')
    persons = doc.pop('persons')
    current_ver = {'file_id':file_id,
                   'update_date': date,
                   'persons': persons}
    if tree:
        # don't add the same version twice
        for i in tree['versions']:
            if i['file_id'] == file_id:
                return
        doc['versions'] = tree['versions']
        doc['versions'].append(current_ver)
        collection.update_one({'num': num},
                              {'$set': doc},
                              upsert=True)

    else:
        doc['versions'] = [current_ver]
        collection.insert_one(doc)
    current_app.redis.set('tree_vers_'+str(num),
                          json.dumps(doc['versions']),
                          300)


def find_version(tree_vers, file_id):
    ''' finding the version index based on an array of version and a file id '''
    for i, ver in enumerate(tree_vers):
        if ver['file_id'] == file_id:
            return i
    # let's assume it's the next version
    return i+1


@celery.task
def update_row(doc, collection_name):
    collection = celery.data_db[collection_name]
    # from celery.contrib import rdb; rdb.set_trace()
    update_doc(collection, doc)
    ensure_indices(collection)

def update_collection(collection, query, doc):
    if MIGRATE_MODE  == 'i':
        return collection.insert(doc)
    else:
{'fraction_non_alphanumeric': 0.12251488545766474, 'fraction_numerical': 0.0018165304268846503, 'mean_word_length': 2.3513696313831587, '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': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2084761', 'n_tokens_mistral': 3059, 'n_tokens_neox': 2996, 'n_words': 684}