instrdatastore.py

from logging import warning

def wrap_func(func, name=None):
    if name is None:
        name = func.__name__
    def wrapper(*args, **kwargs):
        try:
            _DS_INFO[name] += 1
            _DS_INFO[name + '_args'].append((args, kwargs))
        except KeyError:
            _DS_INFO[name] = 1
            _DS_INFO[name + '_args'] = [(args, kwargs)]
        return func(*args, **kwargs)
    wrapper._inner_func = func
    return wrapper

def reset():
    global _DS_INFO
    _DS_INFO = {}

def print_info(print_args=True):
    warning('datastore access information')
    for funcname in ('Get', 'Put', 'Delete', 'Query'):
        try:
            warning('nb %s: %s', funcname, _DS_INFO[funcname])
            if print_args:
                warning('arguments (args, kwargs): \n%s', 
                        '\n'.join(sorted(str(args) for args in _DS_INFO[funcname + '_args'])))
        except KeyError:
            pass
    reset()

from google.appengine.api import datastore
warning('instrumentizing the datastore')
_DS_INFO = {}

if getattr(datastore, '_instrumented', False):
    datastore.Get = wrap_func(datastore.Get._inner_func)
    datastore.Delete = wrap_func(datastore.Delete._inner_func)
    datastore.Put = wrap_func(datastore.Put._inner_func)
    datastore.Query.__init__ = wrap_func(datastore.Query.__init__._inner_func, 'Query')
else:
    datastore.Get = wrap_func(datastore.Get)
    datastore.Delete = wrap_func(datastore.Delete)
    datastore.Put = wrap_func(datastore.Put)
    datastore.Query.__init__ = wrap_func(datastore.Query.__init__, 'Query')
    datastore._instrumented = True