""" Tasks for generating Schematized files from Accounting Statement Export ======================================================================== """ import hashlib import time from garcon import task import smart_open from snowflake_connector.etl_connector import SnowflakeSQLExecutor from processing_accounting.conf.settings import SF_CONFIG from processing_accounting.flows.accounting_statement_export import setting from processing_accounting.flows.accounting_statement_export.sql_generator \ import SqlGenerator from processing_accounting.util import db as db_util from processing_accounting.util import dynamodb as dynamodb_util def _hashFor(data): """Generate random unique key from the data Args: data (str): data to be hashed Returns: str: hashed data """ hashId = hashlib.md5() hashId.update(repr(data).encode('utf-8')) return hashId.hexdigest() @task.decorate(timeout=100) def bootstrap( activity, period_ids, user_type, user_id, transaction_types, payment_interval): """Bootstrap workflow by getting the correct configurations Args: activity (ActivityWorker): the activity worker. period_ids (str): comma delimited string of accounting period ids user_type (str): label, distributor or subaccount user_id (str): labelid or subaccountid transaction_types (str): comma delimited string of transaction types code payment_interval (str): payment interval in the contract (month or quarter) Returns: dict: context; """ assert period_ids and user_type, ( 'Missing required params: (period_ids, user_type).' 'Example: period_ids=203,204 user_type=subaccount/label') if payment_interval is None: payment_interval = 'month' if transaction_types is None: transaction_types = 'all' if user_id is None: user_id = 'all' if payment_interval == 'month': assert len(str(period_ids).split(',')) == 1, ( 'Invalid period_ids {} for payment_interval "month"').format( period_ids) params = '{}_{}_{}_{}_{}'.format( period_ids, user_type, user_id, transaction_types, payment_interval) hive_table_key = _hashFor(params) if user_id != 'all': raw_files_path = setting.config.get('raw_files_path').format( params=params) else: raw_files_path = setting.config.get( 'cross_labels_raw_files_path').format( params=params) track_artists_table_name = setting.TRACK_ARTISTS_AGGREGATION_TABLE_NAME resp = dict( raw_files_path=raw_files_path, hive_table_key=hive_table_key, conversion_hive_step_input=raw_files_path, schematized_file_path=( setting.config.get( 'schematized_file_path').format(params=params)), payment_interval=payment_interval, transaction_types=transaction_types, hdfs_staging_path=setting.hdfs_staging_path.format( params=_hashFor(params)), user_id=user_id, temp_track_artists_table_name=track_artists_table_name, ) resp['emr.logs'] = setting.config.get('emr').get('emr.logs') return resp @task.decorate(timeout=7000) def get_unload_query( activity, period_ids, user_type, user_id, transaction_types, payment_interval): """Get SQL statement for unloading data from a data warehouse to s3 Args: activity (ActivityWorker): the activity worker. period_ids (str): comma delimited string of accounting period ids user_type (str): label, distributor or subaccount user_id (str): labelid or subaccountid transaction_types (str): comma delimited string of transaction types code payment_interval (str): payment interval in the contract (month or quarter) Returns: dict: unload SQL statement with key unload_query """ generator = SqlGenerator( period_ids, user_type, payment_interval, user_id, transaction_types) sql = generator.get_sql() count_query = 'select count(*) from ({})'.format(sql) stop_response = {'stop': True, 'message': 'No data found.'} with SnowflakeSQLExecutor(sf_config=SF_CONFIG) as executor: result = executor.fetchone(count_query, dict_cursor=True) if int(result.get('COUNT(*)')) == 0: return stop_response return {'unload_query': sql} @task.decorate(timeout=7000) def snowflake_unload( activity, query, s3_prefix, aws_access_key, aws_access_secret): """Unload data from Snowflake. Args: activity (ActivityWorker): activity worker. query (str): SELECT query to use in unload statement. s3_prefix (str): S3 path to unload data to. aws_access_key (str): AWS credentials. aws_access_secret (str): AWS credentials. """ with SnowflakeSQLExecutor(sf_config=SF_CONFIG) as executor: unload_sql_template = """ COPY INTO %(s3_prefix)s FROM ({query}) CREDENTIALS=( AWS_KEY_ID=%(aws_access_key)s AWS_SECRET_KEY=%(aws_access_secret)s ) FILE_FORMAT = ( COMPRESSION=GZIP FIELD_DELIMITER='0x07' EMPTY_FIELD_AS_NULL=FALSE ) OVERWRITE=TRUE ;""" unload_sql_template = unload_sql_template.format(query=query) executor.execute( unload_sql_template, params={ 's3_prefix': s3_prefix + 'data_', 'aws_access_key': aws_access_key, 'aws_access_secret': aws_access_secret}) @task.decorate(timeout=7000) def generate_hql(activity, hash_key): """Generate HQL """ hql_config = setting.config.get('emr').get('hql') hql_path = setting.config.get('emr').get('hql_path').format(key=hash_key) with smart_open.smart_open(hql_path, 'wb') as fout: fout.write(hql_config.get('setting')) fout.write( hql_config.get('drop_accounting_statement_export').format( key=hash_key)) fout.write('\n') fout.write( hql_config.get('create_accounting_statement_export').format( key=hash_key)) fout.write('\n') fout.write( hql_config.get('drop_accounting_statement_export_avro').format( key=hash_key)) fout.write('\n') fout.write(hql_config.get('create_avro_table').format(key=hash_key)) fout.write('\n') fout.write(hql_config.get('write_avro_files').format(key=hash_key)) return { 'hql_path': hql_path } @task.decorate(timeout=7000) def update_status( activity, period_ids, user_type, payment_interval, user_id, status): """Update status for all reports Args: activity (ActivityWorker): the activity worker. period_ids (str): comma delimited string of accounting period ids user_type (str): label, distributor or subaccount payment_interval (str): payment interval in the contract (month or quarter) user_id (str): labelid or subaccountid status (str): status of the report """ activity.logger.info('Updating status...') generator = SqlGenerator( period_ids, user_type, payment_interval, user_id) if user_id == 'all': # Iterating through all labels, inserting one entry per label user_params = '__'.join([period_ids, 'all', 'AVRO', 'en_US']) start_time = time.time() for row in db_util.snowflake_query( generator.get_sql_for_bulk_status_update()): # Exclude column user_id_type, periods and status in row resultset row = {k.lower(): v for k, v in row.items()} row_excluded_required_params = { k: v for k, v in row.items() if k not in [ 'user_id_type', 'user_params', 'status'] } # For write capacity 1000 writes/sec, the max time would be 0.001s # So, if the time difference is less than 0.001, we need to wait # for the time difference to prevent exceeding the capacity time_diff = time.time() - start_time max_time = 1 / setting.HIGH_WRITE_CAPACITY_UNITS if time_diff < max_time: time.sleep(max_time - time_diff) dynamodb_util.set_status( row.get('user_id_type'), user_params, status, period_ids=period_ids, **row_excluded_required_params) start_time = time.time() else: user_params = '__'.join([period_ids, 'all', 'AVRO', 'en_US']) dynamodb_util.set_status( '{}{}'.format( user_id, user_type[0].upper()), user_params, status, period_ids=period_ids)