"""Component for loading source files directly to staging_raw via stage.""" from datetime import datetime from feed_ingestion.common.staging_raw_sf.snowflake_stage_loader \ import StageLoader from feed_ingestion.flows.youtube_monthly import util class YouTubeMonthlySL(StageLoader): """StageLoader class.""" def clean_staging_raw_table(self, staging_raw_table, date): """Delete rows from previous unsuccessful workflow run. Args: staging_raw_table (str): A table name in Snowflake. date (str): Date of the data being process (YYYY-MM-DD). """ # Need to get first and last days of period for YouTubeMonthly date_obj = datetime.strptime(date, '%Y-%m-%d') \ if date else datetime.today() start_date, end_date = util.get_first_last_day(date_obj) self.resolve_sql_loader_and_execute( 'delete_from_staging_raw', params=dict( db=self.executor.sf_config['db'], schema=self.executor.sf_config['schema'], staging_raw_table=staging_raw_table, start_date=start_date, end_date=end_date ) ) def load_staging_raw_table( self, staging_raw_table, source_files_dict, date, stage_name): """Load the temp_staging_raw data to the feed's staging_raw table. Args: staging_raw_table (str): A table name in Snowflake. source_files_dict (dict): A dict with source files metadata. date (str): Date of the data being process (YYYY-MM-DD). stage_name (str): Name of Snowflake stage containing source files. """ ingestion_time = datetime.now() for file_dict in source_files_dict['files']: file_name = file_dict['file_name'] if 'orchard' in file_name.lower(): sql_template = 'load_staging_raw_orchard' file_source = 'ORCHARD' elif 'dmgi' in file_name.lower(): sql_template = 'load_staging_raw_dmgi' file_source = 'DMGI' elif 'ioda' in file_name.lower(): sql_template = 'load_staging_raw_ioda' file_source = 'IODA' else: continue self.resolve_sql_loader_and_execute( sql_template, params=dict( db=self.executor.sf_config['db'], schema=self.executor.sf_config['schema'], stage=stage_name, staging_raw_table=staging_raw_table, file_name=file_dict['file_name'], file_size=file_dict['file_size'], file_source=file_source, download_date=date, ingestion_time=ingestion_time ) )