import itertools from concurrent.futures.thread import ThreadPoolExecutor from functools import reduce from typing import Dict, Iterable, List, Union from google.protobuf.json_format import MessageToDict from google.protobuf.pyext._message import MessageMapContainer from structlog import get_logger from delphi_api.const import MAX_THREADS, USE_THREAD_POOL_FOR_AGG LOG = get_logger(__name__) class Models: """ Static data model utility functions for parsing and aggregating data from models """ MULTIKEY_SEP = ';' KEY_VAL_SEP = ':' @classmethod def aggregate_models(cls, models, base_key_name='isrc', country: str = None, group_by: Iterable[str] = None, stats_fields: Iterable[str] = ('country_stats',)) -> dict: """Specific to streams aggregations Args: models: Iterable of model data (rows) base_key_name: Key name of the base field for aggregations (``'isrc'``, ``'playlist_id'``, etc.) stats_fields: model field name(s) containing underlying data metrics to aggregate country: Provide to filter for only one specific ``country_code`` group_by: Iterable of field names for grouping results by (ex: ``date``) Raises: AttributeError: model does not have stats field(s) provided Returns: dict: a dictionary of items with "multikeys" used for grouping """ data = {} for model in models: base_key_val = getattr(model, base_key_name) base_key = f'{base_key_name}:{base_key_val}' if group_by: base_key += cls.get_group_by_multikey(model, group_by=group_by) data[base_key] = data.get(base_key, []) stats_data = [] for field in stats_fields: model_field = getattr(model, field, None) if model_field and hasattr(model_field, 'countries'): stats_data.append(cls.cast_proto_map(model_field.countries, only=country)) if not stats_data: continue # add stats_fields together for multiple (ex: streams + demographics) item = reduce(cls.add_nested_dicts, stats_data, {}) data[base_key].append(item) return cls._reduce_add_nested(data) @classmethod def aggregate_generic(cls, models: Iterable[dict], base_key_name='isrc', group_by: Iterable[str] = None) -> dict: """Similar to :class:`Models.aggregate_models` but used for generic dictionaries Args: models: Iterable of model data (rows) base_key_name: Key name of the base field for aggregations (``'isrc'``, ``'playlist_id'``, etc.) group_by: Iterable of field names for grouping results by (ex: ``date``) Returns: dict: a dictionary of items with "multikeys" used for grouping """ data = {} for model in models: base_key_val = model.get(base_key_name, '_') base_key = f'{base_key_name}:{base_key_val}' if group_by: base_key += cls.get_group_by_multikey(model, group_by=group_by) data[base_key] = data.get(base_key, []) data[base_key].append(model) return cls._reduce_add_nested(data) @classmethod def _reduce_add_nested(cls, data: Dict[str, List[dict]], use_pool: bool = USE_THREAD_POOL_FOR_AGG) -> dict: """"Caller for :meth:`Models.add_nested_dicts` that will optionally use a threadpool see :const:`~delphi_api.const.USE_THREAD_POOL_FOR_AGG` Args: data: Data from one of the previous aggregation methods use_pool: Enable running via ``ThreadPoolExecutor`` Returns: dict: a dictionary of items with "multikeys" used for grouping """ items = {} if use_pool: with ThreadPoolExecutor(max_workers=MAX_THREADS) as executor: for key, val in data.items(): future = executor.submit(reduce, cls.add_nested_dicts, val, {}) items[key] = future.result() return items for key, val in data.items(): items[key] = reduce(cls.add_nested_dicts, val, {}) return items @classmethod def get_group_by_multikey(cls, model, group_by: Union[str, Iterable[str]]) -> str: """ Args: model: Child instance of abstract type :class:`BigTableModel` group_by: Iterable of field names for grouping results by (ex: ``date``) Returns: str: "packed" with ``group_by`` fields in format for unpacking later """ multikey = '' group_by = [group_by] if isinstance(group_by, str) else group_by for field in group_by: if not isinstance(field, str): continue if isinstance(model, dict): value = model.get(field) else: value = getattr(model, field, None) if not value: LOG.debug(f'Warning: Missing or empty field value for model: {field}={value}') continue # example: ';date:2020-01-01' multikey += f'{cls.MULTIKEY_SEP}{field}{cls.KEY_VAL_SEP}{value}' return multikey @classmethod def cast_proto_map(cls, proto_map: Union[MessageMapContainer, dict], only: str = None) -> dict: """ Args: proto_map: dict-like object (map) only: provide to filter for only one key in map (ex: a ``country_code``) Returns: Converted a map of protobuf fields as dictionary """ data = {} keys = proto_map.keys() if only: keys = {only} for key in keys: message = proto_map[key] message_dict = MessageToDict(message, preserving_proto_field_name=True) # INT64 types are serialized to strings instead of integers so we need to cast them # see: https://developers.google.com/protocol-buffers/docs/proto3#json data[key] = cls.cast_nested_str_to_int(message_dict) return data @classmethod def cast_nested_str_to_int(cls, data: dict): """Cast any string property values in nested dictionaries to integers""" new_data = {} for key, val in data.items(): if isinstance(val, dict): new_data[key] = cls.cast_nested_str_to_int(val) else: try: new_data[key] = int(val) except KeyError as e: new_data[key] = val return new_data @classmethod def add_nested_dicts(cls, a: dict, b: dict) -> dict: """Reducer function for adding nested values (summation of integers)""" if (a is None) ^ (b is None): # if a or b is null, but not a and b are null return a if a is not None else b elif a is None: # b is also null in this case return {} data = {} keyset = set(tuple(a.keys()) + tuple(b.keys())) for key in keyset: aval = a.get(key) bval = b.get(key) if isinstance(aval, dict) or isinstance(bval, dict): data[key] = cls.add_nested_dicts(aval, bval) elif isinstance(aval, str) or isinstance(bval, str): data[key] = aval else: try: data[key] = (aval or 0) + (bval or 0) or None except TypeError as e: LOG.warning(e) return data @classmethod def flatten(cls, data: dict, additional_data: dict = None) -> List[dict]: """ Args: data: dictionary with keys in "multikey" formatted strings additional_data: optional dictionary of additional key-value data to add to each object Returns: Flattened data as a list of objects with multi key-val as properties """ items = [] for multikey, country_obj in data.items(): new_obj = {} if not additional_data else additional_data # set keys from parents (isrc, date, etc.) for field in multikey.split(cls.MULTIKEY_SEP): key, val = field.split(cls.KEY_VAL_SEP) new_obj[key] = val for country, stats_obj in country_obj.items(): new_obj['country_code'] = country stats_obj.update(new_obj) items.append(stats_obj) return items @classmethod def flatten_generic(cls, data: dict, additional_data: dict = None) -> List[dict]: """Similar to :class:`Models.flatten` but used for more generic data (not country-based) Args: data: dictionary with keys in "multikey" formatted strings additional_data: optional dictionary of additional key-value data to add to each object Returns: Flattened data as a list of objects with multi key-val as properties """ items = [] for multikey, obj in data.items(): new_obj = {} if not additional_data else additional_data new_obj.update(obj) # set keys from parents (isrc, date, etc.) for field in multikey.split(cls.MULTIKEY_SEP): key, val = field.split(cls.KEY_VAL_SEP) new_obj[key] = val items.append(new_obj) return items @classmethod def flatmap(cls, func, *iterable): return itertools.chain.from_iterable(map(func, *iterable))