"""Model for top countries streams or downloads.""" from ddtrace import tracer from sound_recordings.connectors import snowflake from sound_recordings.constants import cache from sound_recordings.utils import format_sql, store_availability from sound_recordings.utils.cache import cache_in_redis SOUND_RECORDING_COUNTRIES_AGGREGATE_STREAMS_FIELDS = ["country_code", "streams"] SOUND_RECORDING_COUNTRIES_DOWNLOADS_FIELDS = ["country_code", "downloads"] SQLLoader = snowflake.SQLLoader(__file__) def _get_top_countries_result( permissions_filter, isrc, query, fields, distributors, store_ids=[], start_date=None, end_date=None, ): params = { **permissions_filter, "isrc": isrc, "distributors": distributors, "store_ids": store_ids, "start_date": start_date, "end_date": end_date, } sql = SQLLoader.load_query(query) sql = format_sql.format_with_permissions_filter( sql, permissions_filter, start_date=start_date, end_date=end_date ) records = snowflake.fetchall(sql, params) return [dict(zip(fields, record)) for record in records] @tracer.wrap(name="get_top_countries_streams") @cache_in_redis(ttl=cache.SECONDS_PER_HOUR) def get_top_countries_aggregate_streams( permissions_filter, isrc, distributors, store_ids=[], start_date=None, end_date=None ): """Get top countries aggregate streams for a given ISRC. Args: permissions_filter (dict): dict containing resources users can access isrc (str): ISRC of track to fetch downloads for. distributors (list): List of distributors names store_ids (list): List of store ids to filter by start_date (datetime.date): Start date end_date (datetime.date): End date Returns: list: top countries by stream count for sound recording """ if len(store_ids) == 0: store_ids = store_availability.get_store_ids() else: store_ids = sorted( list(set(store_ids).intersection(store_availability.get_store_ids())) ) if len(store_ids) == 0: return [] return _get_top_countries_result( permissions_filter, isrc, "top_countries_aggregate_streams", SOUND_RECORDING_COUNTRIES_AGGREGATE_STREAMS_FIELDS, distributors, store_ids, start_date, end_date, ) @tracer.wrap(name="get_top_countries_downloads") @cache_in_redis(ttl=cache.SECONDS_PER_HOUR) def get_top_countries_downloads( permissions_filter, isrc, distributors, store_ids=[], start_date=None, end_date=None ): """Get top countries downloads for a given ISRC. Args: permissions_filter (dict): dict containing resources users can access isrc (str): ISRC of track to fetch downloads for. distributors (list): List of distributors names store_ids (list): List of store ids to filter by start_date (datetime.date): Start date end_date (datetime.date): End date offset (int): Pagination offset limit (int): Pagination limit Returns: list: top countries by download count for sound recording """ if len(store_ids) == 0: store_ids = store_availability.get_download_store_ids() else: store_ids = sorted( list( set(store_ids).intersection(store_availability.get_download_store_ids()) ) ) if len(store_ids) == 0: return [] return _get_top_countries_result( permissions_filter, isrc, "top_countries_downloads", SOUND_RECORDING_COUNTRIES_DOWNLOADS_FIELDS, distributors, store_ids, start_date, end_date, )