from abc import ABC, abstractmethod from typing import Iterable, List, Optional 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 DEFAULT_PAGE_SIZE from delphi_api.utils import DateUtils, Models, QueryUtils from delphi_api.v2.constants import AMAZON_DSP_SLUGS, COLUMN_FAMILY_ID, UTF8 from delphi_api.v2.enums import ColumnName LOG = get_logger(__name__) class DspViewModel(ABC): def __init__(self, results: Iterable = None, params: Optional[dict] = None): """Abstract base class for DSP view models (legacy v2 Amazon-specific) Args: results: (optional) iterable of rows of data params: (optional) dictionary of parameters from a query Attributes: self.results: iterable of rows self.params: (optional) dictionary of parameters from a query """ self._results = results if results else [] self.params = params if params else {} self._group_by = self.params.get('group_by') self._limit = self.params.get('limit', DEFAULT_PAGE_SIZE) self._offset = self.params.get('offset', 0) self._playlist_id = self.params.get('playlist_id') self._base_key_name = self.params.get('base_key_name', 'isrc') self._num_dsps = len(AMAZON_DSP_SLUGS) @property @abstractmethod def data_model(self) -> BigTableModel: """Subclasses must implement this abstract property. Returns: BigTableModel: Instance of a :class:`BigTableModel` sub-type """ pass # pragma: no cover @property def items(self) -> List[dict]: """ Returns: List[dict]: List of dictionary aggregated items """ return self._items_aggregated[self.offset:self.limit + self.offset] @property def countries(self) -> Iterable[str]: """ Returns: Iterable[str]: the list of country codes """ return self.data_model.Meta.country_codes @property def results(self) -> List[BigTableModel]: return self._results @results.setter def results(self, value): self._results = value @property def group_by(self) -> Optional[str]: return self._group_by @group_by.setter def group_by(self, value): self._group_by = value @property def limit(self) -> Optional[int]: return self._limit @limit.setter def limit(self, value): self._limit = value @property def offset(self) -> Optional[int]: return self._offset @offset.setter def offset(self, value): self._offset = value @property def playlist_id(self) -> Optional[str]: return self._playlist_id @playlist_id.setter def playlist_id(self, value): self._playlist_id = value @property def base_key_name(self) -> Optional[str]: return self._base_key_name @base_key_name.setter def base_key_name(self, value): self._base_key_name = value @property def stats_fields(self) -> Iterable[str]: fields = ('country_stats',) return fields @property def _aggregated_nested_data(self) -> dict: try: return Models.aggregate_models(self.results, stats_fields=self.stats_fields, country=self.params.get('country_code'), base_key_name=self.base_key_name, group_by=self.group_by) except (AttributeError, TypeError) as e: # pragma: no cover LOG.exception(e) return {} @property def _items_aggregated(self) -> List[dict]: additional_data = None if self.playlist_id and not self.base_key_name == 'playlist_id': additional_data = {'playlist_id': self.playlist_id} try: flattened_counts = Models.flatten(self._aggregated_nested_data, additional_data=additional_data) return flattened_counts except (AttributeError, TypeError) as e: # pragma: no cover LOG.exception(e) return [] @property def count(self) -> int: return len(self.items) def to_dict(self) -> dict: """This method calls ``self.items`` which will aggregate and returns a result dictionary""" return { 'items': self.items, 'count': self.count, } def batch_load_tracks_streams(self, isrc: List[str], start_date: str, end_date: str, country_code: str = None, group_by: str = None, limit: int = DEFAULT_PAGE_SIZE, offset: int = 0): """Loads our view model instance with ``results`` by querying underlying data model""" DateUtils.valid_range(start_date, end_date) isrcs = frozenset(isrc) item_keys = QueryUtils.build_item_keys(isrcs) self.group_by = group_by self.limit = limit self.offset = offset model = self.data_model filter_ = None if country_code is not None: filter_ = self.get_column_filter(model, {country_code}) self.results = model.batch_get_range(item_keys, start_date=start_date, end_date=end_date, filter_=filter_, limit=limit * self._num_dsps) return self def batch_load_playlist_tracks_streams(self, isrc: List[str], playlist_id: str, start_date: str, end_date: str, country_code: str = None, group_by: str = None, limit: int = DEFAULT_PAGE_SIZE, offset: int = 0): """Loads our view model instance with ``results`` by querying underlying data model""" DateUtils.valid_range(start_date, end_date) isrcs = frozenset(isrc) item_keys = QueryUtils.build_item_keys(frozenset({playlist_id}), isrcs) self.playlist_id = playlist_id self.group_by = group_by self.limit = limit self.offset = offset model = self.data_model filter_ = None if country_code is not None: filter_ = self.get_column_filter(model, {country_code}) self.results = model.batch_get_range(item_keys, start_date=start_date, end_date=end_date, filter_=filter_, limit=limit * self._num_dsps) return self def batch_load_playlists(self, playlist_id: List[str], start_date: str, end_date: str, country_code: str = None, group_by: str = None, limit: int = DEFAULT_PAGE_SIZE, offset: int = 0): """Loads our view model instance with ``results`` by querying underlying data model""" DateUtils.valid_range(start_date, end_date) item_keys = QueryUtils.build_item_keys(frozenset(playlist_id)) self.group_by = group_by self.limit = limit self.offset = offset model = self.data_model filter_ = None if country_code is not None: filter_ = self.get_column_filter(model, {country_code}) self.results = model.batch_get_range(item_keys, start_date=start_date, end_date=end_date, filter_=filter_, limit=limit * self._num_dsps) return self @classmethod def get_column_filter(cls, model: object, column_names: set) -> RowFilterUnion: """ 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 """ whitelist = { 'isrc': ColumnName.ISRC.value, 'playlist_id': ColumnName.PLAYLIST_ID.value, 'date': ColumnName.DATE.value, } whitelist.update({c: c for c in column_names}) filters = [] for model_attr, col_name in whitelist.items(): if not hasattr(model, model_attr): continue bytes_name = bytes(col_name, UTF8) filters.append( ColumnRangeFilter(COLUMN_FAMILY_ID, start_column=bytes_name, end_column=bytes_name)) return RowFilterUnion(filters=filters)