udata search index profiling

profiling.py

# -*- coding: utf-8 -*-
from __future__ import unicode_literals, print_function

import logging
import sys

from udata.app import create_app, standalone
from udata.commands import init_logging
from udata.search import es, commands as cmd

import objgraph

LIMIT = 10000
MAX_COMMON_TYPES = 15
MAX_GROW = 5
INDEX = 'udata-profiling'

app = standalone(create_app(init_logging=init_logging))
log = logging.getLogger('__main__')


def before():
    log.info('Object graph before')
    objgraph.show_most_common_types(limit=MAX_COMMON_TYPES)
    objgraph.show_growth(limit=MAX_GROW)


def perform():
    log.info('Initiliazing index "{0}"'.format(INDEX))

    if es.indices.exists(INDEX):
        es.indices.delete(INDEX)

    es.initialize(INDEX)

    with cmd.handle_error(INDEX):
        cmd.disable_refresh(INDEX)
        index_some_datasets()
        es.indices.delete(index=INDEX)


def index_some_datasets():
    from udata.core.dataset.models import Dataset
    from udata.core.dataset.search import DatasetSearch

    qs = Dataset.objects.visible()
    qs = qs.exclude(*DatasetSearch.exclude_fields)

    qs = qs

    docs = cmd.iter_qs(qs[:LIMIT], DatasetSearch)
    docs = cmd.iter_for_index(docs, INDEX)

    for ok, info in cmd.streaming_bulk(es.client, docs, raise_on_error=False):
        print('.' if ok else 'x', end='')
        sys.stdout.flush()
    print('')


def after(*types):
    log.info('Object graph after')
    objgraph.show_most_common_types(limit=MAX_COMMON_TYPES)
    objgraph.show_growth()
    
    for typ in types:
        log.info('%ss %s', typ, objgraph.count(typ))
        log.info('%ss without leaking ones: %s', typ,
                 objgraph.count(typ, objgraph.get_leaking_objects()))
        objects = objgraph.by_type(typ)
        leakings = objgraph.get_leaking_objects(objects)[:5]
        objgraph.show_backrefs(leakings, max_depth=4)


before()
with app.app_context():
    perform()
after('Dataset')