"""Snowflake connector class for Amazon Unlimited tasks.""" 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.amazon_datapulse import config as datapulse_config from feed_ingestion.flows.amazon_music import config from feed_ingestion.util.snowflake.errors import JSONParserLoading # Load SQL templates sql_loader = SQLLoader(__file__) class AmazonMusicSFExecutor( SnowflakeSQLExecutorFA, SnowflakeSQLExecutorSR): """Helper class to abstract Snowflake operations.""" @property def report_name(self): """Report name. Returns: str: Report name. """ raise NotImplementedError() @property def licensor(self): """Licensor name. Shall be overridden in ancestors classes/ Returns: str: Licensor. Shall be one of config.licensors. """ 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. 'amazon_Music'. """ return '_'.join([config.feed_name, self.report_name]) @property def feedid(self): """Id of the feed. Returns: str: Feed id """ return config.reports[self.report_name]['feedid'] @property def staging_raw_table(self): """Name of the staging_raw table for the feed. Returns: str: staging_raw_{feed} table name. """ return ( config.reports[self.report_name]['staging_raw_table'] + self.licensor_suffix) @property def aggregated_staging_raw_table(self): """Name of the aggregated_staging_raw table for the feed. Returns: str: staging_raw_{feed} table name. """ return config.aggregated_staging_raw_table @property def storeid(self): """Storeid of feed data. Should match dim_store and feed config value. Returns: integer: Feed's storeid. """ return config.storeid @property def licensor_suffix(self): """Licensor suffix for queries and tables. Returns: string: The empty string by default. """ return '' 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. """ sql_loader = SQLLoader(__file__, date=kwargs.get('date_for_sqlloader')) report_type = kwargs['report_type'].lower() query_name = 'create_temp_staging_raw_{}_table{}'.format( report_type, self.licensor_suffix) 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, **kwargs): """Load temp staging raw table. Args: temp_staging_raw_table (str): A table name in Snowflake. aws (dict): AWS credentials to fill template of COPY SQL statement (deprecated) key_dir (str): Custom arguments. """ sql_loader = SQLLoader(__file__, date=kwargs.get('date_for_sqlloader')) 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, # I concern about error handling. In case of error # in files it returns one string: # "Copy executed with 0 files processed." which # fails of parsing response # I'd rather get exception with clean error message on_error_action='SKIP_FILE_{}'.format(kwargs['error_limit']), error_on_column_count_mismatch=error_on_column_count_mismatch, **aws_params, ) query = f'load_temp_staging_raw{self.licensor_suffix}' return [JSONParserLoading(*e) for e in self.fetchall_query( sql_loader, query, params)] def load_staging_raw_from_datapulse( self, date, report_name, licensor, **kwargs): """Load the unified staging_raw_amazon_music from amazon_datapulse. Renders queries/load_staging_raw_datapulse.sql with the appropriate target/source identifiers and filter values, then executes it. All reports and licensors share the single unified config.DATAPULSE_STAGING_RAW_TABLE target; each run inserts only its own service x licensor slice. Args: date (str): Reporting date (YYYY-MM-DD). report_name (str): One of config.DATAPULSE_SERVICE_BY_REPORT keys. licensor (str): One of config.DATAPULSE_LICENSORS. kwargs (dict): Additional kwargs (currently unused). """ play_events_db, play_events_schema, play_events_table = ( datapulse_config.get_staging_raw_database_schema( 'daily_play_events')) customer_meta_db, customer_meta_schema, customer_metadata_table = ( datapulse_config.get_staging_raw_database_schema( 'daily_customer_metadata')) playlist_meta_db, playlist_meta_schema, playlist_metadata_table = ( datapulse_config.get_staging_raw_database_schema( 'daily_playlist_metadata')) params = dict( # staging_raw_amazon_music lives in the same database/schema as the # datapulse source tables (e.g. consumer_reporting.amazon), not the # analytics db/schema of this flow's sf_config. target_db=play_events_db, target_schema=play_events_schema, staging_raw_table=config.DATAPULSE_STAGING_RAW_TABLE, play_events_db=play_events_db, play_events_schema=play_events_schema, play_events_table=play_events_table, customer_meta_db=customer_meta_db, customer_meta_schema=customer_meta_schema, customer_metadata_table=customer_metadata_table, playlist_meta_db=playlist_meta_db, playlist_meta_schema=playlist_meta_schema, playlist_metadata_table=playlist_metadata_table, download_date=date, service=config.DATAPULSE_SERVICE_BY_REPORT[report_name], entity_name=config.DATAPULSE_ENTITY_NAME_BY_LICENSOR[licensor], licensor=licensor, feedid=config.reports[report_name]['feedid'], ) self.execute_query( sql_loader, 'load_staging_raw_datapulse', params=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. """ sql_loader = SQLLoader(__file__, date=kwargs.get('date_for_sqlloader')) staging_raw_suffix = '' if self.licensor_suffix == '_altafonte' \ else self.licensor_suffix params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_raw_table=f'{staging_raw_table}{staging_raw_suffix}', download_date=date, **kwargs) query = f'load_staging_raw{self.licensor_suffix}' self.execute_query(sql_loader, query, params) 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). """ query = f'update_{table_name}{self.licensor_suffix}' # datapulse dim queries read from the unified staging_raw_amazon_music # table, which lives in the datapulse database/schema (e.g. # consumer_reporting.amazon), not the analytics db/schema of the dims. if table_name.endswith('_datapulse'): staging_raw_db, staging_raw_schema, _ = ( datapulse_config.get_staging_raw_database_schema( 'daily_play_events')) staging_raw_table = config.DATAPULSE_STAGING_RAW_TABLE else: staging_raw_db = self.sf_config['db'] staging_raw_schema = self.sf_config['schema'] staging_raw_table = self.staging_raw_table sql_template = sql_loader.load_query(query) sql, non_identifier_params = self.validator.format_identifiers( sql_template, params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_raw_db=staging_raw_db, staging_raw_schema=staging_raw_schema, date=date, storeid=self.storeid, feedid=self.feedid, staging_raw_table=staging_raw_table)) return self.fetchone( sql, params=non_identifier_params, dict_cursor=True) def staging_fact_table(self, date): """Get name of temp staging_fact_analytics_ table. Args: date (str): Date of the data being process (YYYY-MM-DD). Returns: str: Name of the staging_fact_analytics table for a feed. """ return 'staging_fact_analytics_{feed_name}_{licensor}_{date}'.format( feed_name=self.feed_name, licensor=self.licensor, date=date.replace('-', '')) def load_staging_fact_table(self, date, source=None, **kwargs): """Load staging fact_analytics table from staging_raw table. For datapulse the rows are read from the unified config.DATAPULSE_STAGING_RAW_TABLE (which lives in the datapulse database/schema, not this flow's analytics sf_config, and holds every report x licensor) instead of the per-licensor self.staging_raw_table. A single load_staging_fact_datapulse query serves all licensors; it is scoped by feedid + licensor. The dimension tables stay in the analytics db/schema. Args: date (str): Date of the data being process (YYYY-MM-DD). source (str): Source of the staging data (config.sources). kwargs (dict): Additional kwargs forwarded by the generic task. """ if source == config.SOURCE_DATAPULSE: staging_raw_db, staging_raw_schema, _ = ( datapulse_config.get_staging_raw_database_schema( 'daily_play_events')) self.execute_query( sql_loader, 'load_staging_fact_datapulse', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_fact_table=self.staging_fact_table(date), staging_raw_db=staging_raw_db, staging_raw_schema=staging_raw_schema, staging_raw_table=config.DATAPULSE_STAGING_RAW_TABLE, reportdate=date, storeid=self.storeid, licensor=self.licensor, feedid=self.feedid)) return self.execute_query( sql_loader, f'load_staging_fact_{self.licensor}', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_fact_table=self.staging_fact_table(date), staging_raw_table=self.staging_raw_table, reportdate=date, storeid=self.storeid, licensor=self.licensor, feedid=self.feedid)) def delete_from_aggregated_skips_and_saves(self, date, licensor, **kwargs): """Delete from load_aggregated_skips_and_saves. Args: date (str): Date of the data being process (YYYY-MM-DD). licensor (str): one of config.licensors """ self.execute_query( sql_loader, 'delete_from_aggregated_skips_and_saves', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], reportdate=date, feedid=self.feedid, licensor=licensor)) def load_aggregated_skips_and_saves(self, date, **kwargs): """Load load_aggregated_skips_and_saves from staging fact_analytics. Args: date (str): Date of the data being process (YYYY-MM-DD). kwargs (dict): to hold additional kwargs parameters from task. """ self.execute_query( sql_loader, 'load_aggregated_skips_and_saves', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_fact_table=self.staging_fact_table(date), reportdate=date, storeid=self.storeid, feedid=self.feedid)) def load_fact_data(self, date, **kwargs): """Load matched data into fact_analytics. For datapulse the matched rows are selected straight from the staging fact table (already scoped to feed + licensor + date), dropping the legacy staging_raw join that only served to filter by download_date. Args: date (str): Date of the data being process (YYYY-MM-DD). kwargs (dict): Additional kwargs forwarded by the generic task (includes source). """ if kwargs.get('source') == config.SOURCE_DATAPULSE: sql_template = sql_loader.load_query( 'load_fact_analytics_datapulse') params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], fact_table=self.fact_table, staging_fact_analytics_table=self.staging_fact_table(date)) sql_template, non_identifier_params = ( self.validator.format_identifiers(sql_template, params)) self.execute(sql_template, params=non_identifier_params) return sql_template = sql_loader.load_query('load_fact_analytics') params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], fact_table=self.fact_table, staging_fact_analytics_table=self.staging_fact_table(date), staging_raw_table=self.staging_raw_table, reportdate=date) sql_template, non_identifier_params = ( self.validator.format_identifiers(sql_template, params)) self.execute(sql_template, params=non_identifier_params) def delete_from_fact_table(self, 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). Note: overriding this method happened only because that parent's sql_loader cannot load query from this flow's queries dir. It need to move sql_loader as overridable attribute for root class. Args: 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=self.fact_table, reportdate=date, feedid=self.feedid) 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] def delete_from_fact_error_table(self, date, **kwargs): """Delete rows from fact_analytics_error 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_error_table, date=date, staging_raw_table=self.staging_raw_table, storeid=self.storeid, reportdate=date, feedid=self.feedid, licensor=self.licensor)) def load_fact_error_data(self, date, **kwargs): """Load unmatched data into fact_analytics_error. For datapulse the unified config.DATAPULSE_STAGING_RAW_TABLE (which lives in the datapulse database/schema and holds every report x licensor) is the driving table; it is read via separate staging_raw_db/staging_raw_schema identifiers and scoped by feedid + licensor. The dim and staging fact tables stay in the analytics db/schema. Args: date (str): Date of the data being process (YYYY-MM-DD). kwargs (dict): Additional kwargs forwarded by the generic task (includes source). """ if kwargs.get('source') == config.SOURCE_DATAPULSE: staging_raw_db, staging_raw_schema, _ = ( datapulse_config.get_staging_raw_database_schema( 'daily_play_events')) self.execute_query( sql_loader, 'load_fact_analytics_error_datapulse', 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), staging_raw_db=staging_raw_db, staging_raw_schema=staging_raw_schema, staging_raw_table=config.DATAPULSE_STAGING_RAW_TABLE, reportdate=date, storeid=self.storeid, licensor=self.licensor, feedid=self.feedid)) return self.execute_query( sql_loader, f'load_fact_analytics_error{self.licensor_suffix}', 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), staging_raw_table=self.staging_raw_table, reportdate=date, storeid=self.storeid, licensor=self.licensor, feedid=self.feedid)) def clean_staging_raw_table(self, staging_raw_table, date, **kwargs): """Delete rows from previous unsuccessful workflow run. For datapulse the bootstrap passes the unified config.DATAPULSE_STAGING_RAW_TABLE, which holds every report x licensor and lives in the datapulse database/schema (e.g. consumer_reporting. amazon), not this flow's analytics sf_config db/schema. The delete must therefore target that table there and be scoped by feedid so it only removes this report's slice for the date (not other reports' rows). Args: staging_raw_table (str): A table name in Snowflake. date (str): Date of the data being process (YYYY-MM-DD). """ if staging_raw_table == config.DATAPULSE_STAGING_RAW_TABLE: staging_raw_db, staging_raw_schema, _ = ( datapulse_config.get_staging_raw_database_schema( 'daily_play_events')) self.execute_query( sql_loader, 'delete_from_staging_raw_datapulse', dict( db=staging_raw_db, schema=staging_raw_schema, staging_raw_table=config.DATAPULSE_STAGING_RAW_TABLE, date=date, licensor=self.licensor, feedid=self.feedid)) return query = 'delete_from_staging_raw' params = dict( db=self.sf_config['db'], schema=self.sf_config['schema'], staging_raw_table=self.staging_raw_table, date=date, licensor=self.licensor) if ((self.licensor == 'awal' or self.licensor == 'altafonte') and kwargs.get('use_s3') in ['True', True]): query = f'delete_from_staging_raw{self.licensor_suffix}' params.update(dict(countries=kwargs.get('country_list'))) self.execute_query(sql_loader, query, params) def clean_aggregated_staging_raw_table(self, date, orgs, source=None): """Delete rows from previous unsuccessful workflow run. For datapulse the aggregated load writes the mapped licensor (sme/theorchard/awal/red) rather than a filename-derived store-code org, so the delete is scoped by feedid + licensor instead of the store-code orgs list (see load_aggregated_staging_raw_table). Args: date (str): Date of the data being process (YYYY-MM-DD). orgs (list of str): List of orgs. Org is stored in licensor column. source (str): Source of the staging data (config.sources). """ if source == config.SOURCE_DATAPULSE: self.execute_query( sql_loader, 'delete_from_aggregated_staging_raw_datapulse', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], aggregated_staging_raw_table=( self.aggregated_staging_raw_table), date=date, licensor=self.licensor, feedid=self.feedid)) return self.execute_query( sql_loader, 'delete_from_aggregated_staging_raw', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], aggregated_staging_raw_table=self.aggregated_staging_raw_table, staging_raw_table=self.staging_raw_table, date=date, orgs=orgs, feedid=self.feedid)) def load_aggregated_staging_raw_table(self, date, source=None): """Load rows in aggregated staging raw table. For datapulse the rows are read from the unified config.DATAPULSE_STAGING_RAW_TABLE (which lives in the datapulse database/schema, not this flow's analytics sf_config, and holds every report x licensor) instead of the per-report self.staging_raw_table. It is therefore scoped by feedid + licensor so each per-report/ per-licensor workflow only aggregates its own slice. Args: date (str): Date of the data being process (YYYY-MM-DD). source (str): Source of the staging data (config.sources). """ if source == config.SOURCE_DATAPULSE: staging_raw_db, staging_raw_schema, _ = ( datapulse_config.get_staging_raw_database_schema( 'daily_play_events')) self.execute_query( sql_loader, 'load_aggregated_staging_raw_datapulse', params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], aggregated_staging_raw_table=( self.aggregated_staging_raw_table), staging_raw_db=staging_raw_db, staging_raw_schema=staging_raw_schema, staging_raw_table=config.DATAPULSE_STAGING_RAW_TABLE, date=date, licensor=self.licensor, feedid=self.feedid)) return query_suffix = '' if self.licensor_suffix == '_altafonte' \ else self.licensor_suffix query = f'load_aggregated_staging_raw{query_suffix}' self.execute_query( sql_loader, query, params=dict( db=self.sf_config['db'], schema=self.sf_config['schema'], aggregated_staging_raw_table=self.aggregated_staging_raw_table, staging_raw_table=self.staging_raw_table, date=date, feedid=self.feedid)) class AmazonUnlimitedSFExecutor(AmazonMusicSFExecutor): """Helper class to abstract Snowflake operations.""" @property def report_name(self): """Report name. Returns: str: Report name. """ return 'unlimited' class AmazonPrimeSFExecutor(AmazonMusicSFExecutor): """Helper class to abstract Snowflake operations.""" @property def report_name(self): """Report name. Returns: str: Report name. """ return 'prime' class AmazonAdSupportedSFExecutor(AmazonMusicSFExecutor): """Helper class to abstract Snowflake operations.""" @property def report_name(self): """Report name. Returns: str: Report name. """ return 'adsupported' class SmeLicensorMixin: """Licensor mixin for AmazonMusicSFExecutor.""" @property def licensor(self): """Licensor value.""" return 'sme' class OrchardLicensorMixin: """Licensor mixin for AmazonMusicSFExecutor.""" @property def licensor(self): """Licensor value.""" return 'theorchard' class AwalLicensorMixin: """Licensor mixin for AmazonMusicSFExecutor.""" @property def licensor(self): """Licensor value.""" return 'awal' @property def licensor_suffix(self): """Licensor suffix for queries and tables. Returns: string: The '_awal' suffix for awal queries and tables. """ return '_awal' class AltafonteLicensorMixin: """Licensor mixin for AmazonMusicSFExecutor.""" @property def licensor(self): """Licensor value.""" return 'altafonte' @property def licensor_suffix(self): """Licensor suffix for queries and tables. Returns: string: The '_altafonte' suffix for altafonte queries and tables. """ return '_altafonte' class AmazonSmeUnlimitedSFExecutor(SmeLicensorMixin, AmazonUnlimitedSFExecutor): """Specific snowflake executor class.""" pass class AmazonSmePrimeSFExecutor(SmeLicensorMixin, AmazonPrimeSFExecutor): """Specific snowflake executor class.""" pass class AmazonSmeAdSupportedSFExecutor(SmeLicensorMixin, AmazonAdSupportedSFExecutor): """Specific snowflake executor class.""" pass class AmazonOrchardUnlimitedSFExecutor(OrchardLicensorMixin, AmazonUnlimitedSFExecutor): """Specific snowflake executor class.""" pass class AmazonOrchardPrimeSFExecutor(OrchardLicensorMixin, AmazonPrimeSFExecutor): """Specific snowflake executor class.""" pass class AmazonOrchardAdSupportedSFExecutor(OrchardLicensorMixin, AmazonAdSupportedSFExecutor): """Specific snowflake executor class.""" pass class AmazonAwalUnlimitedSFExecutor( AwalLicensorMixin, AmazonUnlimitedSFExecutor): """Specific snowflake executor class.""" pass class AmazonAwalPrimeSFExecutor( AwalLicensorMixin, AmazonPrimeSFExecutor): """Specific snowflake executor class.""" pass class AmazonAwalAdSupportedSFExecutor( AwalLicensorMixin, AmazonAdSupportedSFExecutor): """Specific snowflake executor class.""" pass class AmazonAltafonteUnlimitedSFExecutor( AltafonteLicensorMixin, AmazonUnlimitedSFExecutor): """Specific snowflake executor class.""" @property def staging_raw_table(self): """Name of the staging_raw table for the feed. Returns: str: staging_raw_{feed} table name. """ return config.reports[self.report_name]['staging_raw_table'] class AmazonAltafontePrimeSFExecutor( AltafonteLicensorMixin, AmazonPrimeSFExecutor): """Specific snowflake executor class.""" @property def staging_raw_table(self): """Name of the staging_raw table for the feed. Returns: str: staging_raw_{feed} table name. """ return config.reports[self.report_name]['staging_raw_table'] class AmazonAltafonteAdSupportedSFExecutor( AltafonteLicensorMixin, AmazonAdSupportedSFExecutor): """Specific snowflake executor class.""" @property def staging_raw_table(self): """Name of the staging_raw table for the feed. Returns: str: staging_raw_{feed} table name. """ return config.reports[self.report_name]['staging_raw_table']