import logging import time from functools import wraps from itertools import chain from typing import Dict, List import numpy as np import pandas as pd from utils.aws_connectors import pd_read_sql from utils.aws_connectors import run_query from utils.base_attributes_generation import get_fan_attribute_values from utils.data_model_utils import ( generate_collection, update_collection, get_table_structure, get_collection_type, write_csv_to_efs, string_to_db, ) from airflow.progress.controller import get_enrichment_progress from utils import appsync_communication as appsync component_logger = logging.getLogger().getChild("enrichment.enrichment") class MissingInputFieldException(Exception): pass def timeit(func): @wraps(func) def timed(self, *args, **kwargs): method_name_length = len(func.__name__) enrichment_name_length = len(self.enrichment_name) print( f'--- {self.enrichment_name} {"".join(["-" for x in range(enrichment_name_length, 34)])} started --- Method: {func.__name__} {"".join(["-" for x in range(method_name_length, 35)])}' ) ts = time.time() try: result = func(self, *args, **kwargs) except Exception as e: te = time.time() print( f'--- {self.enrichment_name} {"".join(["-" for x in range(enrichment_name_length, 34)])} failed ---- Method: {func.__name__} {"".join(["-" for x in range(method_name_length, 35)])} at: {te - ts}' ) raise e te = time.time() print( f'--- {self.enrichment_name} {"".join(["-" for x in range(enrichment_name_length, 34)])} finished -- Method: {func.__name__} {"".join(["-" for x in range(method_name_length, 35)])} took: {te - ts}' ) return result return timed def timeit_static(func): @wraps(func) def timed(*args, **kwargs): method_name = func.__name__ print(f"--- {method_name} started ---") ts = time.time() try: result = func(*args, **kwargs) except Exception as e: te = time.time() print(f"--- {method_name} failed at: {(te - ts):.2f} seconds.") raise e te = time.time() print(f"--- {method_name} took: {(te - ts):.2f} seconds.") return result return timed class Enrichment: """ Parent enrichment class handling all the methods: Main behaviour expectations: - INSERT INTO collection - if needed - INSERT INTO collection_fan - if needed - INSERT INTO fan_attribute_table - if needed - INSERT INTO Attribute table - if needed - Re running same enrichment for the same id several times will not make a difference and will not make new API calls - Total number of available distinct fan_attribute_fan_ids = number of fan_attribute_fan_ids in collection_fan where collection_id = base_collection_id """ source_attribute_ids: List[int] = [] result_attribute_ids: List[int] = [] # Describes the way enrichment is obtaining Source Data # 'loose' - get all the data from fan_attribute ignoring collection_id # 'strict' - get all the data from fan_attribute restricted by base_collection + it's child collections data_gathering_rule = "loose" get_previously_enriched_data = False # Temporary work around for enrGender !TODO Needs Proper Handling def __init__(self, collection_id: int, schema: str, user_id: str, default_value: str = None, **kwargs): self.base_collection_id = int(collection_id) self.schema = schema self.user_id = user_id self.enrichment_name = self.__class__.__name__ self.collection_type = "" self.system_fields = pd.DataFrame({}) self.available_data = {"source": pd.DataFrame({})} self.available_system_fields: List[int] = [] self.collection_ids: Dict[str, List[int]] = {"source": [], "child": []} self.enrichment_collection_id = None self.collection_fan_ids: List[int] = [] self.result_df = pd.DataFrame({}) self.fan_attribute = pd.DataFrame({}) self.default_value = default_value self._async_task = None def _prepare(self, populate_tables=True): """ Running all the required methods to prepare class for actual enrichment""" prepared = True """ Making basic preliminary checks to ensure that we can run enrichment """ self._send_runtime_estimation() self._get_collection_ids() self._check_collection_type() self._get_available_system_fields() ok_to_continue_enrichment = self._check_available_system_fields() if not ok_to_continue_enrichment: prepared = False return prepared """ Data Gathering and preparation """ self._get_all_system_fields() self._get_available_data() self._get_collection_fan_ids() """ Populating main tables """ if populate_tables: self._populate_collection_table() self._populate_collection_fan_table() self._update_collection_status("processing") return prepared def _run(self): """ This method has to be rewritten by child enrichment classes and should contain all the child class logic expected child method outcome: self.result_df = pd.DataFrame({'fan_id': [], 'result_attribute_1_name': [], ....., 'result_attribute_n_name': []}) """ error_string = f"{self.enrichment_name} is missing _run method" raise RuntimeError(error_string) def _finish(self): """ Running all methods to finish enrichment process """ self._generate_fan_attributes() self._populate_collection_attribute_table() self._populate_attribute_table() big_csv_name = self._populate_fan_attribute_table() return big_csv_name def do_everything(self): try: """ Gathering required data, creating required tables or adding data to required tables """ prepared = self._prepare() if not prepared: return """ Running the enrichment """ self._run() """ Wrapping up process and storing the result data""" csv_name = self._finish() return {"csv": csv_name, "collection_id": self.enrichment_collection_id, "schema": self.schema} except MissingInputFieldException as e: """ We don't really use this exception so we'd not pollute alerts monitoring with log errors, and we handle this by returning instead. But leaving this be in case we'll figure out a more elegant solution here. """ self._update_collection_status("failed") return { "error": str(e), } except Exception as e: component_logger.exception(e) self._update_collection_status("failed") return { "error": str(e), } def wait_for_done(self): if self._async_task: try: return self._async_task.wait_for_result() except Exception as e: component_logger.exception(e) return False def _send_runtime_estimation(self): process_status = get_enrichment_progress(self.schema, self.base_collection_id) appsync.send_process_status(self.schema, self.base_collection_id, process_status) def _get_collection_ids(self): """" Getting the base data from collection table (base collection_id + all child collections)""" params = {"collection_id": self.base_collection_id} self.collection_type = get_collection_type(self.schema, self.base_collection_id) if self.collection_type in ["segment"]: set_query = f"SELECT collection_id FROM {self.schema}.set_collection WHERE set_id = %(set_id)s" collection_ids = run_query(set_query, {"set_id": self.base_collection_id}) if len(collection_ids) == 1 and collection_ids[0][0] == -1: # all the sources. sql = f"""SELECT id collection_id, source, collection_type FROM {self.schema}.collection WHERE collection_type IN ('source', 'enrichment');""" else: coll_ids = ",".join(str(cid[0]) for cid in collection_ids) sql = f""" SELECT id collection_id, source, collection_type FROM {self.schema}.collection WHERE (id IN ({coll_ids}) OR parent_id IN ({coll_ids})) AND collection_type IN ('source', 'enrichment');""" else: sql = f""" SELECT id collection_id, source, collection_type FROM {self.schema}.collection WHERE (id = %(collection_id)s OR parent_id = %(collection_id)s) AND collection_type IN ('source', 'enrichment');""" collection_ids_df = pd_read_sql(sql, params=params) if self.collection_type == "source": self.collection_ids["source"] = [self.base_collection_id] condition = collection_ids_df["collection_id"] != self.base_collection_id child_collection_ids_df = collection_ids_df[condition][["collection_id"]] child_collection_ids = child_collection_ids_df.collection_id.unique() self.collection_ids["child"] = child_collection_ids elif self.collection_type == "segment": condition = collection_ids_df["collection_type"] == "source" source_collection_ids_df = collection_ids_df[condition][["collection_id"]] source_collection_ids = source_collection_ids_df.collection_id.unique() self.collection_ids["source"] = source_collection_ids condition = collection_ids_df["collection_type"] != "source" child_collection_ids_df = collection_ids_df[condition][["collection_id"]] child_collection_ids = child_collection_ids_df.collection_id.unique() self.collection_ids["child"] = child_collection_ids else: error_string = ( f'{self.enrichment_name} _get_collection_ids() supports only "source" and "segment" ' f"collection type" ) raise RuntimeError(error_string) def _check_collection_type(self): """ Simple evaluation of the collection_type """ #!TODO need to add better evaluation if self.collection_type not in ["source", "set", "segment"]: self.run_generate_segments = False raise RuntimeError("Cannot enrich this type of collection.") def _get_available_data(self): """ Getting all the data we have for the enrichment to run""" self.available_data["source"] = self._generate_available_data(self.source_attribute_ids) def _generate_available_data(self, attributes_list): """ Generating the query based on attribute_ids and returning dataframe""" data_gathering_condition = self._get_data_gathering_condition_sql(" AND ") attributes_list = [str(x) for x in attributes_list] dynamic_sql = [] for attribute in attributes_list: query_str = f', max(CASE WHEN fa.attribute_id = {attribute} THEN value END) "{attribute}"' dynamic_sql.append(query_str) params = {"collection_id": self.base_collection_id} """Temporary solution for Enrich Gender""" if self.get_previously_enriched_data: row_id_string = "" group_by = "" else: row_id_string = ",s.row_id" group_by = ", 2" if self.collection_type == "segment": # We don't have rows in fan_attribute for segments # NB! We need to limit the collection_fan also by our source collections, else we get some pollution data_cond = self._get_data_gathering_condition_sql("fa.") where_condition = f""" JOIN {self.schema}.collection_fan cf ON cf.fan_id = fa.fan_id AND cf.collection_id = %(collection_id)s {f"WHERE {data_cond}" if data_cond else ""} """ else: where_condition = f"""WHERE fa.collection_id = %(collection_id)s """ sql = f""" SELECT s.fan_id {row_id_string} {''.join(dynamic_sql)} FROM (SELECT DISTINCT fa.fan_id, fa.row_id FROM {self.schema}.fan_attribute fa {where_condition} ) s LEFT JOIN ( SELECT fan_id, row_id, value, attribute_id FROM {self.schema}.fan_attribute fa WHERE attribute_id IN ({', '.join(attributes_list)}) {data_gathering_condition} ) fa ON s.fan_id = fa.fan_id and s.row_id = fa.row_id GROUP BY 1 {group_by}; """ # !TODO - add type casting (try_cast functions) based on commons.system_label table data data = pd_read_sql(sql, params=params) types_list = [{"fan_id": "int64"}, {"row_id": "int32"}, {"collection_id": "int32"}] for type_dict in types_list: try: data = data.astype(type_dict) except: pass return data def _get_data_gathering_condition_sql(self, prefix: str): if self.data_gathering_rule == "strict": collection_ids = ", ".join( [str(x) for x in chain(self.collection_ids["source"], self.collection_ids["child"])] ) if collection_ids: return prefix + f"collection_id in ({collection_ids})" else: return "" else: return "" def _get_available_system_fields(self): """ Get the list of fields available for the collection""" # If "data_gathering_rule" is loose for an enrichment, all attributes are retrieved. # TODO! data_gathering_rule is never loose, remove that handling bit. data_gathering_condition = self._get_data_gathering_condition_sql(" WHERE ") sql = f""" SELECT attribute_id FROM {self.schema}.collection_attribute {data_gathering_condition}; """ self.available_system_fields = pd_read_sql(sql).attribute_id.unique() def _check_available_system_fields(self): """ Basic attributes check - will not run the enrichment if source data doesn't have at least one available self.source_attribute_od """ ok_to_continue_enrichment = True missing_source_attribute_ids = list(set(self.source_attribute_ids) - set(self.available_system_fields)) if len(missing_source_attribute_ids) == len(self.source_attribute_ids): ok_to_continue_enrichment = False info_string = f"{self.schema} - Cannot run {self.enrichment_name} - no source attributes available" component_logger.info(info_string) self._update_collection_status("failed") """ We want to return False because we want to use this to exit enrichment process. We don't want to return an exception because this is handled and ok, and we don't want to spam alerts """ # raise MissingInputFieldException(info_string) return ok_to_continue_enrichment def _get_collection_fan_ids(self): """ Getting the list of fan_attribute_fan_ids for the base_collection_id""" params = {"collection_id": self.base_collection_id} sql = f""" SELECT DISTINCT fan_id FROM {self.schema}.collection_fan WHERE collection_id = %(collection_id)s """ self.collection_fan_ids = pd_read_sql(sql, params=params) def _get_all_system_fields(self): """ Getting all the system fields """ sql = "SELECT * FROM commons.system_label;" self.system_fields = pd_read_sql(sql) def _populate_collection_table(self): """ Generating new collection """ enrichment_collection_id = generate_collection( collection_name=self.enrichment_name, collection_source=self.enrichment_name, schema_name=self.schema, user_id=self.user_id, parent_id=self.base_collection_id, collection_type="enrichment", log_description="initiate_enrichment", status="initiated", ) self.enrichment_collection_id = enrichment_collection_id def _update_collection_status(self, status): """ Updates collection status depending on the stage of enrichment""" if self.enrichment_collection_id: update_collection( schema=self.schema, user_id=self.user_id, collection_id=self.enrichment_collection_id, status=status ) def _populate_collection_fan_table(self): """ Generates the list of collection fans and updates the collection_fan table""" if self.collection_type in ["source", "segment"]: df = self.collection_fan_ids df["collection_id"] = self.enrichment_collection_id df = df[["collection_id", "fan_id"]] string_to_db(df=df, schema=self.schema, table="collection_fan") def _generate_fan_attributes(self): """ Getting a list of attributes that are present in this enrichment""" condition = self.system_fields["a_id"].isin(self.result_attribute_ids) attributes_df = self.system_fields[condition][["a_system_name", "a_id"]] attributes_df.rename(columns={"a_system_name": "name", "a_id": "attribute_id"}, inplace=True) """Special handling for Enrich Gender As we don't have row_id in source data""" # !TODO needs proper handling in the future if self.get_previously_enriched_data and self.collection_type == "source": sql = f""" SELECT DISTINCT fan_id, row_id FROM {self.schema}.fan_attribute WHERE collection_id = {self.base_collection_id} """ original_structure = pd_read_sql(sql) new_result_df = original_structure[["fan_id", "row_id"]].merge(self.result_df, how="left", on="fan_id") else: # We make sure we retain the original row_id new_result_df = self.available_data["source"][["fan_id", "row_id"]].merge( self.result_df, how="left", on="fan_id" ) # If we don't set index, the df default index will be used (and thus wrong row_ids inserted) new_result_df.set_index("row_id", inplace=True) fan_attributes_df = get_fan_attribute_values( source_df=new_result_df, attributes_df=attributes_df, collection_id=self.enrichment_collection_id ) self.fan_attribute = fan_attributes_df def _populate_fan_attribute_table(self): """ Inserts the new fan attribute data if available """ if len(self.fan_attribute) > 0: start = time.time() """ Making sure that after all the iterations DataFrame contains the expected order""" columns = get_table_structure(schema=self.schema, table_name="fan_attribute") fan_attributes_df = self.fan_attribute[columns] """Dropping any row where at least one value is nan because string_to_db can't handle nan As fan_attributes_df has a row per each cell value, we're essentially removing only nan cells """ fan_attributes_df.dropna(how="any", inplace=True) columns_to_int = ["fan_id", "attribute_id", "row_id", "collection_id"] fan_attributes_df = self._convert_to_int64(columns_to_int, fan_attributes_df) return write_csv_to_efs( df=fan_attributes_df, schema=self.schema, table="fan_attribute", attributes=[self.enrichment_name], ) @staticmethod def _convert_to_int64(columns, df): """ Cleaning non convertable values, replacing them with np.nan""" df[columns] = df[columns].apply(pd.to_numeric, axis=0, errors="coerce") type_dict = {} for column in columns: type_dict[column] = "int64" df[columns] = df[columns].astype(type_dict, errors="ignore") return df def _populate_collection_attribute_table(self): """ Adding collection_id - attirbute_id pairs only for new created enrichment_collections""" dynamic_sql = [] for attribute in self.result_attribute_ids: dynamic_sql.append(f"({self.enrichment_collection_id}, {attribute})") if len(dynamic_sql) > 0: sql = f""" INSERT INTO {self.schema}.collection_attribute (collection_id, attribute_id) VALUES {', '.join(dynamic_sql)}""" run_query(sql) def _populate_attribute_table(self): """ INSERTS the missing attributes to attribute_table """ values_list = [] for attribute_id in self.result_attribute_ids: condition = self.system_fields["a_id"] == attribute_id attribute_name = self.system_fields[condition]["a_system_name"].values[0] values_list.append(f"({attribute_id}, '{attribute_name}', 'system_field')") sql = f""" WITH attributes AS ( SELECT * FROM (VALUES {', '.join(values_list)}) AS a(id, name, filter_type) ) INSERT INTO {self.schema}.attribute (id, name, filter_type) SELECT * FROM attributes a WHERE a.id NOT IN (SELECT DISTINCT id FROM {self.schema}.attribute); """ run_query(sql) def _deduplicate_result_data(self): """ Removing duplicate entries: Picks random value from the list (several rows with same fan_id and different names, several genders, etc) """ if len(self.result_attribute_ids) == 1: source_df = self.result_df.groupby(by=["fan_id"]).agg(np.random.choice) source_df.reset_index(inplace=True) self.result_df = source_df else: component_logger.info(f"{self.enrichment_name} can't deduplicate data as more than 1 attribute_id given")