data_source_repeaters_performance.py

from collections import defaultdict

from django.db import models

from corehq import toggles
from corehq.motech.repeaters.const import State
from corehq.motech.repeaters.models import Repeater, RepeatRecord

states_names_by_keys = {
    s.value: s.name for s in State
}
data_source_repeater_name_by_ids = {}
data_source_repeater_status_by_name = {}


def show_stats():
    stats = dict(_get_repeaters_stats())
    for domain in stats.keys():
        print('--------------------')
        print(f"Domain: {domain}")

        domain_stats = stats[domain]
        for repeater_name in domain_stats.keys():
            counts_by_status = []
            repeater_stats = domain_stats[repeater_name]
            print(f"{repeater_name} ({data_source_repeater_status_by_name[repeater_name]})")

            for state_name in repeater_stats.keys():
                counts_by_status.extend([state_name, str(repeater_stats[state_name])])
            print(f"{' '.join(counts_by_status)}")
            print('----------')


def _get_repeaters_stats():
    stats = defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))

    for domain in _get_domains_using_data_source_repeaters():
        data_source_repeaters = _get_data_source_repeaters(domain)
        data_source_repeater_name_by_ids.update({
            repeater.id: repeater.name
            for repeater in data_source_repeaters
        })
        data_source_repeater_status_by_name.update({
            repeater.name: 'Paused' if repeater.is_paused else 'Active'
            for repeater in data_source_repeaters
        })

        for data_source_repeater in data_source_repeaters:
            repeater_name = data_source_repeater_name_by_ids[data_source_repeater.id]
            for repeat_record_state in State:
                stats[domain][repeater_name][repeat_record_state.name] = 0

        query = (
            RepeatRecord.objects.filter(domain=domain, repeater_id__in=data_source_repeater_name_by_ids.keys())
            .values('repeater_id', 'state')
            .order_by()
            .annotate(count=models.Count('*'))
            .values_list('repeater_id', 'state', 'count')
        )

        for repeater_id, state, count in query:
            stats[domain][data_source_repeater_name_by_ids[repeater_id]][states_names_by_keys[state]] = count

    return stats


def _get_domains_using_data_source_repeaters():
    return toggles.SUPERSET_ANALYTICS.get_enabled_domains()


def _get_data_source_repeaters(domain):
    repeaters = Repeater.objects.filter(
        domain=domain,
        repeater_type='DataSourceRepeater'
    )
    return list(repeaters)