"""Model for streams by source of streams.""" 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 import streams as streams_utils from sound_recordings.utils.cache import cache_in_redis SOUND_RECORDING_SOS_STREAMS_FIELDS = [ "isrc", "date", "streams_passive", "streams_active", "streams_collection", "streams", ] SQLLoader = snowflake.SQLLoader(__file__) @tracer.wrap(name="get_streams_by_source_of_streams") @cache_in_redis(ttl=cache.SECONDS_PER_HOUR) def get_streams_by_source_of_streams( permissions_filter, isrc, distributors, countries=[], store_ids=[], start_date=None, end_date=None, ): """Get streams by source of streams timeseries for 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 countries (list): List of country codes to filter by store_ids (list): List of store ids to filter by start_date (datetime.date): Start date end_date (datetime.date): End date Returns: List of dictionaries with keys: source (str): Source of streams date (str): Date in YYYY-MM-DD format streams (int): Total streams for date """ 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 [] params = { **permissions_filter, "isrc": isrc, "distributors": distributors, "store_ids": store_ids, "start_date": start_date, "end_date": end_date, } query_table = "streams_by_sos" if len(countries) > 0: params["country_codes"] = countries query_table = "streams_by_sos_by_country" sql = SQLLoader.load_query(query_table) sql = format_sql.format_with_permissions_filter( sql, permissions_filter, start_date=start_date, end_date=end_date ) records = snowflake.fetchall(sql, params) formatted_records = streams_utils.format_streams_by_sos_model( records, SOUND_RECORDING_SOS_STREAMS_FIELDS ) return formatted_records