"""Snowflake connector class for the Pandora's tasks.""" from snowflake_connector.etl_connector import SQLLoader from feed_ingestion.common.fact_analytics_sf.base_executor \ import SnowflakeSQLExecutorFA from feed_ingestion.flows.pandora import config from feed_ingestion.flows.pandora import util from feed_ingestion.util.snowflake.errors import JSONParserLoading sql_loader = SQLLoader(__file__) class Pandora(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 licensor(self): """Licensor name. Returns str: Licensor name. """ raise NotImplementedError() @property def feed_name(self): """Name of the feed. Should match dir name of this feed, feed_name in config.py of a feed, and a key of executors dict in feed_ingestion/common/ fact_analytics_sf/__init__.py file. Returns: str: Feed name, e.g. 'deezer_daily_sf'. """ return '_'.join([config.feed_name, self.licensor]) @property def feedid(self): """Id of the feed. Returns: str: Feed id """ return config.feedid @property def storeid(self): """Storeid of feed data. Should match dim_store and feed config value. Returns: integer: Feed's storeid. """ return 708 @property def staging_raw_table(self): """Name of the staging_raw table for the feed. Returns: str: staging_raw_{feed} table name. """ return config.snowflake_table_names['staging_raw'] def create_streams_temp_staging_raw_table( self, temp_staging_raw_table, **kwargs): """Create a temporary staging raw table for streams data. Args: temp_staging_raw_table (str): A table name in Snowflake. """ sql_loader = SQLLoader(__file__, date=kwargs.get('date_for_sqlloader')) self.execute_query( sql_loader, 'create_streams_temp_staging_raw', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], streams_temp_staging_raw_table=temp_staging_raw_table)) def create_metadata_temp_staging_raw_table( self, temp_staging_raw_table, **kwargs): """Create a temporary staging raw table for streams data. Args: temp_staging_raw_table (str): A table name in Snowflake. """ sql_loader = SQLLoader(__file__, date=kwargs.get('date_for_sqlloader')) self.execute_query( sql_loader, 'create_metadata_temp_staging_raw', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], metadata_temp_staging_raw_table=temp_staging_raw_table)) def load_temp_staging_raw_table( self, temp_staging_raw_table, aws, key_dir, **kwargs): """Load temp staging raw table with a Pandora raw file. 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. """ error_limit = kwargs.get('error_limit') if error_limit: if not isinstance(error_limit, int): raise ValueError('error_limit should be int') on_error_action = f'SKIP_FILE_{error_limit}' else: on_error_action = 'ABORT_STATEMENT' if kwargs.get('error_on_column_count_mismatch', '').lower() == 'false': error_on_column_count_mismatch = 'FALSE' else: error_on_column_count_mismatch = 'TRUE' aws_params = self.get_aws_params() params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], temp_staging_raw_table=temp_staging_raw_table, s3_path=key_dir, on_error_action=on_error_action, error_on_column_count_mismatch=error_on_column_count_mismatch, **aws_params ) return [JSONParserLoading(*e) for e in self.fetchall_query( sql_loader, 'load_temp_staging_raw', params)] 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). """ self.execute_query( sql_loader, 'delete_from_staging_raw', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_raw_table=staging_raw_table, date=date, licensor=self.licensor)) def load_staging_raw_table( self, date, processed_datetime, filename, staging_raw_table, temp_streams_table, temp_metadata_table, **kwargs): """Load the temp_staging_raw data to the feed's staging_raw table. Args: date (str): Date of the data being process (YYYY-MM-DD). processed_datetime (str): A single processeddaytime to use through all the tables during the workflow run. filename (str): Filename of the streams data from which a particular temp streams table was loaded. staging_raw_table (str): A table name in Snowflake. temp_streams_table (str): Name of the temp streams table. temp_metadata_table (str): Name of the temp metadata table. """ sql_loader = SQLLoader(__file__, date=kwargs.get('date_for_sqlloader')) self.execute_query( sql_loader, 'load_staging_raw', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], filename=filename, date=date, processed_datetime=processed_datetime, streams_temp_staging_raw_table=temp_streams_table, metadata_temp_staging_raw_table=temp_metadata_table, staging_raw_table=staging_raw_table, licensor=self.licensor)) def update_dimension_table(self, date, table_name): """Update dimension table with the new data. Args: date (str): Date of the data being process (YYYY-MM-DD). table_name (str): A table to update (corresponding query should be placed in the queries/ folder of the flow). """ return self._update_dimension_table( date, table_name, sql_loader, licensor=self.licensor) 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). """ sql_template = sql_loader.load_query('load_staging_fact') params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], storeid=self.storeid, staging_fact_table=self.staging_fact_table(date), staging_raw_table=self.staging_raw_table, feedid=self.feedid, reportdate=date, licensor=self.licensor) sql_template, non_identifier_params = ( self.validator.format_identifiers(sql_template, params)) sql = sql_template.format( countryname_cases=util.build_cases_by_country( config.country_mappings), currencyid_cases=util.build_cases_by_country( config.currencyid_mappings)) self.execute(sql, params=non_identifier_params) 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). """ sql_template = sql_loader.load_query('load_fact_analytics_error') params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], fact_error_table=self.fact_error_table, storeid=self.storeid, staging_fact_table=self.staging_fact_table(date), staging_raw_table=self.staging_raw_table, feedid=self.feedid, reportdate=date, licensor=self.licensor) sql_template, non_identifier_params = ( self.validator.format_identifiers(sql_template, params)) sql = sql_template.format( countryname_cases=util.build_cases_by_country( config.country_mappings), currency_cases=util.build_cases_by_country( config.currency_mappings), currencyid_cases=util.build_cases_by_country( config.currencyid_mappings)) self.execute(sql, params=non_identifier_params) def _delete_from_fact_table(self, table, date, **kwargs): """Delete rows in fact table with the current run date. This is required for the workflow to be idempotent, and to avoid row duplication. Before we'll load rows for a specific day to the fact tables, we have to delete rows which were added by previous (allegedly unsuccessful workflow run). Args: table (str): Either fact_analytics, either fact_analytics_error. date (str): Date of the data being process (YYYY-MM-DD). kwargs (dict): Custom arguments. """ sql_template = sql_loader.load_query('delete_from_fact_table') params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], fact_table=table, reportdate=date, feedid=self.feedid, licensor=self.licensor, storeid=self.storeid) params.update(kwargs) sql_template, non_identifier_params = ( self.validator.format_identifiers(sql_template, params)) return self.fetchone(sql_template, params=non_identifier_params)[0] class PandoraTheOrchardFA(Pandora): """Helper class to abstract Snowflake operations.""" @property def licensor(self): """Licensor.""" return 'theorchard' class PandoraSMEFA(Pandora): """Helper class to abstract Snowflake operations.""" @property def licensor(self): """Licensor.""" return 'sme'