Commit effb7ba8 authored by Markus Scheidgen's avatar Markus Scheidgen
Browse files

Remove tmp files after upload.

parent d7f25899
......@@ -95,12 +95,16 @@ def pid_prefix(prefix: int):
@click.argument('paths', nargs=-1)
@click.option('--create-packages', help='Allow migration to create package entries on the fly.', is_flag=True)
@click.option('--local', help='Create local upload files.', is_flag=True)
@click.option('--delete-local', help='Delete created local upload files after upload.', is_flag=True)
@click.option('--parallel', default=1, type=int, help='Use the given amount of parallel processes. Default is 1.')
@click.option('--migration-version', default=0, type=int, help='The version number, only packages with lower or no number will be migrated.')
def upload(paths: list, create_packages, local: bool, parallel: int, migration_version: int):
def upload(
paths: list, create_packages, local: bool, delete_local: bool, parallel: int,
migration_version: int):
infrastructure.setup_logging()
infrastructure.setup_mongo()
migration = NomadCOEMigration(migration_version=migration_version, threads=parallel)
migration.migrate(*paths, local=local, create_packages=create_packages)
migration.migrate(
*paths, local=local, delete_local=delete_local, create_packages=create_packages)
......@@ -124,8 +124,11 @@ class Package(Document):
return iterable_to_stream(zip_file) # type: ignore
def create_package_upload_file(self) -> str:
""" Creates a zip file for the package in tmp and returns its path. """
def create_package_upload_file(self) -> Tuple[str, bool]:
"""
Creates a zip file for the package in tmp and returns its path and whether it
was created (``True``) or existed before (``False``).
"""
upload_filepath = os.path.join(config.fs.nomad_tmp, '%s.zip' % self.package_id)
if not os.path.exists(os.path.dirname(upload_filepath)):
os.mkdir(os.path.dirname(upload_filepath))
......@@ -136,8 +139,11 @@ class Package(Document):
for filename in self.filenames:
filepath = os.path.join(self.upload_path, filename)
zip_file.write(filepath, filename)
created = True
else:
created = False
return upload_filepath
return upload_filepath, created
@classmethod
def index(cls, *upload_paths):
......@@ -544,7 +550,8 @@ class NomadCOEMigration:
self.client.admin.exec_pidprefix_command(payload=dict(prefix=prefix)).response()
def migrate(
self, *args, create_packages: bool = True, local: bool = False) -> utils.POPO:
self, *args, create_packages: bool = True, local: bool = False,
delete_local: bool = False) -> utils.POPO:
"""
Migrate the given uploads.
......@@ -564,7 +571,9 @@ class NomadCOEMigration:
upload_path: A filepath to the upload directory.
create_packages: If True, will create non existing packages.
Will skip with errors otherwise.
local: Instead of streaming an upload, create a local file and use local_path on the upload.
local: Instead of streaming an upload, create a local file and use
local_path on the upload.
delete_local: Delete created local file upload files
Returns: Dictionary with statistics on the migration.
"""
......@@ -577,7 +586,7 @@ class NomadCOEMigration:
def migrate_package(package: Package, of_packages: int):
logger = self.logger.bind(package_id=package.package_id, source_upload_id=package.upload_id)
try:
package_report = self.migrate_package(package, local=local)
package_report = self.migrate_package(package, local=local, delete_local=delete_local)
except Exception as e:
logger.error(
'unexpected exception while migrating packages', exc_info=e)
......@@ -676,7 +685,11 @@ class NomadCOEMigration:
finally:
NomadCOEMigration._client_lock.release()
def migrate_package(self, package: Package, local: bool = False):
def migrate_package(
self, package: Package, local: bool = False,
delete_local: bool = False) -> 'Report':
""" Migrates the given package. For other params see :func:`migrate`. """
source_upload_id = package.upload_id
package_id = package.package_id
......@@ -688,10 +701,13 @@ class NomadCOEMigration:
# upload and process the upload file
from nomad.client import stream_upload_with_client
created_tmp_package_upload_file = None
with utils.timer(logger, 'upload completed'):
try:
if local:
upload_filepath = package.create_package_upload_file()
upload_filepath, created = package.create_package_upload_file()
if created:
created_tmp_package_upload_file = upload_filepath
self.logger.debug('created package upload file')
upload = self.nomad(
'uploads.upload', name=package_id, local_path=upload_filepath)
......@@ -855,6 +871,12 @@ class NomadCOEMigration:
logger.info('migrated package', **report)
if created_tmp_package_upload_file is not None and delete_local:
try:
os.remove(created_tmp_package_upload_file)
except Exception as e:
logger.error('could not remove tmp package upload file', exc_info=e)
return report
def _to_api_metadata(self, calc_with_metadata: CalcWithMetadata) -> dict:
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment