import logging from datetime import datetime from typing import List import numpy as np import pandas as pd from service.db import engine from service.utils.aws_connectors import run_query, s3_to_pandas from service.utils.data_model_utils import ( add_event_log, attach_profile_id, change_collection_status, delete_used_guesses, load_to_rds, update_collection, ) from .data_model import store_file_labels logger = logging.getLogger(__name__) def get_source_data(bucket, file_name) -> pd.DataFrame: """Pulls File from S3 and returns pandas DataFrame.""" # getting the file format file_format = file_name.lower().split(".") index = len(file_format) - 1 file_format = file_format[index] source_df = s3_to_pandas(bucket, file_name, "infer", file_format=file_format) return source_df def find_new_fans( source_unique_fan_df: pd.DataFrame, management_schema ) -> pd.DataFrame: """ Gets additional fans from source that we are missing in our fan table previously """ sql = ( f"INSERT INTO {management_schema}.fan (profile_id) VALUES %s " f"ON CONFLICT (profile_id) DO UPDATE SET profile_id = EXCLUDED.profile_id " f"RETURNING id, profile_id;" ) fans = run_query(sql, [(val,) for val in source_unique_fan_df.profile_id.values]) fans_df = pd.DataFrame(fans, columns=["fan_id", "profile_id"]) return fans_df def get_fans(source_df: pd.DataFrame, fan_identifier_label: str) -> pd.DataFrame: """Returns fan dataframe with unique fan identifier""" fan_df = pd.DataFrame( source_df[fan_identifier_label].unique(), columns=["profile_id"] ) return fan_df def load_new_attributes_to_db( schema: str, source_df: pd.DataFrame, sys_fields: list ) -> pd.DataFrame: """Adding fan_id field to be systemfield""" sys_fields.append("fan_id") all_attributes = pd.DataFrame(source_df.columns).rename(columns={0: "name"}) existing_attributes = run_query( f"SELECT name FROM {schema}.attribute;", return_type="df" ) all_system_attributes = run_query( "SELECT a_id, a_system_name FROM commons.system_label;", return_type="df" ) # Substract source names from what we have already in our db cond = all_attributes["name"].isin(existing_attributes["name"]) new_attributes = all_attributes.drop(all_attributes[cond].index) """ Adding 'system_field' indicator to 'filter_type' column """ new_attributes["filter_type"] = new_attributes["name"].transform( lambda x: "system_field" if x in sys_fields else np.NAN ) """ Splitting attributes in to 2 groups as the new_system_fields will have dedicated attribute_id""" attributes_list = [] new_system_attributes = new_attributes[ new_attributes["filter_type"] == "system_field" ] new_system_attributes = new_system_attributes.merge( all_system_attributes, how="left", left_on="name", right_on="a_system_name" )[["a_id", "name", "filter_type"]] new_system_attributes.rename(columns={"a_id": "id"}, inplace=True) attributes_list.append(new_system_attributes) new_non_system_attributes = new_attributes[new_attributes["filter_type"].isna()] attributes_list.append(new_non_system_attributes) # It might happen that we don't get any new attributes - don't try to insert empty dataframe for df in attributes_list: if not df.empty: """Insert new attributes to attributes table""" df.to_sql( "attribute", engine, schema=schema, if_exists="append", index=False ) """ Since we have new attributes, retrieve attributes table along with new values """ attributes_df = pd.read_sql( f"SELECT id as attribute_id, name FROM {schema}.attribute;", engine ) """ Filter the fields that exist in the current source dataframe to use in collection_attribute table generation """ current_attributes_df = attributes_df[ attributes_df["name"].isin(list(source_df.columns)) ] return attributes_df, current_attributes_df def get_fan_attribute_values( source_df: pd.DataFrame, attributes_df: pd.DataFrame, collection_id: int ) -> pd.DataFrame: """Returns fan_attribute table""" """ We create a row_id """ source_df["row_id"] = source_df.index """ We transpose attributes """ fan_attribute_base = pd.DataFrame( source_df.set_index(["fan_id", "row_id"]).stack(), columns=["value"] ) fan_attribute_base.index.names = ["fan_id", "row_id", "name"] fan_attribute_base = fan_attribute_base.reset_index(level=[0, 1]) """ Left join and map attributes by name to get ids for attributes """ fan_attribute_result = pd.merge( fan_attribute_base, attributes_df, left_on="name", right_on="name", how="left" ) # Limit the columns we want in fan_attribute table fan_attribute_result = fan_attribute_result[ ["fan_id", "attribute_id", "value", "row_id"] ] # Add the field that holds collection_id fan_attribute_result["collection_id"] = collection_id """ Adding timestamp field to successfully stage the file during upload to RDS""" fan_attribute_result["timestamp"] = str(datetime.utcnow()) return fan_attribute_result def generate_collection_fan(fan_df: pd.DataFrame, collection_id: int) -> pd.DataFrame: fan_df["collection_id"] = collection_id return fan_df def generate_collection_attribute( attribute_df: pd.DataFrame, collection_id: int, value=None ) -> pd.DataFrame: collection_attribute_df = attribute_df.copy(deep=True)[["attribute_id"]] collection_attribute_df.columns = ["attribute_id"] collection_attribute_df["collection_id"] = collection_id collection_attribute_df["value"] = value return collection_attribute_df def generate_base_attributes_set( source_df: pd.DataFrame, schema_name: str, management_schema: str, collection_id: int, sys_fields: list, user_id, ) -> None: """Get new fans from source dataframe""" source_unique_fan_df = get_fans(source_df, "profile_id") fan_df = find_new_fans(source_unique_fan_df, management_schema) source_df = source_df.merge(fan_df, how="left", on="profile_id") source_df.drop(columns=["profile_id"], inplace=True) """ Load new attributes we haven't seen yet to attributes table, and return """ attributes_df, current_attributes_df = load_new_attributes_to_db( schema_name, source_df, sys_fields ) """ Generate fan_attribute """ fan_attribute_df = get_fan_attribute_values(source_df, attributes_df, collection_id) """ Generate collection_fan """ collection_fan_df = generate_collection_fan(fan_df[["fan_id"]], collection_id) """ Generate collection_attribute """ without_fan_id = current_attributes_df["name"] != "fan_id" collection_attribute = generate_collection_attribute( current_attributes_df[without_fan_id], collection_id ) """ Put everything together and load to RDS """ table_dict = { "fan_attribute": [fan_attribute_df, ["fan_id", "attribute_id", "row_id"]], "collection_fan": [collection_fan_df, ["collection_id", "fan_id"]], "collection_attribute": [collection_attribute, ["collection_id", "fan_id"]], } load_to_rds(schema=schema_name, table_dict=table_dict) add_event_log( schema=schema_name, collection_id=collection_id, user_id=user_id, log_type="test_log", # ['creation','encrihment'] # TODO! this hsould not be test description="generate_base_attributes_set", ) logger.info( f"{schema_name} - completed base attribute generation for collection {collection_id}" ) def generate_labels_and_attributes( schema_name, management_schema, collection_name, field_map, source_df, user_id, collection_id=None, final_status="finished", ) -> dict: """Rename collection""" update_collection( schema_name, user_id, collection_id, name=collection_name, status="processing" ) try: """Store label information about the file""" source_df, new_column_names, sys_fields = store_file_labels( source_df, field_map ) """ Get fan_id by sha256 obfuscating email address of fan """ source_df = attach_profile_id(df=source_df, fan_identifier_label="userEmail") """ Unpack and load data into data model """ generate_base_attributes_set( source_df, schema_name, management_schema, collection_id, sys_fields, user_id, ) # TODO! why do we need this? """ After running the main generate_base_attributes_set source_df gets 2 new fields""" new_column_names.extend(["fan_id", "row_id"]) """ Generate a view with repacked data """ # generate_collection_source_data_view(schema=schema_name, collection_id=collection_id, suffix='source') """ Remove guesses from guessed_file_upload_fields, now that we successfully did all that processing """ delete_used_guesses(schema_name, collection_id) """ Updating the status in collection table """ if final_status != "processing": change_collection_status( schema=schema_name, user_id=user_id, collection_id=collection_id, status=final_status, ) except Exception as e: change_collection_status( schema=schema_name, user_id=user_id, collection_id=collection_id, status="error: processing failed", ) logger.exception(e) raise e return { "source_df": source_df, "new_column_names": new_column_names, "collection_id": collection_id, } def process_file( schema_name, management_schema, file_name, bucket_name, collection_name, field_map, user_id, collection_id=None, final_status="finished", ) -> str: """ Takes input csv file name, requires one column for fan identifier (e.g. customer email, user hash) """ if not bucket_name: bucket_name = "devel---useruploadoriginals" # 'fansifter-model-data' """ Get dataframe from source csv file in s3 bucket """ source_df = get_source_data(bucket_name, file_name) """ Processing the provided file """ response = generate_labels_and_attributes( schema_name=schema_name, management_schema=management_schema, collection_name=collection_name, field_map=field_map, source_df=source_df, user_id=user_id, collection_id=collection_id, final_status=final_status, ) return response["collection_id"] def update_collection_labels( schema_name, collection_ids, field_map, to_system_fields=True ): """Updates field ids according to field_map for given collection_id (and it's children) - attribute - fan_attribute - collection_attribute Used for external imports (Klaviyo) when user input provided after all the data is extracted """ attribute_pairs_to_update = [] # For updating fan_attribute table attributes_to_add = [] # For updating attribute table for attribute_pair in field_map: old_id = int(attribute_pair["old_attribute_id"]) new_id = int(attribute_pair["system_field_id"]) attribute_pairs_to_update.append( {"old_attribute_id": old_id, "new_attribute_id": new_id} ) attributes_to_add.append( {"id": new_id, "name": attribute_pair["system_field_name"]} ) update_fan_attribute_table( schema=schema_name, collection_ids=collection_ids, attribute_pairs=attribute_pairs_to_update, ) update_attribute_table( schema=schema_name, attributes=attributes_to_add, system_fields=to_system_fields ) update_collection_attribute( schema=schema_name, collection_ids=collection_ids, attribute_pairs=attribute_pairs_to_update, ) pass def update_fan_attribute_table( schema: str, collection_ids: List[int], attribute_pairs: List[dict] ): """Changing old_attribute_id to new_attribute_id in attribute_paris for each collection_id in a collection_ids list """ if len(collection_ids) > 0: if len(attribute_pairs) > 0: set_attribute_statement = generate_case_when_statement(attribute_pairs) where_statement = generate_where_statement(collection_ids) sql = f""" UPDATE {schema}.fan_attribute fa {set_attribute_statement} {where_statement} """ run_query(sql) else: raise RuntimeError("No attribute_pairs provided") else: raise RuntimeError("No collection_ids provided") def generate_case_when_statement(attribute_pairs: List[dict]): """Generating CASE WHEN satements based on attribute_id_pairs - attribute_pairs: list of dicts with `old_attribute_id` and `new_attribtue_id` """ set_attribute_id_lst = [] set_attribute_id_end = [] for attribute_pair in attribute_pairs: if ( "old_attribute_id" in attribute_pair and "new_attribute_id" in attribute_pair ): old_a_id = int(str(attribute_pair["old_attribute_id"])) new_a_id = int(str(attribute_pair["new_attribute_id"])) pair_string = f"CASE WHEN fa.attribute_id = {old_a_id} THEN {new_a_id}" set_attribute_id_lst.append(pair_string) set_attribute_id_end.append("END") else: raise RuntimeError("Wrong attribute_pairs format") set_attribute_id_lst.append("fa.attribute_id") set_attribute_statement = f"SET attribute_id = {' ELSE '.join(set_attribute_id_lst)} {' '.join(set_attribute_id_end)}" return set_attribute_statement def generate_where_statement(collection_ids: List[int]): """Generating WHERE statement that filters collection_id""" where_lst = [] for collection_id in collection_ids: col_id = str(int(str(collection_id))) where_lst.append(col_id) where_statement = f'WHERE fa.collection_id IN ({", ".join(where_lst)});' return where_statement def update_attribute_table( schema: str, attributes: List[dict], system_fields: bool = True ): """Adding missing attribute_ids to attribute_table - attributes: Holds list of dictionaries with `id` and `name` pairs""" if system_fields: system_field = "'system_field'" else: system_field = "null" for attribute in attributes: if "id" in attribute and "name" in attribute: a_id = int(str(attribute["id"])) name = attribute["name"] params = {"name": name} sql = f""" INSERT INTO {schema}.attribute (id, name, filter_type) VALUES ({a_id}, %(name)s, {system_field}) ON CONFLICT (id) DO NOTHING; """ run_query(sql, params) return "done" def update_collection_attribute( schema: str, collection_ids: List[int], attribute_pairs: List[dict] ): """Updating collection_attribute table""" if len(collection_ids) > 0: if len(attribute_pairs) > 0: set_attribute_statement = generate_case_when_statement(attribute_pairs) where_statement = generate_where_statement(collection_ids) sql = f""" UPDATE {schema}.collection_attribute fa {set_attribute_statement} {where_statement} """ run_query(sql) else: raise RuntimeError("No attribute_pairs provided") else: raise RuntimeError("No collection_ids provided")