"""Tasks for sql2sf workflow.""" import datetime import json import subprocess import tempfile import boto3 from garcon import task from garcon_contrib.aws.utils import garcon_s3 from garcon_contrib.snowflake import garcon_snowflake from snowflake_etl.conf.config import merge_configs from snowflake_etl.conf.config import SF_CREDENTIALS from snowflake_etl.flows.s32sf import helpers as s32sf_helpers from snowflake_etl.flows.sql2sf import config from snowflake_etl.flows.sql2sf import helpers from snowflake_etl.flows.sql2sf.config import DYNAMODB_STATUS_TABLE sources = config.SourcesConf() @task.decorate(timeout=1000) def bootstrap( activity, env, db_type, source_db_host, source_schema, source_table, query_type, file_format=None, start_period=None, end_period=None): """Bootstrap the RDBMS to S3 load. Args: activity (ActivityWorker): The activity worker. env (str): Environment name, e.g. ('dev'). db_type (str): Type of a source database ('mysql' or 'redshift'). source_db_host (str): The name of host under db_hosts key in sql2sf_sources.yml. E.g. 'reportsar02'. source_schema (str): The name of the source schema (e.g., 'production'. source_table (str): The name of the source table (e.g., 'dim_zip'. query_type (str): The key to extract SQL query from the config (e.g. 'snapshot'). file_format (str): FILE_FORMAT for COPY command of Snowflake. start_period (str): The date in YYYY-MM-DD format or timestamp in YYYY-MM-DDTHH:MM:SS format. end_period (str): The date in YYYY-MM-DD format or timestamp in YYYY-MM-DDTHH:MM:SS format. Raises: AssertionError: Indicates a a missing value in context. Returns: dict: A new context to be passed down the workflow. """ activity.logger.info('Bootstraping sql2sf workflow...') assert sources.is_schema_supported( source_db_host, source_schema), 'Schema not supported: {}'.format( source_schema) assert sources.is_table_supported( source_db_host, source_schema, source_table), \ 'Table not supported: {}'.format(source_table) source_query = sources.get_query(db_type, query_type) qargs = { 'source_schema': source_schema, 'source_table': source_table, 'source_query': source_query, 'start_period': start_period, 'end_period': end_period } table_params = sources.get_table_params( source_db_host, source_schema, source_table) date_col = table_params.get('date_col') qargs['source_date_col'] = date_col # prepare query sql source_sql = ' '.join(source_query.format(**qargs).split()) # determine the date range if start_period is None or end_period is None: if not sources.is_table_supports_incremental_sync( source_db_host, source_schema, source_table): today = datetime.date.today().isoformat() date_range = '-'.join((today, today)) else: return { 'stop': 'True', 'message': 'No start_period and end_period specified for table which ' 'supports multiple sync types, can not choose defaults ' 'implicitly, stop the workflow'} else: date_range = '-'.join((start_period, end_period)) # data unload s3 path s3_path_data = sources.get_s3_path_data(db_type).format( env=env, date_range=date_range, db_host=source_db_host, source_schema=source_schema, source_table=source_table) # table schema s3 path s3_path_schema = sources.get_s3_path_schema(db_type).format( env=env, date_range=date_range, db_host=source_db_host, source_schema=source_schema, source_table=source_table) s3_bucket, s3_key = garcon_s3.extract_bucket_path(s3_path_data) # clear data for this prefix to get rid of stale data s3_prefix_to_clear = '/'.join(s3_path_data.split('/')[:-3]) + '/' # load data from this prefix to Snowflake s3_prefix_to_load = '/'.join(s3_path_data.split('/')[:-1]) return dict( db_type=db_type, query_type=query_type, source_db_host=source_db_host, source_schema=source_schema, source_table=source_table, unload_source_sql=source_sql, destination_s3_bucket=s3_bucket, destination_s3_key=s3_key, s3_path_data=s3_path_data, s3_path_schema=s3_path_schema, s3_prefix_to_clear=s3_prefix_to_clear, s3_prefix_to_load=s3_prefix_to_load, start_period=start_period, end_period=end_period, file_format=file_format, validate_schema=str(sources.is_schema_validation_needed( source_db_host, source_schema, source_table)), ) @task.decorate(timeout=423000, heartbeat=423000) def pipe_mysql_from_stdin_to_stdout( activity, pipe, source_db_host, source_schema): """Extract data from MySQL. Args: activity (ActivityWorker): The SWF activity worker. pipe (Popen): Store pipe signal. Returns: dict: A dictionary with a Popen object. """ host, port, user, password = sources.get_db_credentials( source_db_host=source_db_host, source_schema=source_schema) if pipe is None: raise Exception('Pipe is empty') temp_stderr = tempfile.TemporaryFile(mode='w+t') p2 = subprocess.Popen( ['mysql', '-q', '-N', '--host={}'.format(host), '--port={}'.format(port), '-u', user, '--password={}'.format(password), '--protocol=TCP'], stdin=pipe.stdout, stdout=subprocess.PIPE, stderr=temp_stderr, close_fds=True) activity.logger.info('Setup of pipe_mysql_to_stdout pipe task is done.') return dict(pipe=p2, mysql_stderr=temp_stderr) @task.decorate(timeout=1000) def describe_mysql_table( activity, source_db_host, source_schema, source_table): """Describe the schema (column types, etc.) of a source MySQL table. Args: activity (ActivityWorker): The activity worker. source_db_host (str): The name of host under db_hosts key in sql2sf_sources.yml. E.g. 'reportsar02'. source_schema (str): The name of the source schema (e.g., 'production'. source_table (str): The name of the source table (e.g., 'contact'. Returns: dict: A new item 'describe_json' to be passed down the workflow. """ _DESCRIBE_TABLE_SQL_TEMPLATE = """ SELECT column_name, is_nullable, data_type, column_default, character_maximum_length, character_octet_length, numeric_precision, NULL AS numeric_precision_radix, numeric_scale, datetime_precision FROM information_schema.columns WHERE table_schema = '{schema}' AND table_name = '{table}' ORDER BY ordinal_position""" host, port, user, password = sources.get_db_credentials( source_db_host, source_schema) activity.logger.info( 'Describing MySQL table: {schema}.{table}'.format( schema=source_schema, table=source_table)) # prepare sql sql = _DESCRIBE_TABLE_SQL_TEMPLATE.format( schema=source_schema, table=source_table) # execute unload sql results = helpers.execute_with_mysql( sql, 'fetchall', host, port, user, password, cursor_type='DictCursor') # prepare results json _json = json.dumps( dict(schema=source_schema, table=source_table, columns=results)) return dict(describe_json=_json) # TODO: Check how it works with a composite primary key @task.decorate(timeout=72000) def get_chunk_query( activity, source_db_host, source_schema, source_table, destination_s3_key, min_id, max_id, columns): """Get chunk query task. This task returns the chunk query (or single SELECT * query if table isn't big) by replacing the min_id an max_id placeholders in query in the context file with the ids retrieved from the generator. This query is being passed forward and for each activity instance to use. Args: activity (ActivityWorker): The activity worker. source_db_host (str): The name of host under db_hosts key in sql2sf_sources.yml. E.g. 'reportsar02'. source_schema (str): The name of the source schema (e.g., 'production'. source_table (str): The name of the source table (e.g., 'contact'. destination_s3_key (str): Destination S3 key path. min_id (int): Min primary key id of the data chunk. max_id (int): Max primary_key id of the data chunk. columns (list): A list of tuples with columns, and to wrap or not to wrap bool. Raises: Exception: If a value is missing from the context, an error is thrown. Returns: dict: A dictionary with a pipe and a destination_s3_key. """ # we must wrap all text columns with a stripSpecialChars or # stripSpecialCharsLong UDF if it's required for the table wrapped_columns = [] strip_special_chars_func_name = sources.get_stripspecialchars_func_name( source_db_host, source_schema, source_table) # we must wrap columns with bad chars with CONVERT ... CHAR UNICODE columns_to_wrap_in_convert_unicode = sources.get_table_params( source_db_host, source_schema, source_table).get( 'wrap_to_convert_unicode') for column in columns: column_name, is_wrap_column = column if (columns_to_wrap_in_convert_unicode and column_name in columns_to_wrap_in_convert_unicode): column_name = 'CONVERT((' + column_name + '), CHAR(5000) UNICODE)' if is_wrap_column and strip_special_chars_func_name: wrapped_columns.append( strip_special_chars_func_name + '(' + column_name + ')') else: wrapped_columns.append(column_name) wrapped_columns = ', '.join(column for column in wrapped_columns) if sources.is_unload_in_chunks( source_db_host, source_schema, source_table): primary_key = sources.get_primary_key( source_db_host, source_schema, source_table) query = """SELECT {columns} FROM `{source_table}` WHERE {primary_key} BETWEEN {min_id} AND {max_id};""".format( columns=wrapped_columns, source_table=source_table, primary_key=primary_key, min_id=min_id, max_id=max_id) destination_s3_key = destination_s3_key.replace( '.gz', str(min_id) + '-' + str(max_id) + '.gz') else: query = 'SELECT {columns} FROM `{source_table}`'.format( columns=wrapped_columns, source_table=source_table) activity.logger.info('Setting up query: {query}'.format(query=query)) command = 'USE {source_schema}; {query}'.format( source_schema=source_schema, query=query) p1 = subprocess.Popen( ['echo', command], stdout=subprocess.PIPE, close_fds=True) return dict(pipe=p1, destination_s3_key=destination_s3_key) @task.decorate(timeout=12000) def sanity_check_rows_count( activity, db_type, sfdb_params, query_type, source_db_host, source_schema, source_table): """Check if row counts in source table and stg_ Snowflake table are equal. For now it's working only for snapshot strategy. Args: activity (ActivityWorker): The activity worker. db_type (str): Type of a source database ('mysql' or 'redshift'). sfdb_params (dict): Snowflake credentials. query_type (str): The type of unload query (e.g. 'snapshot'). source_db_host (str): The name of host under db_hosts key in sql2sf_sources.yml. E.g. 'reportsar02'. source_schema (str): The name of the source schema (e.g., 'production'. source_table (str): The name of the source table (e.g., 'dim_zip'. """ if not query_type == 'snapshot': activity.logger.info( 'Skipping sanity_check_rows_count, sync strategy is not snapshot') return if db_type != 'mysql': activity.logger.info( 'Skipping sanity_check_rows_count, source db_type is not MySQL') return mysql_sql = 'SELECT COUNT(*) FROM `{schema}`.`{table}`;'.format( schema=source_schema, table=source_table) host, port, user, password = sources.get_db_credentials( source_db_host, source_schema) mysql_row_count = helpers.execute_with_mysql( mysql_sql, 'fetchone', host, port, user, password)[0] sf_config = merge_configs(sfdb_params, SF_CREDENTIALS) db, schema = s32sf_helpers.extract_db_and_schema(sf_config) snowflake_sql = 'SELECT COUNT(*) FROM {db}.{schema}.stg_{table};'.format( db=db, schema=schema, table=source_table) snowflake_row_count = garcon_snowflake.execute_with_py_conn( snowflake_sql, sf_config=sf_config)['results'][0] threshold = sources.get_sync_sanity_threshold( source_db_host, source_schema, source_table) if abs(mysql_row_count - snowflake_row_count) > threshold: # this prevent next activity, but error status will be set return { 'stop': 'True', 'message': 'Sync strategy is snapshot, but source table row' ' count is {source}, and staging destination row ' 'count is {dest}'.format(source=mysql_row_count, dest=snowflake_row_count)} @task.decorate(timeout=1000) def set_sync_status( activity, end_period, load_strategy, status, source_schema, source_table): """Check if row counts in source table and stg_ Snowflake table are equal. Args: activity (ActivityWorker): The activity worker. end_period (str): up to what date is synced. Should be YYYY-MM-DDTHH:MM:SS. load_strategy (str): Load strategy (e.g., 'snapshot' or 'incremental'). status (str): status of execution. source_schema (str): The name of the source schema (e.g., 'production'. source_table (str): The name of the source table (e.g., 'dim_zip'. """ # format used by monty ts_format = '%Y-%m-%dT%H:%M:%S' today = datetime.datetime.now().strftime(ts_format) synced_to = today if load_strategy == 'snapshot' else end_period data = { 'schema_name': source_schema, 'table_name': source_table, 'last_sync_timestamp': today, 'synced_to_timestamp': synced_to, 'status': status } try: dynamodb = boto3.client('dynamodb') dynamodb.put_item( TableName=DYNAMODB_STATUS_TABLE, Item=data ) except Exception as e: activity.logger.error( 'Failed to report flow execution status: {}'.format(e)) activity.logger.info('Status {} for workflow was set'.format( status))