"""Sales data ETL Tasks.""" from calendar import monthrange import datetime import json from botocore.exceptions import ClientError from garcon import task from garcon_contrib.dynamo_feed_status import garcon_feed_status from flows import config as flow_config from flows import datastore from flows import datawarehouse from flows import g from flows import util as flow_util from flows.sales_data import config from flows.sales_data import log from flows.sales_data import ows_accounting from flows.sales_data import queries from flows.sales_data import status from flows.sales_data import util from flows.util import correlation_id_hex @task.decorate(timeout=400) def bootstrap(activity, correlation_id, accounting_period_id): """Prepare parameters for flow execution. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. accounting_period_id (str): sales_data period id to process. Returns: dict: Dict with start and end dates for accounting_period_id. """ if not accounting_period_id: accounting_period_id = ows_accounting.get_accounting_period() if util.has_ingested_accounting_period_id(accounting_period_id): return { 'stop': True, 'reason': 'Accounting period ID already processed.'} log.accounting_period_update(correlation_id, accounting_period_id) period_data = util.get_period_data(accounting_period_id) if not period_data: log.update_status(correlation_id, status.TERMINATED_BAD_PERIOD_ID) return {'stop': True, 'reason': 'Bad accounting_period_id.'} # in SFlake: both values are inclusive in between clause. date_start = datetime.date(int(period_data[1]), int(period_data[2]), 1) no_of_days = monthrange(date_start.year, date_start.month)[1] date_end = date_start + datetime.timedelta(days=no_of_days-1) return { 'date_start': date_start.strftime('%Y-%m-%d'), 'date_end': date_end.strftime('%Y-%m-%d'), 'accounting_period_id': accounting_period_id} @task.decorate(timeout=7200) def unload_sales_data(activity, correlation_id, upcs, accounting_period_id): """Unload the data from Fact Sales datawarehouse to s3. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. upcs (list): upcs for unload query. accounting_period_id (str): sales_data period id to process. Returns: dict: unload batch count and unload_path. """ s3_destination = config.UNLOAD_DESTINATION.format( bucket=flow_config.DATA_BUCKET, correlation_id=correlation_id) # @todo go back and change this util fn to take one period_id unload_sqls = util.get_unload_from_snowflake_sql( s3_destination, correlation_id, accounting_period_id, accounting_period_id, upcs) for batch, sql_template in enumerate(unload_sqls): g.log.info('for batch: {}'.format(batch)) g.log.info(sql_template) sql = sql_template.format(batch=batch) datawarehouse.execute(sql) log.update_status(correlation_id, status.UNLOADED_FACT_SALES) return {'batch_count': batch + 1} @task.decorate(timeout=120) def create_temp_table(activity, correlation_id): """Create the temporary table in the datastore to hold unloaded s3 data. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. Returns: dict: temp table name. """ table_name = config.TEMP_TABLE_NAME.format( correlation_hex=correlation_id_hex(correlation_id)) sql = queries.CREATE_TEMP_RAW_TABLE.format(table_name=table_name) datastore.execute(sql) g.log.info('created temp table ' + table_name) return {'table_name': table_name} @task.decorate(timeout=120) def load_temp_table(activity, correlation_id, batch_count, temp_table_name): """Load temporary table in the datastore. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. batch_count (int): No of chunks of upcs that were made. temp_table_name (str): Temp raw table name. """ s3_location_format = config.UNLOAD_DESTINATION.format( bucket=flow_config.DATA_BUCKET, correlation_id=correlation_id) column_names = ', '.join(queries.UNLOAD_COLUMNS) values_placeholder = ('%s, ' * len(queries.UNLOAD_COLUMNS)).rstrip(', ') insert_sql = queries.INSERT_TO_TEMP_RAW_TABLE.format( table_name=temp_table_name, column_names=column_names, values_placeholder=values_placeholder) g.log.info(insert_sql) with datastore.context() as (cursor, connection): for batch_number in range(batch_count): try: s3_path = s3_location_format.format(batch=batch_number) rows = flow_util.read_gzip_csv_from_s3(s3_path) cursor.executemany(insert_sql, rows) except ClientError as err: # If unload this batch file is not found, it just had no data, # and therefore not on s3 (NoSuchKey error). Just move on and # try the next batch. g.log.error('ClientError for batch: {}. Error code: {}'.format( batch_number, err.response['Error']['Code'])) if err.response['Error']['Code'] != 'NoSuchKey': raise err log.update_status(correlation_id, status.TEMP_TABLE_INSERTED) @task.decorate(timeout=800) def load_raw_from_temp( activity, correlation_id, upcs, temp_table_name, accounting_period_id): """Load Raw table from temp table. This will delete existing raw data for selected period ids and insert new data from temp table. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. upcs (list): upcs for unload query. temp_table_name (str): Temp raw table name. accounting_period_id (str): sales_data period id to process. """ delete_sqls = util.delete_raw_sql(upcs, accounting_period_id) with datastore.context() as (cursor, connection): cursor.execute('START TRANSACTION') for delete_sql in delete_sqls: cursor.execute(delete_sql) column_names = ', '.join(queries.UNLOAD_COLUMNS) insert_sql = queries.INSERT_RAW_DATA.format( table_name=temp_table_name, column_names=column_names) cursor.execute(insert_sql) drop_sql = queries.DROP_TABLE.format(table_name=temp_table_name) cursor.execute(drop_sql) log.update_status(correlation_id, status.INSERTED_RAW_TABLE) @task.decorate(timeout=252000) def insert_daily_revenue(activity, correlation_id, upcs, accounting_period_id): """Insert daily calculated raw aggregates into revenue table. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. upcs (list): upcs aggregate calculation query. accounting_period_id (str): sales_data period id for aggregate. """ raw_rows = util.get_aggregate_raw_data(upcs, accounting_period_id) if not raw_rows: g.log.error('No raw sales data found to insert in revenue.') return with datastore.context() as (cursor, connection): cursor.execute('START TRANSACTION;') delete_queries = util.delete_accounting_data_sql(upcs) for delete_query in delete_queries: cursor.execute(delete_query, { 'accounting_period_id': accounting_period_id }) insert_rows = util.get_daily_revenue_query_params(raw_rows) cursor.executemany(queries.INSERT_REVENUE_DATA, insert_rows) log.update_status(correlation_id, status.INSERTED_DAILY_REVENUE) @task.decorate(timeout=200) def prepare_sns_message( activity, correlation_id, upcs, accounting_period_id): """Send notice that this ETL has finished. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. upcs (str): context key used in DatabaseParam object. accounting_period_id (str): sales_data period id to process. """ message = { 'action': config.SNS_ACTION_SUCCESS, 'correlation_id': correlation_id, 'source': config.SNS_SOURCE, 'upcs': upcs, 'accounting_period_id': accounting_period_id} return { 'topic': config.SNS_TOPIC_ARN, 'message': json.dumps(message), 'subject': config.SNS_ACTION_SUCCESS} @task.decorate(timeout=200) def set_dynamo_status(activity, correlation_id, date_start): """Set DynamoDB status to ingested. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. date_start (str): start date for accounting_period_id. """ garcon_feed_status.set_overall_status( config.SWF_WORKFLOW_NAME, date_start, garcon_feed_status.STATUS_INGESTED) log.update_status(correlation_id, status.DYNAMO_STATUS_UPDATED) @task.decorate(timeout=200) def set_final_status(activity, correlation_id): """Set status INGESTED to the db log. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. """ log.update_status(correlation_id, status.INGESTED, True)