"""Source of Streams model.""" 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 SOS_FIELDS = [ "isrc", "streams_passive_7_days", "streams_passive_7_days_growth", "streams_active_7_days", "streams_active_7_days_growth", "streams_collection_7_days", "streams_collection_7_days_growth", "streams_7_days", ] DYNAMIC_SOS_FIELDS = [ "isrc", "streams_passive", "streams_active", "streams_collection", "streams", ] SQLLoader = snowflake.SQLLoader(__file__) @tracer.wrap(name="get_sos_breakdown") @cache_in_redis(ttl=cache.SECONDS_PER_HOUR) def get_sos_breakdown(permissions_filter, isrc, distributors): """Get source of streams breakdown for given ISRC. Args: permissions_filter (dict): dict containing resources users can access isrc (str): ISRC of track to fetch breakdown for distributors (list): list of distributors names Returns: dict: SOS breakdown """ sql = SQLLoader.load_query("source_of_streams") sql = format_sql.format_with_permissions_filter(sql, permissions_filter) breakdown = snowflake.fetchone( sql, {**permissions_filter, "isrc": isrc, "distributors": distributors} ) if not breakdown: return None return dict(zip(SOS_FIELDS, breakdown)) @tracer.wrap(name="get_sos_breakdown") @cache_in_redis(ttl=cache.SECONDS_PER_HOUR) def get_dynamic_sos_breakdown( permissions_filter, isrc, distributors, countries=[], store_ids=[], start_date=None, end_date=None, ): """Get source of streams breakdown for given ISRC. Args: permissions_filter (dict): dict containing resources users can access isrc (str): ISRC of track to fetch breakdown 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: dict: SOS breakdown """ if len(store_ids) == 0: store_ids = store_availability.get_store_ids() else: store_ids = list( set(store_ids).intersection(store_availability.get_store_ids()) ) params = { **permissions_filter, "isrc": isrc, "store_ids": store_ids, "start_date": start_date, "end_date": end_date, "distributors": distributors, } if len(countries) > 0: params["country_codes"] = countries query_table = "source_of_streams_dynamic_by_country" else: query_table = "source_of_streams_dynamic" sql = SQLLoader.load_query(query_table) sql = format_sql.format_with_permissions_filter( sql, permissions_filter, start_date=start_date, end_date=end_date ) breakdown = snowflake.fetchone(sql, params) if not breakdown: return None return dict(zip(DYNAMIC_SOS_FIELDS, breakdown))