import snowflake.connector import config import traceback from datetime import datetime from typing import List, Dict, Optional from utils.list_utils import chunks from utils.reporting.worker_logger import WorkerLogger from utils import list_utils from utils.snowflake.models import ( GRASProjectsQueryResult, GRASRepOwnersQueryResult, GRASProjectModel, GRASArtist, GRASMultiRepOwnersProjectsQueryResult, GRASProjectRepertoireOwner, GRASProductsQueryResult, GRASProductModel, GRASTrack, ) from utils.snowflake.constants import LABEL_EXTERNAL_IDS from workers.gras_project_worker import GRASProjectWorker from utils.snowflake.queries import Query snowflake_connector = snowflake.connector.connect( user=config.SNOWFLAKE_USER, password=config.SNOWFLAKE_PASS, account=config.SNOWFLAKE_ACCOUNT, database=config.PROJECTS_SNOWFLAKE_DB, schema=config.PROJECTS_SNOWFLAKE_SCHEMA, warehouse=config.SNOWFLAKE_WAREHOUSE, ) class GRASProjectsImporter: def fetch(self) -> List[GRASProjectsQueryResult]: cursor = snowflake_connector.cursor() data = cursor.execute(Query.fetch_gras_projects_query()).fetchall() return [GRASProjectsQueryResult._make(result) for result in data] def fetch_projects_with_multiple_rep_owners(self) -> List[GRASMultiRepOwnersProjectsQueryResult]: cursor = snowflake_connector.cursor() data = cursor.execute(Query.fetch_rep_owners_for_various_tracks_projects()).fetchall() return [GRASMultiRepOwnersProjectsQueryResult._make(result) for result in data] class GRASRepOwnersImporter: def fetch_popular_rep_owners(self) -> List[GRASRepOwnersQueryResult]: cursor = snowflake_connector.cursor() data = cursor.execute(Query.fetch_rep_owners()).fetchall() return [GRASRepOwnersQueryResult._make(result) for result in data] class GRASDataImporter: MAX_PRIMARY_ARTISTS = 3 MAX_FEATURED_ARTISTS = 20 PRIMARY_COEFFICIENT = 0.5 rep_owner_mapping = {} MULTIPLE_REP_OWNER_ID = "9999" projects_importer: GRASProjectsImporter logger: WorkerLogger = WorkerLogger() products_per_project: Dict[str, List[GRASProductModel]] = {} def import_projects(self): self.logger.info("Import started", f"GRAS projects import started at {datetime.now()}") try: self.__import_data() except Exception as error: self.logger.error("Import error", f"{error}\n{traceback.format_exc()}") finally: snowflake_connector.close() def __import_data(self): self.projects_importer = GRASProjectsImporter() self.__import_projects() def fetch_projects_with_multiple_rep_owners(self) -> List[GRASProjectRepertoireOwner]: importer = GRASProjectsImporter() data = importer.fetch_projects_with_multiple_rep_owners() return [GRASProjectRepertoireOwner(v.project_id, v.rep_owner_id, v.rep_owner_popularity) for v in data] def __import_projects(self): results = self.projects_importer.fetch() if not results: self.logger.error("Import error", "No projects returned from GRAS") projects_dict: Dict[int, GRASProjectModel] = {} for item in results: project = projects_dict.get(item.project_id, self.__gras_project_from_item(item)) artists = self.__artists_from_item(item) self.__add_artists(artists, project) projects_dict[item.project_id] = project if projects_dict: self.__process_projects(projects_dict) def __product_from_item(self, item: GRASProductsQueryResult): return GRASProductModel( id=item.product_id, title=item.product_title, project_id=item.project_id, type=item.product_type ) def __track_from_item(self, item: GRASProductsQueryResult) -> GRASTrack: date = datetime.strptime(str(item.track_release_date), "%Y%m%d") if item.track_release_date else None isrc = item.track_isrc track = GRASTrack(isrc, item.track_name, item.track_name_suppl, item.track_type, date) return track def __artist_from_item(self, item: GRASProductsQueryResult) -> GRASArtist: if item.member_id: return GRASArtist(item.member_id, item.member_name) return GRASArtist(item.artist_id, item.artist_name) def __add_artists(self, artists: List[GRASArtist], project: GRASProjectModel): for artist in artists: artist_in_list = list_utils.first_or_none(project.artists, lambda x: x.id == artist.id) if not artist_in_list: project.artists.append(artist) artist_in_list = artist else: artist_in_list.popularity += artist.popularity if artist.is_main: artist_in_list.is_main = True if artist.is_primary: artist_in_list.is_primary = True def __artists_from_item(self, item: GRASProjectsQueryResult) -> List[GRASArtist]: artists = [] if item.main_member_id: is_primary = item.main_member_type == "Primary" artists.append( GRASArtist(item.main_member_id, item.main_member_name, True, is_primary, item.main_member_popularity) ) elif item.main_artist_id and item.main_artist_type == "Individual": artists.append( GRASArtist(item.main_artist_id, item.main_artist_name, True, True, item.main_artist_popularity) ) if item.member_id: artists.append(GRASArtist(item.member_id, item.member_name, False, False, item.member_popularity)) elif item.artist_id and item.artist_type == "Individual": artists.append(GRASArtist(item.artist_id, item.artist_name, False, False, item.artist_popularity)) return artists def __gras_project_from_item(self, item: GRASProjectsQueryResult) -> GRASProjectModel: return GRASProjectModel( item.project_id, item.project_title, item.rep_owner_key, item.project_start_date, item.is_deleted ) def __process_projects(self, projects_dict: Dict[int, GRASProjectModel]): rep_owners_metadata = self.fetch_projects_with_multiple_rep_owners() projects = map(lambda x: self.__map_project(x, rep_owners_metadata), projects_dict.values()) projects = list(filter(lambda x: x is not None, projects)) for chunk in chunks(projects, 500): GRASProjectWorker().perform_async(chunk) if len(projects_dict.keys()) > 0: self.logger.success( "Import finished", f"GRAS projects import finished at {datetime.now()} count: {len(projects_dict.keys())}", ) else: self.logger.error("Import error", "No projects imported") def __map_project(self, project, rep_owners_metadata) -> Optional[GRASProjectModel]: self.__fix_rep_owner_if_needed(project, rep_owners_metadata) if project.rep_owner_id == self.MULTIPLE_REP_OWNER_ID: return None project.primary_artists = [] project.featured_artists = [] sorted_artists = project.sorted_artists() if sorted_artists: base_popularity = sorted_artists[0].popularity primary = list(filter(lambda x: self.__is_primary(x, base_popularity), sorted_artists))[:3] featured_artists_candidates = list(filter(lambda x: x not in primary, sorted_artists)) featured = featured_artists_candidates[:20] project.featured_artists = featured if featured else [] project.primary_artists = primary if primary else [] project.products = self.products_per_project.get(project.id, []) return project def __fix_rep_owner_if_needed( self, project: GRASProjectModel, rep_owners_metadata: List[GRASProjectRepertoireOwner] ): metadata_items = list(filter(lambda x: x.project_id == project.id, rep_owners_metadata)) if len(metadata_items) > 0: most_popular = metadata_items[:2] has_major_rep_owner = len(most_popular) > 1 and most_popular[0].popularity > most_popular[1].popularity if len(most_popular) == 1 or has_major_rep_owner: if most_popular[0].rep_owner_id not in list(LABEL_EXTERNAL_IDS): project.rep_owner_id = None else: project.rep_owner_id = most_popular[0].rep_owner_id def __is_primary(self, artist: GRASArtist, base_popularity: int): return artist.is_main and artist.is_primary and artist.popularity > self.PRIMARY_COEFFICIENT * base_popularity