"""Snowflake connector class for the iTunes Hides tasks.""" from snowflake_connector.etl_connector import SQLLoader from feed_ingestion.common.staging_raw_sf.base_executor import \ SnowflakeSQLExecutorSR from feed_ingestion.flows.itunes_hides import config sql_loader = SQLLoader(__file__) class ITunesHidesSF(SnowflakeSQLExecutorSR): """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): """Create a temporary staging raw table. Args: temp_staging_raw_table (str): A table name in Snowflake. """ sql_template = 'create_temp_table' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], temp_staging_raw_table=temp_staging_raw_table) self.execute_query(sql_loader, sql_template, params) def load_temp_staging_raw_table( self, temp_staging_raw_table, aws, key_dir): """Load temp staging raw table. Args: temp_staging_raw_table (str): A table name in Snowflake. aws (dict): AWS credentials to fill a template of COPY SQL statement. key_dir (str): A S3 path to load files from. """ aws_params = self.get_aws_params() sql_template = 'load_temp_table' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], temp_staging_raw_table=temp_staging_raw_table, s3_path=key_dir, **aws_params ) self.execute_query(sql_loader, sql_template, params) def clean_staging_raw_table(self, staging_raw_table, date, **kwargs): """Delete rows from previous unsuccessful workflow run. The flow is cumulative. staging_raw_table is not required. Args: staging_raw_table (str): A table name in Snowflake. date (str): Date of the data being process (YYYY-MM-DD). """ return None def clean_snapshot_table(self, snapshot_table, date): """Delete rows from previous unsuccessful workflow run. Args: snapshot_table (str): A table name in Snowflake. date (str): Date of the data being process (YYYY-MM-DD). """ sql_template = 'clean_snapshot_table' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], snapshot_table=snapshot_table, ingest_date=date ) self.execute_query(sql_loader, sql_template, params) def load_snapshot_table(self, temp_staging_raw_table, snapshot_table): """Load rows from temp staging raw table into snapshot table. Args: temp_staging_raw_table (str): A table name in Snowflake. snapshot_table (str): A table name in Snowflake. """ sql_template = 'load_snapshot_table' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], temp_staging_raw_table=temp_staging_raw_table, snapshot_table=snapshot_table ) self.execute_query(sql_loader, sql_template, params) def remove_previous_snapshots( self, snapshot_table, snapshots_amount_to_keep): """Remove previous snapshots leaving only specified amount. Args: snapshot_table (str): A table name in Snowflake. snapshots_amount_to_keep (int): Amount of snapshots to keep. """ sql_template = 'remove_previous_snapshots' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], snapshot_table=snapshot_table, snapshots_amount=str(snapshots_amount_to_keep) ) self.execute_query(sql_loader, sql_template, params) def load_staging_raw_table( self, temp_staging_raw_table, staging_raw_table, date, **kwargs): """Load the temp_staging_raw data to the feed's staging_raw table. Args: temp_staging_raw_table (str): A table name in Snowflake. staging_raw_table (str): A table name in Snowflake. date (str): Date of the data being process (YYYY-MM-DD). """ query_name = 'load_staging_raw' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], temp_staging_raw_table=temp_staging_raw_table, staging_raw_table=staging_raw_table) self.execute_query(sql_loader, query_name, params)