import datetime from typing import Callable, Iterable, List, Optional, Set from apollo_main_db.apollo.models import StarredContent from apollo_main_db.apple.models import AppleMusicChart, AppleMusicChartSongEntry, AppleMusicSong from apollo_main_db.push_notifications.models import UserDeviceToken, VendorEnum from apollo_main_db.spotify.models import SpotifyCharts from marshmallow import Schema from sqlalchemy import and_, distinct, func from sqlalchemy.orm import Query, Session from apollo_notifications.charts.config import ChartsConfig from apollo_notifications.charts.utils import is_filtered from apollo_notifications.client import NotificationsBaseClient from apollo_notifications.constants import (APPLE_CHART_NAME, PUSH_SEARCH_MODE_TO_MODEL, SPOTIFY_URL_PREFIX, PushSearchMode) from apollo_notifications.utils import filter_by class NotificationsChartClient(NotificationsBaseClient): """Base class to get charts vendor-oriented data.""" vendor: str def __init__(self, config: ChartsConfig): super().__init__() self.topic = config.TOPIC self.push_search_mode = PushSearchMode(config.PUSH_SEARCH_MODE) def _get_markets(self) -> Iterable[str]: pass def get_markets(self, by: Iterable[str] = None) -> Iterable: """Get chart markets""" return filter_by(self._get_markets(), by) def filter_by_starred_and_active( self, tracks_sub_query: Query, user_id_list: List[str]) -> Query: """Filtering tracks query by starred tracks, join with User devices. Args: tracks_sub_query: subquery with track_id, position, track_name, artist_name aliased columns. user_id_list: list of user id to filter by. """ query = self.session.query( StarredContent.user_id, func.group_concat( UserDeviceToken.token.distinct().op('ORDER BY')(UserDeviceToken.created_at.desc()) ).label('device_tokens'), tracks_sub_query ).join( UserDeviceToken, UserDeviceToken.user_id == StarredContent.user_id ).join( tracks_sub_query, func.upper(StarredContent.isrc) == func.upper(tracks_sub_query.c.isrc) ).filter( UserDeviceToken.is_active.is_(True), UserDeviceToken.user_id.in_(user_id_list), StarredContent.isrc.isnot(None), StarredContent.isrc.op('regexp')(r'[A-Za-z]{2}[A-Za-z0-9]{3}[0-9]{7}') ) return query def get_existing_push_messages(self, push_date: str) -> Set[str]: """Return set of existing messages ids for specific date, topic and vendor. Args: push_date: str (Y-m-d format) date to generate push-messages for. """ model = PUSH_SEARCH_MODE_TO_MODEL[self.push_search_mode] q = self.session.query( model.id ).filter( model.date == push_date, model.topic == self.topic, model.vendor == self.vendor ) return set([str(pm.id) for pm in q]) def get_push_messages( self, query: Iterable, push_date: str, market: str, push_schema: Schema, *args, filter_function: Callable = is_filtered, existing_messages: Set[str] = None, **kwargs ): """Returns dumped filtered push messages.""" if existing_messages is None: existing_messages = self.get_existing_push_messages(push_date) messages = [] for track in query: if not filter_function( track, push_date, self.topic, market, self.vendor, existing_messages, *args, **kwargs): continue messages.append(push_schema.dump(track)) return messages class AppleChart(NotificationsChartClient): """Class to get Apple Music vendor-oriented data.""" vendor = VendorEnum.APPLE.name.lower() def filter_by_starred_and_active( self, tracks_query: Query, user_id_list: List[str]) -> Query: """Filtering tracks query by starred tracks, join with User devices, exclude already sent pushes. Args: tracks_query: Query with track_id, position, track_name, artist_name aliased columns. user_id_list: list of user id to filter by. """ tracks_sub_query = tracks_query.subquery() q = super().filter_by_starred_and_active(tracks_sub_query, user_id_list)\ .group_by(tracks_sub_query.c.track_id, StarredContent.user_id) return q def _get_markets(self) -> Iterable[str]: """Get chart markets""" return [i.lower() for i, in self.session.query(distinct(AppleMusicChart.storefront))] def _get_base_query(self, session: Session, market: str, date: str) -> Query: """Get base song entities query Args: session (Session): DB session. market: str market. date: str date in format %Y-%m-%d. """ return session.query( AppleMusicChartSongEntry.song_id.label('track_id'), AppleMusicChartSongEntry.name.label('track_name'), AppleMusicChartSongEntry.artist_name.label('artist_name'), (AppleMusicChartSongEntry.position + 1).label('position'), AppleMusicSong.isrc.label('isrc'), ).join( AppleMusicChart, AppleMusicChart.id == AppleMusicChartSongEntry.chart_id ).join( AppleMusicSong, AppleMusicSong.id == AppleMusicChartSongEntry.song_id ).filter( AppleMusicChart.chart_name == APPLE_CHART_NAME, AppleMusicChart.storefront == market, AppleMusicChart.date == date, ) def get_diff_tracks_by_dates(self, market: str, start_date: str, end_date: str) -> Query: """Get song entities difference between two dates. Using this method we can get added or removed tracks in charts. Args: market: str market. start_date: str date in format %Y-%m-%d. end_date: str date in format %Y-%m-%d. """ query_for_yesterday = self._get_base_query(self.session, market, end_date).subquery() query_for_today = self._get_base_query(self.session, market, start_date) return query_for_today.outerjoin( query_for_yesterday, query_for_yesterday.c.track_id == AppleMusicChartSongEntry.song_id ).filter(query_for_yesterday.c.track_id.is_(None)) def get_dates(self, market: str) -> List[Optional[datetime.date]]: """Get two most recent dates of SpotifyChart chart data. Args: market: str market. Returns: List of up to two dates in format %Y-%m-%d. """ dates = self.session.query(distinct(AppleMusicChart.date)).filter_by( storefront=market, chart_name=APPLE_CHART_NAME ).order_by(AppleMusicChart.date.desc())[:2] try: return [d.strftime("%Y-%m-%d") for d, in dates] except AttributeError: return [] class SpotifyChart(NotificationsChartClient): """Class to get Spotify vendor-oriented data.""" vendor = VendorEnum.SPOTIFY.name.lower() def filter_by_starred_and_active( self, tracks_query: Query, user_id_list: List[str]) -> Query: """Filtering tracks query by starred tracks, join with User devices, exclude already sent pushes. Args: tracks_query: Query with track_id, position, track_name, artist_name aliased columns. user_id_list: list of user id to filter by. """ tracks_sub_query = tracks_query.subquery() q = super().filter_by_starred_and_active(tracks_sub_query, user_id_list) \ .group_by( tracks_sub_query.c.isrc, tracks_sub_query.c.position, StarredContent.user_id ) return q def _get_markets(self) -> Iterable[str]: """Get chart markets""" return [i.lower() for i, in self.session.query(distinct(SpotifyCharts.market)).filter( and_(SpotifyCharts.market != '', SpotifyCharts.market.isnot(None)))] def _get_base_query(self, session: Session, market: str, date: str) -> Query: """Get base song entities query Args: session: DB session. market: str market. date: str date in format %Y-%m-%d. """ return session.query( func.REPLACE(SpotifyCharts.track_url, SPOTIFY_URL_PREFIX, '').label('track_id'), SpotifyCharts.track_name.label('track_name'), SpotifyCharts.artist_name.label('artist_name'), SpotifyCharts.position.label('position'), SpotifyCharts.isrc.label('isrc'), ).filter( SpotifyCharts.time_window_id == 0, SpotifyCharts.list_type_id == 0, SpotifyCharts.market == market, SpotifyCharts.date == date, ) def get_diff_tracks_by_dates(self, market: str, start_date: str, end_date: str) -> Query: """Get track entities difference between two dates. Using this method we can get added or removed tracks in charts. Args: market: str market. start_date: str date in format %Y-%m-%d. end_date: str date in format %Y-%m-%d. """ query_for_yesterday = self._get_base_query(self.session, market, end_date).subquery() query_for_today = self._get_base_query(self.session, market, start_date) return query_for_today.outerjoin( query_for_yesterday, query_for_yesterday.c.isrc == SpotifyCharts.isrc ).filter(query_for_yesterday.c.position.is_(None)) def get_dates(self, market: str) -> List[Optional[datetime.date]]: """Get two most recent dates of SpotifyChart chart data. Args: market: str market. Returns: List of up to two dates in format %Y-%m-%d. """ dates = self.session.query(distinct(SpotifyCharts.date)).filter_by( market=market, time_window_id=0, list_type_id=0 ).order_by(SpotifyCharts.date.desc())[:2] try: return [d.strftime("%Y-%m-%d") for d, in dates] except AttributeError: return []