"""General ETL utility functions.""" from argparse import ArgumentTypeError from collections import defaultdict import csv import functools import gzip import inspect import json import uuid import boto3 from flows import art_relations from flows import config from flows import constant from flows import log from flows import queries from flows import s3 from flows.flow import DatabaseParam def get_serviced_upcs(query_based=False): """Return a tuple of UPCs covered under the Film Transparency service. Args: query_based (bool): use query heuristic for UPCs (for compatibility). Returns: tuple: UPC strings. """ mapping = get_serviced_vendor_id_upc_map(query_based) upcs = [] for values in mapping.values(): upcs.extend(list(values)) return tuple(upcs) def get_serviced_vendor_id_upc_map(query_based=False): """Return a mapping of vendor ids to upcs covered under the service. Args: query_based (bool): use query heuristic for UPCs (for compatibility). Returns: dict: vendor_id: set(upcs) """ if query_based: rows = art_relations.query(queries.serviced_vendor_id_upc_map_sql) mapping = defaultdict(set) for vendor_id, upc in rows: mapping[vendor_id].add(str(upc)) return mapping return constant.FILM_WHITELIST def correlation_id_hex(correlation_id): """Extract only the hex characters of a correlation ID. Args: correlation_id (str): correlation ID. Returns: str: hex characters in the correlation ID. """ base = correlation_id.split('.')[0] parts = base.split('-') return ''.join(parts) def upc_cli_type(value): """Assert the command line input value is numeric. This goes beyond just the .isnumeric() or .isdecimal() methods because they support unicode values like fractions. Args: value (str): input string from command line. Returns: str: same string if valid. Raises: ArgumentTypeError: the input value is not in the proper format. """ try: assert value.isdecimal() assert len(value) in (12, 13) assert all(48 <= ord(c) <= 57 for c in value) except AssertionError: raise ArgumentTypeError('UPC is not in valid format') return value def date_cli_type(value): """Assert the input value is in YYYY-MM-DD format. Args: value (str): input string from command line. Returns: str: same string if valid. Raises: ArgumentTypeError: the input value is not in the proper format. """ try: parts = value.split('-') assert len(parts[0]) == 4 assert len(parts[1]) == 2 assert len(parts[2]) == 2 for part in parts: assert part.isdecimal() except AssertionError: raise ArgumentTypeError('Date is not in valid format') return value def create_swf_execution_params( correlation_id, timeout, workflow_name, workflow_version, workflow_id=None): """Create SWF workflow execution parameters based on correlation_id. Args: correlation_id (str): Correlation ID of workflow. timeout (int): workflow execution lifetime timeout. workflow_name (str): workflow name. workflow_version (str): workflow version. workflow_id (str): optional workflow_id. Returns: dict: params for boto3.client('swf').start_workflow_execution. """ if workflow_id: workflow_id_value = workflow_id else: workflow_id_value = '{name}_{cid}'.format( name=workflow_name, cid=correlation_id) return { 'domain': config.SWF_DOMAIN, 'executionStartToCloseTimeout': str(timeout), 'tagList': ['cid:{cid}'.format(cid=correlation_id)], 'taskList': {'name': workflow_name}, 'workflowId': workflow_id_value, 'workflowType': {'name': workflow_name, 'version': workflow_version}} def start_swf_execution( context, timeout, workflow_name, workflow_version, workflow_id=None): """Start a SWF workflow with the correlation ID and context. Args: context (dict): context to pass to the execution. timeout (int): workflow execution lifetime timeout. workflow_name (str): workflow name. workflow_version (str): workflow version. workflow_id (str): optional workflow_id. Returns: dict: dict containing the runId of the execution. """ correlation_id = context.get('correlation_id', str(uuid.uuid1())) params = create_swf_execution_params( correlation_id, timeout, workflow_name, workflow_version, workflow_id) swf = boto3.client('swf', region_name=config.SWF_REGION_NAME) context['correlation_id'] = correlation_id for key, value in context.items(): if isinstance(value, DatabaseParam): value.put_data(correlation_id) context[key] = value.context_key info_log = ( 'Starting SWF workflow execution with correlation ID: ' '{correlation_id}') logger = log.get_logger() logger.info(info_log.format(correlation_id=correlation_id)) response = swf.start_workflow_execution( input=json.dumps(context), **params) return response def send_sns_message(subject, message, topic_arn): """Send a SNS message. Args: subject (str): SNS subject. message (str): SNS message body. topic_arn (str): SNS Topic ARN. Returns: dict: response with message ID. """ client = boto3.client('sns', region_name=config.SWF_REGION_NAME) return client.publish(TopicArn=topic_arn, Message=message, Subject=subject) def batch_param_calls(parameter, batch_size): """Parameterized decorator to call multiple times given a param. Args: parameter (str): parameter name to batch on. batch_size (int): size of the parameter to call per batch. Returns: decorator: specified decorator to wrap callable around. """ def decorator(callee): """Decorator to call a callable multiple times in batches. Args: callee (callable): callable with the param to wrap. Returns: Wrapper: batching callable if parameter exists, else original. """ signature = inspect.signature(callee) if parameter not in signature.parameters: return callee class Wrapper: """Decorating class to generate batch calls of a callable. The self.batch_size is stored and recalled during call time to allow changes during testing. """ def __init__(self, callee): """Decorator initialization. Args: callee: callable to wrap. """ self.callee = callee self.batch_size = batch_size def __call__(self, *args, **kwargs): """Generation of callable results based on params and size. The decorated callable returns a generator object and not the results usually expected. Args: args (tuple): call parameters passing through. kwargs (dict): call parameters passing through. Yields: any: results of batched calls to callable. """ start = 0 end = self.batch_size binded_call = signature.bind(*args, **kwargs) upcs = binded_call.arguments[parameter] upc_slice = upcs[start:end] while upc_slice: binded_call.arguments[parameter] = upc_slice yield self.callee( *binded_call.args, **binded_call.kwargs) start = end end += self.batch_size upc_slice = upcs[start:end] wrapper = Wrapper(callee) functools.update_wrapper(wrapper, callee) return wrapper return decorator def batch_upc_calls(callee): """Decorator to call a callable multiple times for many UPCs. This is a convenience decorator that uses batch_param_calls for backwards compatibility. Args: callee (callable): function with "upcs" param to wrap. Returns: callable: batching callable if "upcs" parameter, else original. """ decorator = batch_param_calls('upcs', config.ITERATION_BATCH_SIZE) return decorator(callee) @batch_upc_calls def get_upc_vendor_id(upcs): """Get the UPCs' vendor IDs from the database. This function is a generator due to the batching decorator. Args: upcs (list): UPCs to get vendor IDs of. Yields: dict: upc to vendor_id mapping in batches. """ sql = queries.upc_vendor_id_lookup(upcs) rows = art_relations.query(sql) return {str(r[0]): str(r[1]) for r in rows} def read_gzip_csv_from_s3(s3_path): """Read gzipped csv file from s3. Try to use garcon-contrib.mysql fns if this is all your task does. This is for cases when you need to load multiple files in single transaction. Note: dupe from digital etl. Todo to update digital to use this common one. Args: s3_path (str): path to csv on s3. Yields: list: csv rows: [['v1', 'v2'], ['v3', 'v4']]. """ raw = s3.get_object(s3_path).get()['Body'].read() preprocessed = filter( lambda f: len(f) > 0, map( lambda m: m.strip(), gzip.decompress(raw).decode('utf8').split('\n') ) ) yield from csv.reader(preprocessed)