"""Snowflake connector class for the AWA workflow.""" from datetime import datetime from snowflake_connector.etl_connector import SQLLoader from feed_ingestion.common.fact_analytics_sf.base_executor \ import SnowflakeSQLExecutorFA from feed_ingestion.common.staging_raw_sf.base_executor \ import SnowflakeSQLExecutorSR from feed_ingestion.flows.awa import config from feed_ingestion.util.snowflake.errors import JSONParserLoading # Load SQL templates sql_loader = SQLLoader(__file__) class AWA(SnowflakeSQLExecutorFA, 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 '_'.join( [config.feed_name, self.licensor]) @property def staging_raw_table(self): """Name of staging raw table. Returns: str: staging_raw table name """ raise Exception('Implement in child class') @property def feedid(self): """ID of the feed. Returns: int: Feed id """ return config.feedid @property def storeid(self): """ID of the store. Returns: int: Store id """ return config.storeid @property def licensor(self): """Licensor. Returns: str: licensor """ raise Exception('Implement in child class') def temp_staging_fact_table(self, date): """Name of temp staging fact table. Returns: str: temp_staging_fact_unpivot table name """ raise Exception('Implement in child class') 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. """ report = kwargs['report'] query_name = f'create_temp_staging_raw_{report}' 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, query_name, params) def load_temp_staging_raw_table( self, temp_staging_raw_table, aws, key_dir, files, **kwargs): """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): Custom arguments. files (str): List of files to ingest. """ aws_params = self.get_aws_params() query_filename = f'load_temp_staging_raw_{self.licensor}' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], temp_staging_raw_table=temp_staging_raw_table, s3_path=key_dir, files=files, on_error_action='ABORT_STATEMENT', **aws_params) return [JSONParserLoading(*e) for e in self.fetchall_query( sql_loader, query_filename, params)] def load_staging_raw_table( self, temp_staging_raw_table, staging_raw_table, date, **kwargs): """Load staging raw table with activity, user, and playlist files. Args: temp_staging_raw_table (str): A table name in Snowflake. staging_raw_table (str): Name of staging raw table. date (str): Date in YYYY-MM-DD format. kwargs (dict): Custom arguments with temporary table names and filenames for user, playlist and activity. """ report = kwargs['report'] params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_raw_table=staging_raw_table, temp_staging_raw_table=temp_staging_raw_table, download_date=date, **kwargs) query_name = f'load_staging_raw_{report}' self.execute_query(sql_loader, query_name, params) def clean_staging_raw_table_(self, date): """Delete data from staging raw table for specific date. Args: date (str): Date of the data being process (YYYY-MM-DD). """ self.clean_staging_raw_table( db=self.sf_config['db'], schema=self.sf_config['schema'], date=date, staging_raw_table=self.staging_raw_table) def drop_temp_staging_raw_table(self, table_name): """Delete temp staging raw table. Args: table_name (str): Name of table to be deleted. """ self.execute_query( sql_loader, 'delete_tmp_table', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], table_name=table_name)) def create_temp_staging_fact_table(self, report, date): """Create a temporary and unpivoted staging fact table. Args: report (str): report name. date (str): Date of the data being process (YYYY-MM-DD). """ query_name = f'create_temp_staging_fact_{report}' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], temp_staging_fact_table=self .temp_staging_fact_table(date)) self.execute_query(sql_loader, query_name, params) def load_temp_staging_fact_table(self, date, report, **kwargs): """Load the temporary and unpivoted staging fact table. Args: report (str): report name. date (str): Date of the data being process (YYYY-MM-DD). kwargs (dict): Custom arguments with report type. """ query_name = f'load_temp_staging_fact_{report}' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], temp_staging_fact_table=self .temp_staging_fact_table(date), staging_raw_table=self.staging_raw_table, licensor=self.licensor, download_date=date) self.execute_query(sql_loader, query_name, params) def load_staging_fact_table(self, date): """Load staging fact_analytics table from staging_raw table. Args: date (str): Date of the data being process (YYYY-MM-DD). """ query_name = 'load_staging_fact' if self.licensor == 'altafonte': query_name = 'load_staging_fact_altafonte' self.execute_query( sql_loader, query_name, params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_fact_table=self.staging_fact_table(date), temp_staging_fact_table=self .temp_staging_fact_table(date), feedid=self.feedid, storeid=self.storeid, licensor=self.licensor, download_date=date)) def load_fact_error_data(self, date): """Load unmatched data into fact_analytics_error. Args: date (str): Date of the data being process (YYYY-MM-DD). """ query_name = 'load_fact_analytics_error' if self.licensor == 'altafonte': query_name = 'load_fact_analytics_error_altafonte' self.execute_query( sql_loader, query_name, params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], fact_error_table=self.fact_error_table, staging_fact_table=self.staging_fact_table(date), temp_staging_fact_table=self .temp_staging_fact_table(date), reportdate=date, feedid=self.feedid, licensor=self.licensor, storeid=self.storeid)) def delete_from_fact_table(self, date, *args): """Delete rows from fact_analytics table for the given date. Args: date (str): Date of the data being process (YYYY-MM-DD). """ query_name = 'delete_from_fact_table' self.execute_query( sql_loader, query_name, params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], fact_table=self.fact_table, storeid=self.storeid, reportdate=date, feedid=self.feedid, licensor=self.licensor)) class AwaTheOrchard(AWA): """Helper class to abstract Snowflake operations.""" @property def licensor(self): """Licensor.""" return 'theorchard' def temp_staging_fact_table(self, date): """Name of staging fact table. Returns: str: staging_fact table name """ date_obj = datetime.strptime(date, '%Y-%m-%d') return config.licensors['theorchard']['reports']['play_summary'][ 'temp_staging_fact_table'].format(date=date_obj) @property def staging_raw_table(self): """Name of staging raw table. Returns: str: staging_raw table name """ return config.licensors['theorchard']['reports']['play_summary'][ 'staging_raw_table'] class AwaSmej(AWA): """Helper class to abstract Snowflake operations.""" @property def licensor(self): """Licensor.""" return 'smej'