"""Aggregate missing playlists history data and update existing.""" from datetime import date, timedelta import math from threading import Thread from typing import Dict import sentry_sdk from sentry_sdk.utils import BadDsn from athena import AthenaClient import config import constants.common as consts import constants.sql as sql_consts from logger import logger from main_db import MainDB from redis_db import get_redis_lock from s3 import S3Client import utils try: sentry_sdk.init(dsn=config.SENTRY_DSN, environment=config.ENVIRONMENT) except BadDsn: pass @utils.timing() def drop_excess_data(main_db: MainDB, vendor: str, last_top_date: date) -> int: """Drop playlists history data that is not in top anymore. Args: main_db (MainDB): DB. vendor (str): Vendor name. last_top_date (date): Last playlists top date. Returns: int: Rows count. """ row_count = 0 while True: logger.debug(f'{vendor}: getting non top playlists chunk from summary') non_top_playlists = main_db.get_non_top_playlists_history(vendor, last_top_date, config.DEFAULT_BATCH_LIMIT) if non_top_playlists: logger.debug(f'{vendor}: deleting non top playlists summary records') row_count = row_count + main_db.delete_history_by_playlists(vendor, non_top_playlists) if len(non_top_playlists) < config.DEFAULT_BATCH_LIMIT: logger.debug(f'{vendor}: {row_count} excess records were found and removed from summary') return row_count @utils.timing(3) def add_missing_data( main_db: MainDB, athena_client: AthenaClient, s3_client: S3Client, vendor: str, last_top_date: date) -> int: """Add missing top playlists history data. Args: main_db (MainDB): DB. athena_client (AthenaClient): Athena client. s3_client (S3Client): S3 client. vendor (str): Vendor name. last_top_date (date): Last playlists top date. Returns: int: Rows count. """ logger.debug(f'{vendor}: getting missing playlists') missing_playlists = main_db.get_missing_playlists(vendor, last_top_date) if not missing_playlists: return 0 logger.debug(f'{vendor}: {len(missing_playlists)} playlists to aggregate') chunks_count = math.ceil(len(missing_playlists) / config.PLAYLIST_CHUNK_SIZE) count = 1 result = 0 try: athena_client.create_table(vendor, '') for playlists_chunk in utils.split_chunks(missing_playlists, config.PLAYLIST_CHUNK_SIZE): max_playlist = utils.max_playlist(playlists_chunk) row_id = max_playlist if isinstance(max_playlist, str) else "-".join(max_playlist) logger.debug(f'{vendor}: ({count}/{chunks_count}) {row_id}') tr_playlist_id = utils.trim_playlist_id(row_id) athena_client.update_table(vendor, tr_playlist_id) upload_path = s3_client.get_full_path(vendor, row_id, consts.FOLDER_UPLOAD) download_path = s3_client.get_full_path(vendor, row_id, consts.FOLDER_DOWNLOAD) logger.debug(f'{vendor}: uploading') main_db.upload_to_s3(vendor, playlists_chunk, upload_path, s3_client) logger.debug(f'{vendor}: aggregating') athena_client.execute_aggregate_query(vendor, tr_playlist_id) input_s3_path = s3_client.get_relative_path(vendor, tr_playlist_id, consts.FOLDER_UPLOAD) output_s3_path = s3_client.get_relative_path(vendor, tr_playlist_id, consts.FOLDER_DOWNLOAD) s3_client.delete_s3_files_except_csv(output_s3_path) s3_client.delete_folder_files(input_s3_path) logger.debug(f'{vendor}: downloading') result = result + main_db.download_from_s3(vendor, download_path) s3_client.delete_folder_files(output_s3_path) count = count + 1 finally: athena_client.drop_table(vendor) s3_client.delete_temp() return result @utils.timing() def update_vendor_one_date(main_db: MainDB, vendor: str, last_top_date: date, current_date: date) -> Dict[str, int]: """Update vendor data for specific date. Args: main_db (MainDB): DB. vendor (str): Vendor name. last_top_date (date): Last playlists top date. current_date (date): Date to update for. Returns: Dict[str, int]: Command executing statistics, changed rows per command type. """ logger.debug(f'{vendor}: {current_date}') logger.debug(f'{vendor}: creating table') main_db.drop_temp_table(vendor) main_db.exec_command( sql_consts.SQL_CREATE_TEMP_TABLE[vendor], replacements=dict(temp_table=sql_consts.TABLE_HISTORY_TEMP[vendor])) results = {consts.COMMAND_SELECT: 0, consts.COMMAND_UPDATE: 0, consts.COMMAND_INSERT: 0} prev_playlist = '' if vendor == consts.VENDOR_SPOTIFY else ('', '') count = 1 logger.debug(f'{vendor}: getting changes') while True: playlists = main_db.get_playlists(vendor, last_top_date, prev_playlist, config.PLAYLIST_CHUNK_SIZE) if playlists: max_playlist = utils.max_playlist(playlists) logger.debug(f'{vendor}: ({count}) {prev_playlist}-{max_playlist}') rows_count = main_db.select_history_chunk(vendor, current_date, playlists) results[consts.COMMAND_SELECT] = results[consts.COMMAND_SELECT] + rows_count logger.debug(f'{vendor}: rows count = {rows_count}') prev_playlist = max_playlist if count % config.EXEC_CHANGE_CHUNK_COUNT == 0 or len(playlists) < config.PLAYLIST_CHUNK_SIZE: logger.debug(f'{vendor}: updating') end_date = current_date - timedelta(days=1) start_date = end_date - timedelta(days=config.DAYS_COUNT_FIX) results[consts.COMMAND_UPDATE] = main_db.exec_command( sql_consts.SQL_HISTORY_UPDATE[vendor], arguments=(current_date, start_date, end_date), replacements=dict(temp_table=sql_consts.TABLE_HISTORY_TEMP[vendor]), ) logger.debug(f'{vendor}: inserting') prev_date = current_date - timedelta(days=(1 + config.DAYS_COUNT_FIX)) results[consts.COMMAND_INSERT] = main_db.exec_command( sql_consts.SQL_HISTORY_INSERT[vendor], arguments=(current_date, current_date, prev_date, current_date), replacements=dict(temp_table=sql_consts.TABLE_HISTORY_TEMP[vendor]), ) logger.debug(f'{vendor}: truncating') main_db.exec_command( sql_consts.SQL_TRUNCATE_TEMP_TABLE, replacements=dict(history_temp=sql_consts.TABLE_HISTORY_TEMP[vendor]) ) if len(playlists) < config.PLAYLIST_CHUNK_SIZE: break count = count + 1 logger.debug(f'{vendor}: dropping table') main_db.drop_temp_table(vendor) return results def update_vendor(vendor: str): """Update vendor data for each date from range. Args: vendor (str): Vendor name. """ main_db = MainDB() last_top_date = main_db.get_top_playlists_last_date(vendor) logger.debug(f'{vendor}: playlists top max date {last_top_date}') results = {} s3_client = S3Client() if config.AGGREGATE_MISSING_PLAYLISTS and ( config.CHECK_MISSING_DAY < 0 or config.DATE_TO.weekday() == config.CHECK_MISSING_DAY): athena_client = AthenaClient() if config.ATHENA_TRY_CREATE_DATABASE: athena_client.create_database() results[consts.COMMAND_DELETE] = drop_excess_data(main_db, vendor, last_top_date) results[consts.COMMAND_INSERT] = add_missing_data(main_db, athena_client, s3_client, vendor, last_top_date) logger.info(f'{vendor}: aggregate stats {results}') if config.UPDATE_AGGREGATED_DATA: current_date = min(config.DATE_FROM, s3_client.get_last_date(vendor, config.DATE_FROM)) while current_date <= config.DATE_TO: stats = update_vendor_one_date(main_db, vendor, last_top_date, current_date) logger.info(f'{vendor}: update {current_date} stats {stats}') s3_client.set_last_date(vendor, current_date) current_date = current_date + timedelta(days=1) def process_history_data(): """Aggregate history data for missing playlists and update. """ lock = None try: if config.REDIS_LOCK_ENABLED: # set lock to avoid possible race condition lock = get_redis_lock() if not lock.acquire(blocking=False): logger.info('Can not get lock') return logger.debug(f'Session ID {config.CURRENT_RUN_ID}') if config.CURRENT_VENDOR: update_vendor(config.CURRENT_VENDOR) else: vendor_threads = [] for vendor in consts.VENDORS: thread = Thread(target=update_vendor, args=(vendor,)) thread.start() vendor_threads.append(thread) for thread in vendor_threads: thread.join() except Exception as ex: sentry_sdk.capture_exception(ex) raise finally: if config.REDIS_LOCK_ENABLED and lock: lock.release() if __name__ == '__main__': process_history_data()