def test_get_record_class_by_pid_type(app):
"""Test get record class by PID type."""
record = SonarRecord.get_record_class_by_pid_type('doc')
assert record.__name__ == 'DocumentRecord'
def export(pid_type, serializer_key, output_dir):
"""Export records for the given record type.
:param pid_type: record type
:param output_dir: Output directory
"""
click.secho('Export "{pid_type}" records in {dir}'.format(
pid_type=pid_type, dir=output_dir.name))
try:
# Get the correct record class
record_class = SonarRecord.get_record_class_by_pid_type(pid_type)
if not record_class:
raise Exception('No record class found for type "{type}"'.format(
type=pid_type))
# Load the serializer
serializer_class = current_app.config.get(
'SONAR_APP_EXPORT_SERIALIZERS', {}).get(pid_type)
if serializer_class:
serializer = obj_or_import_string(serializer_class)()
else:
serializer = None
pids = record_class.get_all_pids()
records = []
# Create ouptut directory if not exists
if pids:
pathlib.Path(output_dir.name).mkdir(mode=0o755,
parents=True,
exist_ok=True)
for pid in pids:
record = record_class.get_record_by_pid(pid)
if serializer:
record = serializer.dump(record)
else:
record = record.dumps()
for file in record.get('files', []):
if file.get('uri'):
target_path = join(output_dir.name, pid, file['key'])
pathlib.Path(dirname(target_path)).mkdir(mode=0o755,
parents=True,
exist_ok=True)
shutil.copyfile(file['uri'], target_path)
file.pop('uri')
file['path'] = './{pid}/{key}'.format(pid=pid,
key=file['key'])
records.append(record)
if records:
# Write data
output_file = join(output_dir.name, 'data.json')
f = open(output_file, 'w')
f.write(json.dumps(records))
f.close()
click.secho('Finished', fg='green')
except Exception as exception:
click.secho('An error occured during export: {error}'.format(
error=str(exception)),
fg='red')