def _process_artists_dump(self, dump_file_path, resolve): LOGGER.info( "Starting import of musicians and bands from Discogs dump '%s'", dump_file_path, ) start = datetime.now() tables = [ DiscogsMusicianEntity, DiscogsMusicianNlpEntity, DiscogsMusicianLinkEntity, DiscogsGroupEntity, DiscogsGroupNlpEntity, DiscogsGroupLinkEntity, ] db_manager = DBManager() LOGGER.info('Connected to database: %s', db_manager.get_engine().url) db_manager.drop(tables) db_manager.create(tables) LOGGER.info( 'SQL tables dropped and re-created: %s', [table.__tablename__ for table in tables], ) extracted_path = '.'.join(dump_file_path.split('.')[:-1]) # Extract dump file if it has not yet been extracted if not os.path.exists(extracted_path): LOGGER.info('Extracting dump file') with gzip.open(dump_file_path, 'rb') as f_in: with open(extracted_path, 'wb') as f_out: shutil.copyfileobj(f_in, f_out) # count number of entries n_rows = sum( 1 for _ in self._g_process_et_items(extracted_path, 'artist')) session = db_manager.new_session() entity_array = [] # array to which we'll add the entities for _, node in tqdm(self._g_process_et_items(extracted_path, 'artist'), total=n_rows): if not node.tag == 'artist': continue infos = self._extract_from_artist_node(node, resolve) if infos is None: continue if 'groups' in infos: entity = DiscogsMusicianEntity() self._populate_musician(entity_array, entity, infos) # Band elif 'members' in infos: entity = DiscogsGroupEntity() self._populate_band(entity_array, entity, infos) # commit in batches of `self._sqlalchemy_commit_every` if len(entity_array) >= self._sqlalchemy_commit_every: LOGGER.info('Adding batch of entities to the database, ' 'this will take a while. ' 'Progress will resume soon.') insert_start_time = datetime.now() session.bulk_save_objects(entity_array) session.commit() session.expunge_all() # clear session entity_array.clear() # clear entity array LOGGER.debug( 'It took %s to add %s entities to the database', datetime.now() - insert_start_time, self._sqlalchemy_commit_every, ) # finally commit remaining entities in session # (if any), and close session session.bulk_save_objects(entity_array) session.commit() session.close() end = datetime.now() LOGGER.info( 'Import completed in %s. ' 'Total entities: %d - %d musicians with %d links - %d bands' ' with %d links - %d discarded dead links.', end - start, self.total_entities, self.musicians, self.musician_links, self.bands, self.band_links, self.dead_links, ) # once the import process is complete, # we can safely delete the extracted discogs dump os.remove(extracted_path)
def _process_masters_dump(self, dump_file_path): LOGGER.info("Starting import of masters from Discogs dump '%s'", dump_file_path) start = datetime.now() tables = [DiscogsMasterEntity, DiscogsMasterArtistRelationship] db_manager = DBManager() LOGGER.info('Connected to database: %s', db_manager.get_engine().url) db_manager.drop(tables) db_manager.create(tables) LOGGER.info( 'SQL tables dropped and re-created: %s', [table.__tablename__ for table in tables], ) extracted_path = '.'.join(dump_file_path.split('.')[:-1]) # Extract dump file if it has not yet been extracted if not os.path.exists(extracted_path): LOGGER.info('Extracting dump file') with gzip.open(dump_file_path, 'rb') as f_in: with open(extracted_path, 'wb') as f_out: shutil.copyfileobj(f_in, f_out) # count number of entries n_rows = sum( 1 for _ in self._g_process_et_items(extracted_path, 'master')) session = db_manager.new_session() entity_array = [] # array to which we'll add the entities relationships_set = set() self.total_entities = 0 for _, node in tqdm(self._g_process_et_items(extracted_path, 'master'), total=n_rows): if not node.tag == 'master': continue self.total_entities += 1 entity = self._extract_from_master_node(node, relationships_set) entity_array.append(entity) # commit in batches of `self._sqlalchemy_commit_every` if len(entity_array) >= self._sqlalchemy_commit_every: LOGGER.info('Adding batch of entities to the database, ' 'this will take a while.' 'Progress will resume soon.') insert_start_time = datetime.now() session.bulk_save_objects(entity_array) session.commit() session.expunge_all() # clear session entity_array.clear() # clear entity array LOGGER.debug( 'It took %s to add %s entities to the database', datetime.now() - insert_start_time, self._sqlalchemy_commit_every, ) # finally commit remaining entities in session # (if any), and close session session.bulk_save_objects(entity_array) session.bulk_save_objects([ DiscogsMasterArtistRelationship(id1, id2) for id1, id2 in relationships_set ]) session.commit() session.close() end = datetime.now() LOGGER.info( 'Import completed in %s. Total entities: %d. ' 'Total relationships %s.', end - start, self.total_entities, len(relationships_set), ) # once the import process is complete, # we can safely delete the extracted discogs dump os.remove(extracted_path)
def extract_and_populate(self, dump_file_paths: List[str], resolve: bool) -> None: """ Extracts the data in the dumps (person and movie) and processes them. It then proceeds to add the appropriate data to the database. See :ref:`soweego.importer.models.imdb_entity` module to see the SQLAlchemy definition of the entities we use to save IMDB data. :param dump_file_paths: the absolute paths of the already downloaded dump files. """ # the order of these files is specified in `self.get_dump_download_urls` person_file_path = dump_file_paths[0] movies_file_path = dump_file_paths[1] LOGGER.debug('Path to movie info dump: %s', movies_file_path) LOGGER.debug('Path to person info dump: %s', person_file_path) start = datetime.datetime.now() tables = [ imdb_entity.ImdbActorEntity, imdb_entity.ImdbDirectorEntity, imdb_entity.ImdbMovieEntity, imdb_entity.ImdbMusicianEntity, imdb_entity.ImdbProducerEntity, imdb_entity.ImdbWriterEntity, imdb_entity.ImdbMoviePersonRelationship, ] db_manager = DBManager() LOGGER.info('Connected to database: %s', db_manager.get_engine().url) db_manager.drop(tables) db_manager.create(tables) LOGGER.info( 'SQL tables dropped and re-created: %s', [table.__tablename__ for table in tables], ) LOGGER.info('Starting import of movies ...') # Here we open the movie dump file, and add everything to the DB for movie_info, entity_array in self._loop_through_entities( movies_file_path): # create the movie SQLAlchemy entity and populate it movie_entity = imdb_entity.ImdbMovieEntity() movie_entity.catalog_id = movie_info.get('tconst') movie_entity.title_type = movie_info.get('titleType') if movie_info.get('primaryTitle') is not None: movie_entity.name = movie_info.get('primaryTitle') movie_entity.name_tokens = ' '.join( text_utils.tokenize(movie_info.get('primaryTitle'))) movie_entity.is_adult = (True if movie_info.get('isAdult') == '1' else False) try: movie_entity.born = datetime.date(year=int( movie_info.get('startYear')), month=1, day=1) movie_entity.born_precision = 9 except (KeyError, TypeError): LOGGER.debug('No start year value for %s', movie_entity) try: movie_entity.died = datetime.date(year=int( movie_info.get('endYear')), month=1, day=1) movie_entity.died_precision = 9 except (KeyError, TypeError): LOGGER.debug('No end year value for %s', movie_entity) movie_entity.runtime_minutes = movie_info.get('runtimeMinutes') if movie_info.get('genres'): # if movie has a genre specified movie_entity.genres = ' '.join( text_utils.tokenize(movie_info.get('genres'))) # Creates entity for alias alias = movie_info.get('originalTitle') if alias is not None and movie_entity.name != alias: alias_entity = copy.deepcopy(movie_entity) alias_entity.name = alias alias_entity.name_tokens = ' '.join(text_utils.tokenize(alias)) entity_array.append(alias_entity) entity_array.append(movie_entity) self.n_movies += 1 # mark end for movie import process end = datetime.datetime.now() LOGGER.info( 'Movie import completed in %s. ' 'Total movies imported: %d', end - start, self.n_movies, ) LOGGER.info('Starting import of people ...') # reset timer for persons import start = datetime.datetime.now() for person_info, entity_array in self._loop_through_entities( person_file_path): # IMDb saves the list of professions as a comma separated # string professions = person_info.get('primaryProfession') # if person has no professions then ignore it if not professions: LOGGER.debug('Person %s has no professions', person_info.get('nconst')) continue professions = professions.split(',') # each person can be added to multiple tables in the DB, # each table stands for one of the main professions types_of_entities = [] if 'actor' in professions or 'actress' in professions: self.n_actors += 1 types_of_entities.append(imdb_entity.ImdbActorEntity()) if 'director' in professions: self.n_directors += 1 types_of_entities.append(imdb_entity.ImdbDirectorEntity()) if 'producer' in professions: self.n_producers += 1 types_of_entities.append(imdb_entity.ImdbProducerEntity()) if any(prof in [ 'sound_department', 'composer', 'music_department', 'soundtrack', ] for prof in professions): self.n_musicians += 1 types_of_entities.append(imdb_entity.ImdbMusicianEntity()) if 'writer' in professions: self.n_writers += 1 types_of_entities.append(imdb_entity.ImdbWriterEntity()) # if the only profession a person has is `miscellaneous` then we # add it to all tables if professions == ['miscellaneous']: self.n_misc += 1 types_of_entities = [ imdb_entity.ImdbActorEntity(), imdb_entity.ImdbDirectorEntity(), imdb_entity.ImdbMusicianEntity(), imdb_entity.ImdbProducerEntity(), imdb_entity.ImdbWriterEntity(), ] # add person to every matching table for etype in types_of_entities: self._populate_person(etype, person_info, entity_array) # if person is known for any movies then add these to the # database as well if person_info.get('knownForTitles'): self.n_person_movie_links += 1 self._populate_person_movie_relations(person_info, entity_array) self.n_persons += 1 # mark the end time for the person import process end = datetime.datetime.now() LOGGER.info( 'Person import completed in %s. ' 'Total people imported: %d - ' 'Actors: %d - Directors: %d - Musicians: %d - ' 'Producers: %d - Writers: %d - Misc: %d', end - start, self.n_persons, self.n_actors, self.n_directors, self.n_musicians, self.n_producers, self.n_writers, self.n_misc, )