from typing import Iterable, List, Optional, Tuple, Type from google.cloud.bigtable.row_filters import ColumnRangeFilter, RowFilterUnion from structlog import get_logger from delphi_api.bigtable import BigTableModel from delphi_api.const import BIGTABLE_MAX_ROWS_READ from delphi_api.core.constants import UTF8 from delphi_api.utils import DateUtils, Models from delphi_api.v3.enums import AggByParam, ColumnFamily, IncludeParamStreams, SubsetParam from delphi_api.v3.view_models.params import Params LOG = get_logger(__name__) class GenericDspViewModel: def __init__(self, params: Params, base_key_name: str, data_model: Type[BigTableModel] = None, results=None, stats_fields: Iterable[str] = ('streams',)): """This class is utilized closely with :class:`StreamsViewModel` Refactor from :class:`delphi_api.v2.view_models.dsp.DspViewModel` includes: - data model is passed as a param instead of an abstract property - simplified logic, params, and consolidated extraneous methods Args: params: Container of request params and any additional API loaded params base_key_name: Base field name for the aggregation (typically ISRC/track_id/playlist_id) data_model: (optional) Pass either a data_model or results for a grouped model results: (optional) Pass either ``data_model`` or ``results`` for a grouped model stats_fields: (optional) Specify the fields that will be used for metrics/aggregations """ self.params = params self.base_key_name = base_key_name self.data_model = data_model self._results = results if results is not None else [] self._stats_fields = stats_fields @property def model_instance(self) -> BigTableModel: return self.data_model() if self.data_model else self.data_model @property def results(self) -> List[BigTableModel]: return self._results @results.setter def results(self, value: List[BigTableModel]): self._results = value @property def items(self) -> List[dict]: """ Returns: List[dict]: List of dictionary aggregated items """ return self.get_aggregate_data_flat() @property def countries(self) -> Iterable[str]: """ Returns: Iterable[str]: the list of country codes """ return self.data_model.Meta.country_codes @property def row_limit(self) -> int: # we cannot limit results accurately with bigtable based on possible query sizes and ranges return BIGTABLE_MAX_ROWS_READ @property def stats_fields(self) -> Iterable[str]: fields = self._stats_fields if self.params.include == IncludeParamStreams.DEMOGRAPHICS.value: fields += ('demographics',) return fields @property def group_by_fields(self) -> Iterable[str]: group = self.params.group_by or [] if isinstance(self.params.group_by, str): group = [self.params.group_by] group_bys = set(group) if self.params.playlist_id and len(self.params.playlist_id) == 1: group_bys.add('playlist_id') elif self.params.subset == SubsetParam.PLAYLISTS.value: group_bys.add('playlist_id') if self.params.dsp and len(self.params.dsp) == 1: group_bys.add('dsp') if self.params.agg_by == AggByParam.ISRC.value: group_bys.add('isrc') elif self.params.agg_by == AggByParam.ARTIST.value: group_bys.add('artist_id') group_bys.discard(self.base_key_name) return list(group_bys) @property def additional_data(self) -> dict: """Some data we know from linked queries but does not exist in bigtable results""" data = {} if (self.params.playlist_id and not self.base_key_name == 'playlist_id' and len(self.params.playlist_id) == 1): data.update({'playlist_id': self.params.playlist_id[0]}) if self.params.dsp and len(self.params.dsp) == 1: data.update({'dsp': self.params.dsp[0]}) if self.params.artist_id and isinstance(self.params.artist_id, str): data.update({'artist_id': self.params.artist_id}) return data def get_aggregate_data_nested(self) -> dict: """See :meth:`delphi_api.utils.models.Models.aggregate_models`""" country_code = self.params.country_code[0] if self.params.country_code else None try: return Models.aggregate_models(self.results, base_key_name=self.base_key_name, country=country_code, stats_fields=self.stats_fields, group_by=self.group_by_fields) except (AttributeError, TypeError) as e: # pragma: no cover LOG.exception(e) return {} def get_aggregate_data_flat(self) -> List[dict]: """See :meth:`delphi_api.utils.models.Models.flatten`""" try: flattened_counts = Models.flatten(self.get_aggregate_data_nested(), additional_data=self.additional_data) return flattened_counts except (AttributeError, TypeError) as e: # pragma: no cover LOG.exception(e) return [] @classmethod def get_column_filters(cls, model: object, column_names: frozenset) -> Optional[RowFilterUnion]: """Creates a row filter to only return columns specified Args: model: The model to check if an attribute exists before including it in the query column_names: One or more column names to include in the results (ex: a country code) Returns: RowFilterUnion: RowFilter including the always included columns, and any passed columns """ filters = [ColumnRangeFilter(ColumnFamily.META.value)] for col_name in column_names: if not hasattr(model, col_name): LOG.warning('Unable to get model attribute for column name: %s' % col_name) continue # model_attr = getattr(model, col_name) bytes_name = bytes(col_name, UTF8) filters.append( ColumnRangeFilter(ColumnFamily.METRICS.value, start_column=bytes_name, end_column=bytes_name)) if not len(filters) > 1: return None return RowFilterUnion(filters=filters) def bigtable_load_results(self, item_keys: List[Tuple[bytes, bytes]]): """Loads our view model instance ``results`` by querying via the underlying data models Args: item_keys: A list of tuples in format (start_key, end_key) Returns: GenericDspViewModel: Instance (self) with ``results`` property loaded """ params = self.params model = self.model_instance if not model: raise AttributeError( 'Missing data model on instance. Ensure it is passed before calling load results') DateUtils.valid_range(params.start_date, params.end_date) filter_ = None self.results = model.batch_get_range_multiset(item_keys, filter_, self.row_limit) return self