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')