"""Cable Ingestion ETL Tasks.""" import glob import gzip import json import os from os import path import shutil import tempfile from garcon import task from garcon_contrib.dynamo_feed_status import garcon_feed_status import pysftp from flows import datastore from flows import s3 from flows import util from flows.cable_ingestion import config from flows.cable_ingestion import log from flows.cable_ingestion import status from flows.cable_ingestion import util as etl_util from flows.exceptions import EmptyGeneratorError @task.decorate(timeout=1800) def download_from_server( activity, correlation_id, date_end, date_start, sftp_credentials): """Download all paired files that are in the date range. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. date_end (str): YYYY-MM-DD exclusive date end range for download. date_start (str): YYYY-MM-DD inclusive date start range for download. sftp_credentials (dict): SFTP login credentials for source data. Returns: dict: location of temporary files. """ temp_directory = path.join(tempfile.gettempdir(), correlation_id) if os.path.isdir(temp_directory): shutil.rmtree(temp_directory) os.mkdir(temp_directory) date_end = date_end.replace('-', '') date_start = date_start.replace('-', '') # this is to disable checking for host keys. cnopts = pysftp.CnOpts(knownhosts=None) cnopts.hostkeys = None with pysftp.Connection(cnopts=cnopts, **sftp_credentials) as sftp: sftp.cwd('OUT') filenames = sftp.listdir() bound_files = etl_util.filter_files_by_date_range( date_end, date_start, filenames) file_pairs = etl_util.filter_file_pairs(bound_files) for pair in file_pairs: for filename in pair.values(): destination = path.join(temp_directory, filename) with open(destination, 'wb'): sftp.get(filename, destination) log.update_status(correlation_id, status.DOWNLOADED_FROM_FTP) return {'directory': temp_directory} @task.decorate(timeout=900) def validate_source_files(activity, correlation_id, directory): """Use the Rentrak file pairs downloaded to verify integrity. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. directory (str): temporary local directory where files live. Returns: dict: valid files to move to S3 for storage. """ glob_path = path.join(directory, '*') files = glob.glob(glob_path) file_pairs = etl_util.filter_file_pairs(files) valid_files = [] for pair in file_pairs: with gzip.open(pair['ctl'], mode='rt') if config.BACKFILL_RUN \ else open(pair['ctl']) as ctl_fh, \ open(pair['tar.gz'], mode='r+b') as data_fh: ctl_data = etl_util.extract_ctl_data(ctl_fh) tar_data = etl_util.extract_rentrak_metadata(data_fh) tar_map = {tar_file['name']: tar_file for tar_file in tar_data} for ctl_file in ctl_data: filename = ctl_file['name'] tar_file = tar_map[filename] if filename in tar_map and ctl_file == tar_file: del tar_map[filename] if not tar_map: valid_files.extend(pair.values()) else: bad_files = list(tar_map.keys()) message = { 'message': ( 'Bad files in Cable Ingestion ETL,' 'task: "validate_source_files'), 'correlation_id': correlation_id, 'files': bad_files} util.send_sns_message( config.SNS_ACTION_FAILURE, json.dumps(message), config.SNS_TOPIC_ARN) return {'valid_files': False} log.update_status(correlation_id, status.VALIDATED_SOURCE_FILES) return {'valid_files': valid_files} @task.decorate(timeout=900) def upload_to_s3_archive( activity, bucket, correlation_id, destination, valid_files): """Upload verified temporary files to S3 for archiving. Args: activity (ActivityWorker): activity worker. bucket (str): s3 destination bucket name. correlation_id (str): etl correlation ID. destination (str): s3 destination location with bucket and path. valid_files (list): valid files to move to S3 for storage. Returns: dict: files moves to S3. """ if not valid_files: return archives = [] for filename in valid_files: with open(filename, 'rb') as fh: s3_destination = path.join( destination.format( bucket=bucket, correlation_id=correlation_id), path.basename(filename)) s3_object = s3.get_object(s3_destination) s3_object.put(Body=fh) archives.append(s3_destination) log.update_status(correlation_id, status.UPLOADED_RAW_TO_S3_ARCHIVES) return {'s3_files': archives} @task.decorate(timeout=300) def cleanup_local_raw(activity, correlation_id, directory, valid_files): """Cleanup temporary files and anything else from the extract activity. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. directory (str): temporary local directory where files live. valid_files (list): valid files to move to S3 for storage. """ if not valid_files: return shutil.rmtree(directory) log.update_status(correlation_id, status.CLEANED_UP_LOCAL_RAW) @task.decorate(timeout=60) def set_final_status(activity, correlation_id): """Set status COMPLETED to the db log. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. """ log.update_status(correlation_id, status.COMPLETED, True) @task.decorate(timeout=300) def create_temp_raw_table(activity, correlation_id, sql, table_name): """Create temporary raw table for data staging. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. sql (str): create table query. table_name (str): format string for the temp table name. Returns: dict: temp table name. """ correlation_hex = util.correlation_id_hex(correlation_id) temp_table_name = table_name.format(correlation_hex=correlation_hex) query = sql.format(table_name=temp_table_name) datastore.execute(query) log.update_status(correlation_id, status.CREATED_TEMP_RAW_TABLE) return {'table_name': temp_table_name} @task.decorate(timeout=3600) def insert_to_temp_raw_table( activity, correlation_id, delete, insert, s3_files, table_name): """Insert S3 CSV tarballs into temp raw table. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. delete (str): delete by date range query. insert (str): insert into query. s3_files (list): S3 tarball and ctl urls. table_name (str): the temp table name. """ insert_query = insert.format(table_name=table_name) s3_files = sorted( (s3_file for s3_file in s3_files if s3_file.endswith('.tar.gz'))) with datastore.context() as (cursor, connection): for s3_file in s3_files: try: rows = etl_util.get_rows_from_tarball(s3_file, delimiter='|') # clear rows in temp file_dates = etl_util.extract_filename_dates( path.basename(s3_file)) delete_query = delete.format(table_name=table_name) delete_params = { 'date_end': file_dates['date_end'], 'date_start': file_dates['date_start']} cursor.execute(delete_query, delete_params) # fill into temp cursor.executemany(insert_query, rows) except EmptyGeneratorError: continue log.update_status(correlation_id, status.INSERTED_TO_TEMP_RAW_TABLE) @task.decorate(timeout=900) def insert_select_to_raw_table( activity, correlation_id, delete, drop, insert, raw_table_name, table_date_range, temp_table_name): """Insert rows from temp table into the raw table. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. delete (str): delete from table query. drop (str): drop table query. insert (str): insert select query. raw_table_name (str): raw table name. table_date_range (str): query to determine replacing data's date range. temp_table_name (str): the temp table name. """ drop = drop.format(table_name=temp_table_name) insert = insert.format(table_name=temp_table_name) date_range_query = table_date_range.format(table_name=temp_table_name) date_start, date_end = datastore.query(date_range_query).fetchone() delete_query = delete.format(table_name=raw_table_name) with datastore.context() as (cursor, connection): cursor.execute('START TRANSACTION') cursor.execute( delete_query, {'date_end': date_end, 'date_start': date_start}) cursor.execute(insert) cursor.execute(drop) log.update_status(correlation_id, status.INSERTED_TO_RAW_TABLE) @task.decorate(timeout=60) def update_dashboard_status(activity, correlation_id, status): """Update the Orchard global level overall status for the dashboard. Args: activity (ActivityWorker): activity worker. correlation_id (str): etl correlation ID. """ row = log.get_etl_log(correlation_id) date = str(row[4].date()) garcon_feed_status.set_overall_status( config.SWF_WORKFLOW_NAME, date, status)