""" Extract data from any SQL db to S3 and load to Snowflake. Now supports MySQL, Aurora and Redshift as the sources. A Garcon workflow to extract data from relational db to S3 and from there load it into Snowflake. """ import copy from garcon import runner from garcon.param import StaticParam from garcon_contrib.aws import garcon_aws_cli_tool from garcon_contrib.aws import garcon_s3 from garcon_contrib.dynamo_feed_status.garcon_feed_status \ import STATUS_INGESTED from garcon_contrib.gzip import garcon_gzip from garcon_contrib.pipe import garcon_pipe from garcon_contrib.pipe import garcon_pipe_task_runner from snowflake_etl.conf import config as cfg from snowflake_etl.conf import flat_dict from snowflake_etl.flows import base from snowflake_etl.flows.s32sf import tasks as tasks_s32sf from snowflake_etl.flows.sql2sf import config from snowflake_etl.flows.sql2sf import generators from snowflake_etl.flows.sql2sf import tasks env = cfg.getconf('env')['env'] _PARAMS_AWS = {k: StaticParam(v) for k, v in flat_dict(cfg.getconf('aws'), '_').items()} class Flow(base.FlowBase): """SQL to Snowflake Workflow Flow class. This flow unloads data from relational databases to Snowflake. """ timeout = 423000 # 5 days def __init__(self): """Initialise the flow instance with the config parameters.""" if env == 'test': domain = 'test' # for the integration test purposes else: domain = '{env}_snowflake_etl'.format(env=env) super(Flow, self).__init__( feed_name=config.feed_name, version='2.0', domain=domain) def workflow_id(self, initial_context=None): """Calculate the workflow's locally unique identifier.""" template = ('{db_type}-{source_schema}.{source_table}-{query_type}-' 'into-{snowflake_db}.{snowflake_schema}') ctx_copy = copy.deepcopy(initial_context) # for date ranges, use start_period and end_period discriminants if ctx_copy['query_type'] in ('date_range', 'sync_incremental'): for key in ('start_period', 'end_period'): ctx_copy[key] = ctx_copy[key].replace(':', '_') template = '-'.join( (template, '{start_period}', '{end_period}')) sf_db = ctx_copy['sfdb_params']['db'] sf_schema = ctx_copy['sfdb_params']['schema'] return template.format( snowflake_schema=sf_schema, snowflake_db=sf_db, **ctx_copy) def decider(self, schedule, context): """Activity decider. Args: schedule (callable): The scheduler method. context (dict): Initial context of workflow. """ bootstrap_sql2sf = schedule('bootstrap_sql2sf', self.bootstrap_sql2sf) if bootstrap_sql2sf.result.get('bootstrap_sql2sf.stop'): return bootstrap_s32sf = schedule( 'bootstrap_s32sf', self.bootstrap_s32sf, requires=[bootstrap_sql2sf]) cleanup_s3 = schedule( 'cleanup_s3', self.cleanup_s3, requires=[bootstrap_s32sf]) db_type = bootstrap_sql2sf.result.get( 'bootstrap_sql2sf.db_type') if db_type == 'mysql': unload_data = schedule( 'unload_mysql_data', self.unload_mysql_data, requires=[bootstrap_s32sf, cleanup_s3]) else: raise ValueError( 'db_type {} not supported'.format(db_type)) if db_type == 'mysql': describe_source_table = schedule( 'describe_source_table', self.describe_mysql_table, requires=[unload_data]) else: raise ValueError( 'db_type {} not supported'.format(db_type)) write_describe_table_to_s3 = schedule( 'write_describe_table_to_s3', self.write_describe_table_to_s3, requires=[describe_source_table]) validate_compatibility = schedule( 'validate_compatibility', self.validate_compatibility, requires=[bootstrap_s32sf, write_describe_table_to_s3]) start_load_snowflake = schedule( 'start_load_snowflake', self.start_load_snowflake, requires=[unload_data, validate_compatibility]) if not start_load_snowflake.result.get('sanity_check_rows_count.stop'): finish_load_snowflake = schedule( 'finish_load_snowflake', self.finish_load_snowflake, requires=[start_load_snowflake]) schedule( 'set_sync_status', self.set_sync_status, requires=[finish_load_snowflake]) @property def bootstrap_sql2sf(self): """Bootstrap the common parameters.""" return self.create( name='bootstrap_sql2sf', tasks=runner.Sync(tasks.bootstrap.fill( namespace='bootstrap_sql2sf', env=StaticParam(env), db_type='db_type', source_db_host='source_db_host', source_schema='source_schema', source_table='source_table', query_type='query_type', start_period='start_period', end_period='end_period', file_format='file_format'))) @property def bootstrap_s32sf(self): """Bootstrap the s32sf parameters.""" return self.create( name='bootstrap_s32sf', tasks=runner.Sync( tasks_s32sf.bootstrap.fill( namespace='bootstrap_s32sf', sfdb_params='sfdb_params', table='bootstrap_sql2sf.source_table', s3_path_data='bootstrap_sql2sf.s3_path_data', s3_path_schema='bootstrap_sql2sf.s3_path_schema', load_strategy='load_strategy'))) @property def cleanup_s3(self): """Remove stale files from the previos unsuccessful flow run.""" return self.create( name='cleanup_s3', retry=2, tasks=runner.Sync( garcon_s3.remove_files_from_path.fill( path='bootstrap_sql2sf.s3_prefix_to_clear'))) @property def unload_mysql_data(self): """Unload data from MySQL to S3.""" return self.create( name='unload_mysql_data', timeout=423000, # 5 days schedule_to_start=423000, # 5 days generators=[generators.unload_mysql_data_params], tasks=garcon_pipe_task_runner.Pipe( tasks.get_chunk_query.fill( namespace='get_chunk_query', source_db_host='bootstrap_sql2sf.source_db_host', source_schema='bootstrap_sql2sf.source_schema', source_table='bootstrap_sql2sf.source_table', columns='generator.columns', destination_s3_key='generator.destination_s3_key', min_id='generator.min_id', max_id='generator.max_id'), tasks.pipe_mysql_from_stdin_to_stdout.fill( namespace='pipe_mysql_from_stdin_to_stdout', source_db_host='bootstrap_sql2sf.source_db_host', source_schema='bootstrap_sql2sf.source_schema', pipe='get_chunk_query.pipe'), garcon_gzip.pipe_gzip_from_stdin_to_stdout.fill( namespace='gzip', pipe='pipe_mysql_from_stdin_to_stdout.pipe'), garcon_aws_cli_tool.upload_to_s3_from_stdin.fill( namespace='s3', pipe='gzip.pipe', destination_s3_key='get_chunk_query.destination_s3_key', bucket='generator.bucket'), garcon_pipe.communicate.fill( pipe='s3.pipe', mysql_stderr=('pipe_mysql_from_stdin_to_stdout.' 'mysql_stderr')))) @property def describe_mysql_table(self): """Extract the schema from MySQL table.""" return self.create( name='describe_mysql_table', tasks=runner.Sync( tasks.describe_mysql_table.fill( namespace='describe_source_table', source_db_host='source_db_host', source_schema='bootstrap_sql2sf.source_schema', source_table='bootstrap_sql2sf.source_table', **_PARAMS_AWS))) @property def write_describe_table_to_s3(self): """Unload the extracted schema on S3.""" return self.create( name='write_describe_table_to_s3', tasks=runner.Sync( garcon_s3.create_object.fill( namespace='write_describe_table_to_s3', path='bootstrap_sql2sf.s3_path_schema', content='describe_source_table.describe_json', **_PARAMS_AWS))) # s32sf activities @property def validate_compatibility(self): """Validate compatibility of the schemas.""" return self.create( name='validate_compatibility', tasks=runner.Sync( tasks_s32sf.source_schema.fill( namespace='validate_compatibility', s3_path_schema='bootstrap_sql2sf.s3_path_schema'), tasks_s32sf.create_destination_table.fill( namespace='validate_compatibility', sfdb_params='bootstrap_s32sf.sfdb_params', db_type='bootstrap_sql2sf.db_type', table='bootstrap_sql2sf.source_table', source_schema='validate_compatibility.source_schema'), tasks_s32sf.target_schema.fill( namespace='validate_compatibility', sfdb_params='bootstrap_s32sf.sfdb_params', table='bootstrap_sql2sf.source_table'), tasks_s32sf.validate_compatibility.fill( namespace='validate_compatibility', db_type='bootstrap_sql2sf.db_type', source_schema='validate_compatibility.source_schema', target_schema='validate_compatibility.target_schema', strict='bootstrap_sql2sf.validate_schema'))) @property def start_load_snowflake(self): """Load data to staging Snowflake table.""" return self.create( name='start_load_snowflake', timeout=72000, # 20 hrs tasks=runner.Sync( # creates or replaces table if exists, so there is no need to # drop old staging table before tasks_s32sf.create_staging_table.fill( table='bootstrap_s32sf.table', db_type='bootstrap_sql2sf.db_type', sfdb_params='bootstrap_s32sf.sfdb_params'), tasks_s32sf.load_staging_table.fill( table='bootstrap_s32sf.table', sfdb_params='bootstrap_s32sf.sfdb_params', db_type='bootstrap_sql2sf.db_type', s3_path_data='bootstrap_sql2sf.s3_prefix_to_load', file_format='bootstrap_sql2sf.file_format'), tasks.sanity_check_rows_count.fill( namespace='sanity_check_rows_count', db_type='db_type', table='bootstrap_s32sf.table', sfdb_params='bootstrap_s32sf.sfdb_params', query_type='bootstrap_sql2sf.query_type', source_db_host='bootstrap_sql2sf.source_db_host', source_schema='bootstrap_sql2sf.source_schema', source_table='bootstrap_sql2sf.source_table'), tasks.set_sync_status.fill( end_period='bootstrap_sql2sf.end_period', load_strategy='bootstrap_s32sf.load_strategy', status=StaticParam('SYNC_ERROR'), source_schema='bootstrap_sql2sf.source_schema', source_table='bootstrap_sql2sf.source_table'))) @property def finish_load_snowflake(self): """Load data to permanent Snowflake table.""" return self.create( name='finish_load_snowflake', timeout=7200, # 2 hrs tasks=runner.Sync( tasks_s32sf.swap_snowflake_tables.fill( table='bootstrap_s32sf.table', sfdb_params='bootstrap_s32sf.sfdb_params', load_strategy='bootstrap_s32sf.load_strategy'), tasks_s32sf.insert_into_destination_table.fill( table='bootstrap_s32sf.table', sfdb_params='bootstrap_s32sf.sfdb_params', load_strategy='bootstrap_s32sf.load_strategy'))) @property def set_sync_status(self): """Set status in the {env}_sql2sf_table_sync_status DynamoDB table.""" return self.create( name='set_sync_status', timeout=300, tasks=runner.Sync( tasks.set_sync_status.fill( namespace='set_sync_status', end_period='bootstrap_sql2sf.end_period', load_strategy='bootstrap_s32sf.load_strategy', status=StaticParam(STATUS_INGESTED), source_schema='bootstrap_sql2sf.source_schema', source_table='bootstrap_sql2sf.source_table')))