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 service.async_task_manager import io_task from service.tasks.base_attributes_generation import get_fan_attribute_values from service.utils.aws_connectors import pd_read_sql, run_query from service.utils.data_model_utils import ( generate_collection, get_collection_type, get_table_structure, string_to_db_async, update_collection, ) logger = logging.getLogger(__name__) 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 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! def __init__( self, collection_id: int, schema: str, user_id: str, default_value: str = None ): 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.source_data = pd.DataFrame({}) self.collection_ids: Dict[str, List[int]] = {"source": [], "child": []} self.enrichment_collection_id = None self.collection_fan_ids: List[int] = [] self.fan_attribute_fan_ids: Dict[str, List[int]] = {"source": []} self.result_df = pd.DataFrame({}) self.fan_attribute = pd.DataFrame({}) self.default_value = default_value self._async_task = None def _prepare(self): """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._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() self._get_fan_attribute_fan_ids() """ Populating main 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_fan_attribute_table() self._populate_collection_attribute_table() self._populate_attribute_table() self._update_collection_status("finished") 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""" self._finish() except MissingInputFieldException: """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") except Exception as e: logger.exception(e) self._update_collection_status("failed") @io_task def _do_everything_async(self, **kwargs): self.do_everything() return True def do_everything_async(self): self._async_task = self._do_everything_async() return self._async_task def wait_for_done(self): if self._async_task: try: return self._async_task.wait_for_result() except Exception as e: logger.exception(e) return False def async_done(self): if self._async_task: return self._async_task.done else: return False def async_failed(self): if self._async_task: return self._async_task.failed else: return False 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 = """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 Exception: 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" 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_fan_attribute_fan_ids(self): """We need to split the fan_attribute_fan_ids source - have the source_attribute_ids Allows enrichment to skipp fan_attribute_fan_ids (if it needs to) and use existing enriched data instead """ source_fan_ids = self.available_data["source"].fan_id.unique() self.fan_attribute_fan_ids = {"source": source_fan_ids} 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_async(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""" 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 ) string_to_db_async( df=fan_attributes_df, schema=self.schema, table="fan_attribute" ) logger.info( f"{self.schema} - It took {time.time()-start} to load enrichment fan_attribute data to DB" ) @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=1, 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: logger.info( f"{self.enrichment_name} can't deduplicate data as more than 1 attribute_id given" )