def handle(file_path: str, mappings: dict, entity_type, link_type):
"""Generic .nt file dump handler for the dump import into the DB.
Assumptions: each entity must be represented in a compact block of ''adjacent'' lines
"""
db_manager = DBManager(
get_path('soweego.importer.resources', 'db_credentials.json'))
db_manager.drop(link_type)
db_manager.drop(entity_type)
session = db_manager.new_session()
if not os.path.isfile(file_path):
LOGGER.warning("file: %s not found", file_path)
with open(file_path) as file:
current_key = None
current_entity = entity_type()
for row in file.readlines():
try:
# split triples
row_chunks = []
split_regex = r'(?<=["\<])[^<>"]*(?=["\>])'
for match in re.finditer(split_regex, row, re.I):
row_chunks.append(match.group(0))
if len(row_chunks) != 3:
raise Exception(loc.MALFORMED_ROW)
if row_chunks[1] in mappings:
current_field = mappings[row_chunks[1]]
else:
raise Exception(loc.FIELD_NOT_MAPPED % row_chunks[1])
if current_field == 'url':
link = link_type()
link.url = row_chunks[2]
link.tokens = row_chunks[2].replace("/", " ") # TODO
link.catalog_id = current_entity.catalog_id
session.add(link)
if current_key is not None and current_key != row_chunks[0]:
session.add(current_entity)
current_entity = entity_type()
current_key = row_chunks[0]
current_value = row_chunks[2]
setattr(current_entity, current_field, current_value)
except Exception as e:
LOGGER.warning('Error at row %s \n %s', row, str(e))
session.add(current_entity)
try:
db_manager.create(entity_type)
db_manager.create(link_type)
session.commit()
except Exception as e:
LOGGER.warning(loc.WRONG_MAPPINGS, str(e))
def extract_and_populate(self, dump_file_path):
dump_path = os.path.join(os.path.dirname(
os.path.abspath(dump_file_path)), "%s_%s" % (os.path.basename(dump_file_path), 'extracted'))
if not os.path.isdir(dump_path):
with tarfile.open(dump_file_path, "r:bz2") as tar:
tar.extractall(dump_path)
tables = [MusicbrainzArtistEntity,
MusicbrainzBandEntity]
db_manager = DBManager()
db_manager.drop(tables)
db_manager.create(tables)
artist_count = 0
for artist in self._artist_generator(dump_path):
artist_count = artist_count + 1
session = db_manager.new_session()
session.add(artist)
session.commit()
LOGGER.debug("Added %s artist records" % artist_count)
db_manager.drop([MusicbrainzArtistLinkEntity,
MusicbrainzBandLinkEntity])
db_manager.create([MusicbrainzArtistLinkEntity,
MusicbrainzBandLinkEntity])
link_count = 0
for link in self._link_generator(dump_path):
link_count = link_count + 1
session = db_manager.new_session()
session.add(link)
session.commit()
LOGGER.debug("Added %s link records" % link_count)
isni_link_count = 0
for link in self._isni_link_generator(dump_path):
isni_link_count = isni_link_count + 1
session = db_manager.new_session()
session.add(link)
session.commit()
LOGGER.debug("Added %s ISNI link records" % isni_link_count)
db_manager.drop([MusicBrainzArtistBandRelationship])
db_manager.create([MusicBrainzArtistBandRelationship])
relationships_count = 0
relationships_total = 0
for relationship in self._artist_band_relationship_generator(dump_path):
try:
relationships_total = relationships_total + 1
session = db_manager.new_session()
session.add(relationship)
session.commit()
relationships_count = relationships_count + 1
except IntegrityError as i:
LOGGER.warning(str(i))
LOGGER.debug("Added %s/%s relationships records" %
(relationships_count, relationships_total))
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 _loop_through_entities(
self, file_path: str) -> Generator[Tuple[Dict, List], None, None]:
"""
Generator that given an IMDb dump file (which
should be ".tsv.gz" format) it loops through every
entry and yields it.
:return: a generator which yields a Tuple[entity_info, entity_array]
the consumer of this generator will take `entity_info`, create an
SQLAlchemy entity, and append this to the `entity_array`
"""
db_manager = DBManager()
with gzip.open(file_path, 'rt') as ddump:
session = db_manager.new_session()
# count number of rows for TQDM, so we can display how
# much is missing to complete the process. Then go back
# to the start of the file with `.seek(0)`
n_rows = sum(1 for line in ddump)
ddump.seek(0)
entity_array = []
LOGGER.debug('Dump "%s" has %d entries', file_path, n_rows)
reader = csv.DictReader(ddump, delimiter='\t')
# for every entry in the file..
for entity_info in tqdm(reader, total=n_rows):
# clean the entry
self._normalize_null(entity_info)
# yield the cleaned dict
yield entity_info, entity_array
# every `_sqlalchemy_commit_every` loops we commit the
# session to the DB. This is more efficient than commiting
# every loop, and is not so hard on the memory requirements
# as would be adding everything to session and commiting once
# the for loop is done
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.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.datetime.now() - insert_start_time,
len(entity_array),
)
# commit remaining entities
session.bulk_save_objects(entity_array)
session.commit()
# clear list reference since it might still be available in
# the scope where this generator was used.
entity_array.clear()