from airflow.models import Variable import requests import json from airflow.contrib.operators.ssh_operator import SSHOperator from airflow.contrib.operators.bigquery_operator import BigQueryOperator from airflow.hooks.base_hook import BaseHook from airflow.contrib.operators.dataflow_operator import DataflowTemplateOperator from airflow.operators.bash_operator import BashOperator from airflow.operators.python_operator import PythonOperator import gcs_storage import os # This is to get the spotify access token for a licensor def get_access_token(licensor): secret = BaseHook.get_connection('spotify_sony_conn').password oauth_url = Variable.get('spotify_oauth_base_api_url') + '&client_id=' + \ Variable.get(licensor['client_id_param']) + '&client_secret=' + \ secret print('Before calling OAuth url') r = requests.post(url = oauth_url) print('status ==', r.status_code) access_token = r.json()['access_token'] return access_token # This is to get licensor details for a given licensor name def get_licensor_details(name): licensor = None licensors = json.loads(Variable.get('licensors')) for ref_licensor in licensors: if ref_licensor['name'] == name: licensor = ref_licensor break return licensor # This is to check if the downloads are available for a given date def are_downloads_available(licensor_name, date_to_download): doNotUseAvailableCountryURL = Variable.get('spotify_do_not_use_available_country_url') licensor = get_licensor_details(licensor_name) access_token = get_access_token(licensor) downloads_available = False if doNotUseAvailableCountryURL == 'No': available_countries_streams_url = Variable.get('spotify_base_api_url') + '/' + \ licensor['name'] + 'v2/streams/' + date_to_download.replace('-', '/') + \ '?oauth_token=' + access_token r = requests.get(available_countries_streams_url) available_country_streams = r.json(); if len(available_country_streams) > 0: downloads_available = True else: tracks_url = Variable.get('spotify_base_api_url') + '/' + \ licensor['name'] + 'v2/tracks/' + date_to_download.replace('-', '/') + \ '?oauth_token=' + access_token r = requests.get(tracks_url) if r.status_code == 200: downloads_available = True else: streams_url = Variable.get('spotify_base_api_url') + '/' + \ licensor['name'] + 'v2/streams/' + date_to_download.replace('-', '/') + \ '/US?oauth_token=' + access_token r = requests.get(streams_url) if r.status_code == 200: downloads_available = True return downloads_available def get_streams_bq_operator(dag, task_id, licensor_name, countries_to_include, countries_to_exclude, date_to_download): where_clause = '' countries_to_consider = countries_to_include.split() count = 1 for country in countries_to_consider: filter = '"%spotify_streams_v2_' + licensor_name + '_' + date_to_download + '_' + country + '.json"' if count == 1: where_clause = where_clause + ' WHERE _FILE_NAME like ' + filter else: where_clause = where_clause + ' or _FILE_NAME like ' + filter count = count + 1 countries_not_to_consider = countries_to_exclude.split() count = 1 for country in countries_not_to_consider: filter = '"%spotify_streams_v2_' + licensor_name + '_' + date_to_download + '_' + country + '.json"' if count == 1: where_clause = where_clause + ' WHERE _FILE_NAME not like ' + filter else: where_clause = where_clause + ' and _FILE_NAME not like ' + filter count = count + 1 if where_clause == '': where_clause = ' WHERE _FILE_NAME like "%spotify_streams_v2_' + licensor_name + '_' + date_to_download + '_%' + '.json"' # Load from external table into BQ internal streams table bq_load_streams = BigQueryOperator( task_id=task_id, sql=''' INSERT `dsp-automation.spotify.spotify_streams` ( new_user_id, shuffle, os, repeat_play, device_type, message, version, user_id, track_id, completion_flag, content_type_flag, cached, source_uri, source, utc_timestamp_offset, length, timestamp, offline_timestamp, report_date, report_licensor, report_country ) SELECT * EXCEPT(timestamp, offline_timestamp, filename,country), SAFE.PARSE_DATETIME('%Y%m%dT%H:%M:%S', timestamp) AS timestamp, SAFE.PARSE_DATETIME('%Y%m%dT%H:%M:%S',offline_timestamp ) AS offline_timestamp, CAST(REGEXP_EXTRACT(filename, r'_([0-9]*-[0-9]*-[0-9]*)_') AS date) AS report_date, REGEXP_EXTRACT(filename, r'v[0-9]_([a-z]*)_') AS report_licensor, REGEXP_EXTRACT(filename, r'_([a-z]*)\.') AS report_country FROM ( SELECT *, _FILE_NAME AS filename FROM `dsp-automation.spotify.external_streams` ''' + where_clause + ''' ) ''', write_disposition='WRITE_APPEND', use_legacy_sql=False, dag=dag ) return bq_load_streams def get_tracks_bq_operator(dag, licensor_name, date_to_download): # Load from external table into BQ internal tracks table filter = '%spotify_tracks_v2_' + licensor_name + '_' + date_to_download + '.json' bq_load_tracks = BigQueryOperator( task_id='bqexternal_to_bqinternal_tracks_' + licensor_name, sql=''' INSERT `dsp-automation.spotify.spotify_tracks` ( album_name, track_name, track_artists, album_code, isrc, album_artist, uri, track_id, version, message, report_date, report_licensor ) SELECT * EXCEPT(filename), CAST(REGEXP_EXTRACT(filename, r'_([0-9]*-[0-9]*-[0-9]*)\.') AS date) AS report_date, REGEXP_EXTRACT(filename, r'v[0-9]_([a-z]*)_') AS report_licensor FROM ( SELECT *, _FILE_NAME AS filename FROM `dsp-automation.spotify.external_tracks` WHERE _FILE_NAME like "''' + filter + '''" ) ''', write_disposition='WRITE_APPEND', use_legacy_sql=False, dag=dag ) return bq_load_tracks def get_users_bq_operator(dag, licensor_name, date_to_download): # Load from external table into BQ internal users table filter = '%spotify_users_v2_' + licensor_name + '_' + date_to_download + '.json' bq_load_users = BigQueryOperator( task_id='bqexternal_to_bqinternal_users_' + licensor_name, sql=''' INSERT `dsp-automation.spotify.spotify_users` ( new_user_id, region, version, user_id, access, type, gender, product, message, partner, birth_year, referral, zipcode, country, report_date, report_licensor ) SELECT * EXCEPT(filename), CAST(REGEXP_EXTRACT(filename, r'_([0-9]*-[0-9]*-[0-9]*)\.') AS date) AS report_date, REGEXP_EXTRACT(filename, r'v[0-9]_([a-z]*)_') AS report_licensor FROM ( SELECT *, _FILE_NAME AS filename FROM `dsp-automation.spotify.external_users` WHERE _FILE_NAME like "''' + filter + '''" ) ''', write_disposition='WRITE_APPEND', use_legacy_sql=False, dag=dag ) return bq_load_users def set_move_task(dag, date, upstream_task, downstream_task, filters): for filter in filters: move_files = BashOperator( task_id=filter[0], bash_command='gsutil -m mv gs://' + Variable.get("gcs_bucket") + '/' + Variable.get("spotify_folder_in_bucket") + '/incoming/' + filter[1] + \ ' gs://' + Variable.get("gcs_bucket") + '/' + Variable.get("spotify_folder_in_bucket") + '/processed/' + date + '|| true', dag=dag ) move_files.set_upstream(upstream_task) move_files.set_downstream(downstream_task) def move_files(**kwargs): countries_to_ignore = ['US', 'MX', 'NL', 'GB', 'SE', 'DE', 'BR', 'TR', 'PL'] list = gcs_storage.get_list_blobs_with_partial_file_name(Variable.get("gcs_bucket"), Variable.get("spotify_folder_in_bucket") + '/incoming/', kwargs['filter']) for file in list: ignore_move =False for country in countries_to_ignore: if country in file: ignore_move =True break if not ignore_move: os.system('gsutil -m mv gs://' + Variable.get("gcs_bucket") + '/' + file + \ ' gs://' + Variable.get("gcs_bucket") + '/' + Variable.get("spotify_folder_in_bucket") + '/processed/' + kwargs['date'] + '|| true') def set_move_task_for_rest_of_the_countries(dag, date, upstream_task, downstream_task, filter): move_files_task = PythonOperator( task_id=filter[0], python_callable=move_files, op_kwargs={'date': date, 'filter':filter[1]}, provide_context=True, dag=dag ) move_files_task.set_upstream(upstream_task) move_files_task.set_downstream(downstream_task) # This function is for creating the download tasks and it spreads out the SSH tasks based on the country def create_download_tasks(date_to_download, dag, fork_dwnlds, wait_task, stop_vms_task): licensors = json.loads(Variable.get('licensors')) # Download the extracted files and spread out the downloads across 4 VMs for licensor in licensors: if licensor['name'] == 'sonybmgmusicentertainment': # 1 - (Only SonyBMG) This task creates download for only US stream file countries_include = 'US' licensor_task1 = SSHOperator( ssh_conn_id='remote_vm_conn', task_id=licensor['name'] + '_' + '_dwnlds_US', #command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \ #' -dd ' + date_to_download + ' -dt extracted_files -ci \'' + countries_include + \ #'\' -dncf False', command='java -cp collect-analytics-0.0.1-SNAPSHOT-jar-with-dependencies.jar com.sony.analytics.spotify.SpotifyDownloadExecutor ' + \ '-doNotUseAvailableCountryURL -downloadByLicensorAndCountry -licensorName ' + licensor['name'] + \ ' -downloadDate ' + date_to_download + ' -countries ' + countries_include + ' -doNotDownloadNonCountryFiles true', dag=dag) task_id_for_bq_operator = licensor['name'] + '_' + 'bqexternal_to_bqinternal_streams' + '_US' us_streams_bq_operator = get_streams_bq_operator(dag, task_id_for_bq_operator, licensor['name'], 'US', '', date_to_download) licensor_task1.set_upstream(fork_dwnlds) licensor_task1.set_downstream(us_streams_bq_operator) filter_str = 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'US' + '.json' filter = ('Move_US_file', filter_str) filters = [filter] set_move_task(dag, date_to_download, us_streams_bq_operator, wait_task, filters) # 2 - (Only SonyBMG) This task creates download for countries that have large stream files countries_include = 'GB SE DE BR TR PL NL MX' licensor_task2 = SSHOperator( ssh_conn_id='remote_vm_conn1', task_id=licensor['name'] + '_' + '_dwnlds_GB_SE_DE_BR_TR_PL_NL_MX', #command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \ #' -dd ' + date_to_download + ' -dt extracted_files -ci \'' + countries_include + \ #'\' -dncf False', command='java -cp collect-analytics-0.0.1-SNAPSHOT-jar-with-dependencies.jar com.sony.analytics.spotify.SpotifyDownloadExecutor ' + \ '-doNotUseAvailableCountryURL -downloadByLicensorAndCountry -licensorName ' + licensor['name'] + \ ' -downloadDate ' + date_to_download + ' -countries ' + countries_include + ' -doNotDownloadNonCountryFiles false', dag=dag) task_id_for_bq_operator = licensor['name'] + '_' + 'bqexternal_to_bqinternal_streams' + '_GB_SE_DE_BR_TR_PL_NL_MX' big_countries_streams_bq_operator = get_streams_bq_operator(dag, task_id_for_bq_operator, licensor['name'], countries_include, '', date_to_download) tracks_bq_operator = get_tracks_bq_operator(dag, licensor['name'], date_to_download) users_bq_operator = get_users_bq_operator(dag, licensor['name'], date_to_download) licensor_task2.set_upstream(fork_dwnlds) licensor_task2.set_downstream(big_countries_streams_bq_operator) licensor_task2.set_downstream(tracks_bq_operator) licensor_task2.set_downstream(users_bq_operator) filter1 = ('Move_GB_file', 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'GB' + '.json') filter2 = ('Move_SE_file', 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'SE' + '.json') filter3 = ('Move_DE_file', 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'DE' + '.json') filter4 = ('Move_BR_file', 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'BR' + '.json') filter5 = ('Move_TR_file', 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'TR' + '.json') filter6 = ('Move_PL_file', 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'PL' + '.json') filter7 = ('Move_NL_file', 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'NL' + '.json') filter8 = ('Move_MX_file', 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + 'MX' + '.json') filters = [filter1, filter2, filter3, filter4, filter5, filter6, filter7, filter8] set_move_task(dag, date_to_download, big_countries_streams_bq_operator, wait_task, filters) filter1 = ('Move_tracks_file', 'spotify_tracks_v2_' + licensor['name'] + '_' + date_to_download + '.json') set_move_task(dag, date_to_download, tracks_bq_operator, wait_task, [filter1]) filter1 = ('Move_users_file', 'spotify_users_v2_' + licensor['name'] + '_' + date_to_download + '.json') set_move_task(dag, date_to_download, users_bq_operator, wait_task, [filter1]) # 3 - (Only SonyBMG) This task creates download for rest of the countries # and all files including tracks, users & agg streams countries_exclude = 'US MX NL GB SE DE BR TR PL' licensor_task3 = SSHOperator( ssh_conn_id='remote_vm_conn2', task_id=licensor['name'] + '_' + '_dwnlds_REST_COUNTRIES', #command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \ #' -dd ' + date_to_download + ' -dt extracted_files -ce \'' + countries_exclude + \ #'\' -dncf True', command='java -cp collect-analytics-0.0.1-SNAPSHOT-jar-with-dependencies.jar com.sony.analytics.spotify.SpotifyDownloadExecutor ' + \ '-doNotUseAvailableCountryURL -downloadByLicensorWithExcludeCountries -licensorName ' + licensor['name'] + \ ' -downloadDate ' + date_to_download + ' -countries ' + countries_exclude + ' -doNotDownloadNonCountryFiles true', dag=dag) task_id_for_bq_operator = licensor['name'] + '_' + 'bqexternal_to_bqinternal_streams' + '_REST_COUNTRIES' rest_streams_bq_operator = get_streams_bq_operator(dag, task_id_for_bq_operator, licensor['name'], '', countries_exclude, date_to_download) licensor_task3.set_upstream(fork_dwnlds) licensor_task3.set_downstream(rest_streams_bq_operator) #rest_streams_bq_operator.set_downstream(wait_task) filter_str = 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download filter = ('Move_rest_of_the_countries_file', filter_str) set_move_task_for_rest_of_the_countries(dag, date_to_download, rest_streams_bq_operator, wait_task, filter) else: # 4 - (Orchard, sonymarketing & intl) This task creates download for all countries # and all files including tracks, users & agg streams licensor_task = SSHOperator( ssh_conn_id='remote_vm_conn3', task_id=licensor['name'] + '_' + '_dwnlds_ALL', #command='python3 spotify_multi_process_dwld_upld_v1_1.py -ln ' + licensor['name'] + \ #' -dd ' + date_to_download + ' -dt extracted_files', command='java -cp collect-analytics-0.0.1-SNAPSHOT-jar-with-dependencies.jar com.sony.analytics.spotify.SpotifyDownloadExecutor ' + \ '-doNotUseAvailableCountryURL -downloadByLicensor -licensorName ' + licensor['name'] + \ ' -downloadDate ' + date_to_download, dag=dag) task_id_for_bq_operator = licensor['name'] + '_' + 'bqexternal_to_bqinternal_streams' + '_ALL' all_licensors_streams_bq_operator = get_streams_bq_operator(dag, task_id_for_bq_operator, licensor['name'], '', '', date_to_download) licensor_task.set_upstream(fork_dwnlds) licensor_task.set_downstream(all_licensors_streams_bq_operator) #all_licensors_streams_bq_operator.set_downstream(wait_task) filter_str = 'spotify_streams_v2_' + licensor['name'] + '_' + date_to_download + '_' + '*' + '.json' filter = ('Move_'+ licensor['name'] + '_files', filter_str) filters = [filter] set_move_task(dag, date_to_download, all_licensors_streams_bq_operator, wait_task, filters) tracks_bq_operator = get_tracks_bq_operator(dag, licensor['name'], date_to_download) users_bq_operator = get_users_bq_operator(dag, licensor['name'], date_to_download) licensor_task.set_downstream(tracks_bq_operator) licensor_task.set_downstream(users_bq_operator) filter1 = ('Move_tracks_file_' + licensor['name'], 'spotify_tracks_v2_' + licensor['name'] + '_' + date_to_download + '.json') set_move_task(dag, date_to_download, tracks_bq_operator, wait_task, [filter1]) filter1 = ('Move_users_file' + licensor['name'], 'spotify_users_v2_' + licensor['name'] + '_' + date_to_download + '.json') set_move_task(dag, date_to_download, users_bq_operator, wait_task, [filter1]) # Download the archive files and spread out the downloads across 4 VMs lc_idx = 0 for licensor in licensors: connection_id = 'remote_vm_conn' if lc_idx > 0: connection_id = 'remote_vm_conn' + str(lc_idx) licensor_task = SSHOperator( ssh_conn_id=connection_id, task_id=licensor['name'] + '_' + '_gz_dwnlds', command='java -cp collect-analytics-0.0.1-SNAPSHOT-jar-with-dependencies.jar com.sony.analytics.spotify.SpotifyDownloadExecutor ' + \ '-doNotUseAvailableCountryURL -downloadByLicensor -downloadGzipsOnly -licensorName ' + licensor['name'] + \ ' -downloadDate ' + date_to_download, dag=dag) licensor_task.set_upstream(wait_task) licensor_task.set_downstream(stop_vms_task) agg_streams_to_bq = DataflowTemplateOperator( task_id='sony_spotify_agg_streams_load_bq', template='gs://' + Variable.get('dataflow-template-bucket') + '/templates/aggregate-streams-etl', gcp_conn_id='google_cloud_default', parameters={ 'inputFile': 'gs://' + Variable.get('gcs_bucket') + '/' + Variable.get('spotify_folder_in_bucket') +\ '/incoming/spotify_aggregatedstreams*' + date_to_download + '*', 'outputDataset': Variable.get('spotify_agg_streams_output_dataset') }, dag=dag) agg_streams_to_bq.set_upstream(wait_task) for licensor in licensors: filter = ('Move_agg_stream_files_' + licensor['name'], 'spotify_aggregatedstreams_v2_' + licensor['name'] + '_' + date_to_download + '.json') filters = [filter] set_move_task(dag, date_to_download, agg_streams_to_bq, stop_vms_task, filters)