<commit_before>from kafka.common import OffsetRequest

from corehq.apps.change_feed.connection import get_kafka_client
from .document_types import CASE, FORM, DOMAIN, META, APP

# this is redundant but helps avoid import warnings until nothing references these
CASE = CASE
FORM = FORM
DOMAIN = DOMAIN
META = META
APP = APP

# new models
CASE_SQL = 'case-sql'
FORM_SQL = 'form-sql'
SMS = 'sms'
LEDGER = 'ledger'
COMMCARE_USER = 'commcare-user'
GROUP = 'group'
WEB_USER = 'web-user'


ALL = (
    CASE,
    CASE_SQL,
    COMMCARE_USER,
    DOMAIN,
    FORM,
    FORM_SQL,
    GROUP,
    LEDGER,
    META,
    SMS,
    WEB_USER,
    APP,
)


def get_topic(document_type_object):
    return document_type_object.primary_type



def get_topic_offset(topic):
    assert topic in ALL
    return get_all_offsets()[topic]


def get_multi_topic_offset(topics):
    offsets = get_all_offsets()
    return {topic: offsets[topic] for topic in topics}


def get_all_offsets():
    client = get_kafka_client()
    offset_requests = [OffsetRequest(topic, 0, -1, 1) for topic in ALL]
    responses = client.send_offset_request(offset_requests)
    return {
        r.topic: r.offsets[0] for r in responses
    }<commit_msg>minor refactor<commit_after>from kafka.common import OffsetRequest

from corehq.apps.change_feed.connection import get_kafka_client
from .document_types import CASE, FORM, DOMAIN, META, APP

# this is redundant but helps avoid import warnings until nothing references these
CASE = CASE
FORM = FORM
DOMAIN = DOMAIN
META = META
APP = APP

# new models
CASE_SQL = 'case-sql'
FORM_SQL = 'form-sql'
SMS = 'sms'
LEDGER = 'ledger'
COMMCARE_USER = 'commcare-user'
GROUP = 'group'
WEB_USER = 'web-user'


ALL = (
    CASE,
    CASE_SQL,
    COMMCARE_USER,
    DOMAIN,
    FORM,
    FORM_SQL,
    GROUP,
    LEDGER,
    META,
    SMS,
    WEB_USER,
    APP,
)


def get_topic(document_type_object):
    return document_type_object.primary_type


def get_topic_offset(topic):
    """
    :returns: The kafka offset for the topic."""
    return get_multi_topic_offset([topic])[topic]


def get_multi_topic_offset(topics):
    """
    :returns: A dict of offsets keyed by topic"""
    assert set(topics) < set(ALL)
    client = get_kafka_client()
    offset_requests = [OffsetRequest(topic, 0, -1, 1) for topic in topics]
    responses = client.send_offset_request(offset_requests)
    return {
        r.topic: r.offsets[0] for r in responses
    }


def get_all_offsets():
    """
    :returns: A dict of offsets keyed by topic"""
    return get_multi_topic_offset(ALL)
{'fraction_non_alphanumeric': 0.093603744149766, 'fraction_numerical': 0.0031201248049922, 'mean_word_length': 3.6133093525179856, 'pattern_counts': {'":': 0, '<': 4, '<?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': '951680', 'n_tokens_mistral': 939, 'n_tokens_neox': 884, 'n_words': 240}