"""Snowflake connector class for the BandsInTown tasks.""" from snowflake_connector.etl_connector import SQLLoader from feed_ingestion.common.fact_analytics_sf.base_executor \ import SnowflakeSQLExecutorFA from feed_ingestion.flows.bandsintown import config sql_loader = SQLLoader(__file__) class BandsInTownSF(SnowflakeSQLExecutorFA): """Helper class to abstract Snowflake operations. This class inherits from SnowflakeSQLExecutor class, which provides basic set of methods. This class extends SnowflakeSQLExecutor with some specific methods, which are useful to encapsulate some flow specific operations. """ @property def feed_name(self): """Name of the feed. Should match dir name of this feed, feed_name in config.py of a feed. Returns: str: Feed name """ return config.feed_name def create_temp_staging_raw_table(self, temp_staging_raw_table, **kwargs): """Create a temporary staging raw table. Args: temp_staging_raw_table (str): A table name in Snowflake. kwargs (dict): Custom arguments with report type. """ query_name = config.create_temp_table_query_name.format( report_name=kwargs['report_name']) params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], table_name=temp_staging_raw_table) self.execute_query(sql_loader, query_name, params) def load_temp_staging_raw_table( self, temp_staging_raw_table, aws, key_dir, **kwargs): """Load temp staging raw table. Args: temp_staging_raw_table (str): A table name in Snowflake. aws (dict): AWS credentials. key_dir (str): Custom arguments. """ aws_params = self.get_aws_params() params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], table_name=temp_staging_raw_table, s3_path=key_dir, file_pattern=kwargs['file_pattern'], on_error_action='SKIP_FILE_{}'.format(kwargs['error_limit']), **aws_params, ) self.execute_query(sql_loader, 'load_temp_table', params) def load_staging_raw_table( self, date, staging_raw_table, temp_staging_raw_table, report_name, **kwargs): """Create a temporary staging raw table. Args: date (str): The date to load. staging_raw_table (str): A staging raw table in Snowflake. temp_staging_raw_table (str): A temp table name in Snowflake where data is staged. report_name (str): The report name. kwargs (dict): Custom arguments with report type. """ query_name = config.load_query_name.format( report_name=report_name) params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_raw_table=staging_raw_table, download_date=date, filename=kwargs['filename'], licensor=kwargs['licensor'], temp_table_name=temp_staging_raw_table, ) self.execute_query(sql_loader, query_name, params) def clean_staging_raw_table(self, staging_raw_table, date, **kwargs): """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). """ params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_raw_table=staging_raw_table, download_date=date, vendor_id=kwargs['vendor_id'] ) self.execute_query(sql_loader, 'delete_from_staging_raw', params) def delete_old_snapshots(self, staging_raw_table): """Delete old snapshots from staging raw tables. Args: staging_raw_table (str): A table name in Snowflake. """ params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_raw_table=staging_raw_table ) self.execute_query( sql_loader, 'delete_from_staging_raw_old_snapshots', params)