"""Projections ETL Tasks.""" import json from garcon import task from flows import datastore from flows import s3 from flows import util from flows.projections import config from flows.projections import log from flows.projections import queries from flows.projections import status from flows.projections import util as etl_util @task.decorate(timeout=900) def bootstrap( activity, correlation_id, workflow_run_id, s3_bucket, s3_key, projection_type): """Prepare parameters for flow execution. UPC value and file date are taken from s3_key. Initial log record is created since all required data is ready for it. Generate required data for SWF context. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. workflow_run_id (str): workflow run id generated by AWS SWF service. s3_bucket (str): S3 bucket name of ingested CSV file. s3_key (str): S3 key path of ingested CSV file. projection_type (str): Type of projection file we are processing. Returns: dict: Dict with update to SWF context. """ upc, file_date = etl_util.parse_file_name(s3_key) drop_url = 's3://{bucket}/{key}'.format(bucket=s3_bucket, key=s3_key) filename = drop_url.split('/')[-1] return_params = { 'upc': upc, 'drop_url': drop_url, 'revenue_table': config.PROJECTIONS_REVENUE.format( projections=config.REGULAR_PROJECTION), 'log_table': config.PROJECTIONS_REVENUE_ETL_LOG.format( projections=config.REGULAR_PROJECTION), 's3_archive_path': config.S3_ARCHIVE_PATH.format( projections=config.REGULAR_PROJECTION, filename=filename), 'transaction_types': config.TRANSACTION_TYPE} if projection_type == 'original': return_params['revenue_table'] = config.PROJECTIONS_REVENUE.format( projections=config.ORIGINAL_PROJECTION) return_params['log_table'] = config.PROJECTIONS_REVENUE_ETL_LOG.format( projections=config.ORIGINAL_PROJECTION) return_params['s3_archive_path'] = config.S3_ARCHIVE_PATH.format( projections=config.ORIGINAL_PROJECTION, filename=filename) return_params['transaction_types'] = \ config.ORIGINAL_PROJECTION_TRANSACTION_TYPE log.create( correlation_id, workflow_run_id, upc, file_date, return_params['log_table']) return return_params @task.decorate(timeout=900) def determine_anchor_date( activity, correlation_id, upc, log_table, projection_type): """Determine the anchor date for a upc. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. upc (str): release UPC. log_table (str): ETL log table name. projection_type (str): Type of projection file we are processing. Returns: dict: YYYY-MM-DD format anchor date string. """ date = etl_util.get_upc_projection_anchor_date(upc, projection_type) if date: log.update_status( correlation_id, status.ANCHOR_DATE_DETERMINED, log_table) return {'anchor_date': date.isoformat()} log.update_status( correlation_id, status.ANCHOR_DATE_UNDETERMINED, log_table) message = { 'correlation_id': correlation_id, 'message': 'Could not determine anchor date', 'source': config.SNS_SOURCE, 'upc': upc} util.send_sns_message( config.SNS_ACTION_FAILURE, json.dumps(message), config.SNS_TOPIC_ARN) return {'anchor_date': None, 'stop': True} @task.decorate(timeout=100) def create_temp_table(activity, correlation_id, log_table, revenue_table): """Create temp table for projections. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. log_table (str): ETL log table name. revenue_table (str): Revenue data table name. Returns: dict: Dict with temp table name for SWF context. """ table_name = '{table_name}_{correlation_id}'.format( table_name=revenue_table, correlation_id=util.correlation_id_hex(correlation_id)) create_table_sql = queries.CREATE_TEMP_TABLE.format( temp_table=table_name, revenue_table=revenue_table) datastore.execute(create_table_sql) log.update_status(correlation_id, status.TEMP_TABLE_CREATED, log_table) return {'temp_table_name': table_name} @task.decorate(timeout=900) def download_to_db( activity, anchor_date, correlation_id, drop_url, temp_table_name, upc, log_table, transaction_types): """Download data from S3 bucket to db. Args: activity (ActivityWorker): activity worker. anchor_date (str): iso formatted date 'YYYY-MM-DD' correlation_id (str): etl correlation ID. drop_url (str): S3 url of source file. temp_table_name (str): temp table name. upc (str): release UPC. log_table (str): ETL log table name. transaction_types (dict): transaction type mapping. """ insert_sql = queries.INSERT_TO_TEMP_TABLE.format( table_name=temp_table_name) rows = s3.download_csv(drop_url) with datastore.context() as (cursor, connection): cursor.execute('START TRANSACTION;') for row in etl_util.transform( rows, anchor_date, upc, transaction_types): cursor.execute(insert_sql, row) log.update_status( correlation_id, status.DATA_DOWNLOADED_TO_TEMP_TABLE, log_table) @task.decorate(timeout=900) def insert_new_data( activity, correlation_id, temp_table_name, log_table, revenue_table): """Copy ingested data from temp to prod table. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. temp_table_name (str): name for temp projections revenue table. log_table (str): ETL log table name. revenue_table (str): Revenue data table name. """ copy_sql = queries.COPY_TO_PROD.format( temp_table=temp_table_name, prod_table=revenue_table) datastore.execute(copy_sql) log.update_status(correlation_id, status.TEMP_TABLE_MADE_LIVE, log_table) @task.decorate(timeout=900) def drop_temp_table(activity, correlation_id, temp_table_name, log_table): """Drop temp projections revenue table. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. temp_table_name (str): name for temp projections revenue table. log_table (str): ETL log table name. """ datastore.execute(queries.DROP_TABLE.format(table_name=temp_table_name)) log.update_status(correlation_id, status.TEMP_TABLE_DROPPED, log_table) @task.decorate(timeout=900) def archive_dropped_csv( activity, archive_path, bucket, correlation_id, drop_url, log_table): """Grab dropped CSV file into the archive location. Args: activity (ActivityWorker): activity worker. archive_path (str): location to store csv file to on S3. bucket (str): S3 bucket of source and destination files. correlation_id (str): etl correlation ID. drop_url (str): dropped csv source location on S3. log_table (str): ETL log table name. Returns: dict: full archive s3 location. """ drop_obj = s3.get_object(drop_url) archive_url = 's3://{bucket}/{key}'.format(bucket=bucket, key=archive_path) archive_obj = s3.get_object(archive_url) archive_obj.copy({'Bucket': drop_obj.bucket_name, 'Key': drop_obj.key}) drop_obj.delete() log.update_status(correlation_id, status.DROP_FILE_ARCHIVED, log_table) return {'archive_url': archive_url} @task.decorate(timeout=100) def set_final_status(activity, correlation_id, log_table): """Set status INGESTED to the db log. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. log_table (str): ETL log table name. """ log.update_status(correlation_id, status.INGESTED, log_table) @task.decorate(timeout=900) def send_success_notification(activity, correlation_id, upc): """Send notification about success ingestion. Args: activity (Activity): SWF Activity. correlation_id (str): correlation id. upc (str): release UPC. """ if config.SNS_TOPIC_ARN: message = { 'correlation_id': correlation_id, 'message': 'Successfully ingested', 'source': config.SNS_SOURCE, 'upc': upc} util.send_sns_message( config.SNS_ACTION_SUCCESS, json.dumps(message), config.SNS_TOPIC_ARN)