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

Added archive storage of json to files.py.

parent 116d58d5
......@@ -38,7 +38,7 @@ Uploads
.. autoclass:: Upload
"""
from typing import Callable, List, Any, Generator, IO
from typing import Callable, List, Any, Generator, IO, TextIO
import sys
import os
from os.path import join
......@@ -50,6 +50,9 @@ import logging
import itertools
import hashlib
import base64
from contextlib import contextmanager
import gzip
import io
import nomad.config as config
......@@ -64,12 +67,15 @@ if _client is None and 'sphinx' not in sys.modules:
secure=False)
# ensure all neccessary buckets exist
try:
_client.make_bucket(bucket_name=config.s3.uploads_bucket)
logger.info("Created uploads bucket with name %s." % config.s3.uploads_bucket)
except minio.error.BucketAlreadyOwnedByYou:
logger.debug(
"Uploads bucket with name %s already existed." % config.s3.uploads_bucket)
def ensure_bucket(name):
try:
_client.make_bucket(bucket_name=name)
logger.info("Created uploads bucket with name %s." % name)
except minio.error.BucketAlreadyOwnedByYou:
pass
ensure_bucket(config.s3.uploads_bucket)
ensure_bucket(config.s3.archive_bucket)
def get_presigned_upload_url(upload_id: str) -> str:
......@@ -259,3 +265,32 @@ class Upload():
def get_path(self, filename: str) -> str:
""" Returns the tmp directory relative version of a filename. """
return join(self.upload_extract_dir, filename)
@contextmanager
def write_archive_json(archive_id) -> Generator[IO, None, None]:
""" Context manager that yiels a file-like to write the archive json. """
binary_out = io.BytesIO()
gzip_wrapper = gzip.open(binary_out, 'wt')
try:
yield gzip_wrapper
finally:
gzip_wrapper.flush()
binary_out.seek(0)
length = len(binary_out.getvalue())
_client.put_object(
config.s3.archive_bucket, archive_id, binary_out, length=length,
content_type='application/json',
metadata={'Content-Encoding': 'gzip'})
gzip_wrapper.close()
binary_out.close()
def open_archive_json(archive_id) -> IO:
""" Returns a file-like to read the archive json. """
# The result already is a file-like and due to the Content-Encoding metadata is
# will automatically be un-gzipped.
return _client.get_object(config.s3.archive_bucket, archive_id)
# Copyright 2018 Markus Scheidgen
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an"AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
from minio import ResponseError
from threading import Thread
import subprocess
import shlex
import time
from typing import Generator
import json
import nomad.files as files
import nomad.config as config
example_file = './data/examples_vasp.zip'
example_upload_id = '__test_upload_id'
@pytest.fixture
def uploaded_id() -> str:
def uploaded_id() -> Generator[str, None, None]:
example_upload_id = '__test_upload_id'
files._client.fput_object(config.s3.uploads_bucket, example_upload_id, example_file)
return example_upload_id
yield example_upload_id
try:
files._client.remove_object(config.s3.uploads_bucket, example_upload_id)
except ResponseError:
pass
@pytest.fixture(autouse=True)
def tear_down():
yield
@pytest.fixture
def upload_id() -> Generator[str, None, None]:
example_upload_id = '__test_upload_id'
yield example_upload_id
try:
files._client.remove_object(config.s3.uploads_bucket, example_upload_id)
except ResponseError:
pass
def test_presigned_url():
url = files.get_presigned_upload_url(example_upload_id)
@pytest.fixture
def archive_id() -> Generator[str, None, None]:
example_archive_id = '__test_archive_id'
yield example_archive_id
# try:
# files._client.remove_object(config.s3.archive_bucket, example_archive_id)
# except ResponseError:
# pass
def test_presigned_url(upload_id):
url = files.get_presigned_upload_url(upload_id)
assert url is not None
assert isinstance(url, str)
upload_url = files.get_presigned_upload_url(example_upload_id)
upload_url = files.get_presigned_upload_url(upload_id)
cmd = files.create_curl_upload_cmd(upload_url).replace('<ZIPFILE>', example_file)
subprocess.call(shlex.split(cmd))
stat = files._client.stat_object(config.s3.uploads_bucket, example_upload_id)
stat = files._client.stat_object(config.s3.uploads_bucket, upload_id)
assert stat.content_type.startswith('application/octet-steam')
......@@ -52,20 +83,20 @@ def test_upload(uploaded_id: str):
@pytest.mark.timeout(10)
def test_upload_notification():
def test_upload_notification(upload_id):
@files.upload_put_handler
def handle_upload_put(upload_id: str):
assert upload_id == example_upload_id
def handle_upload_put(received_upload_id: str):
assert upload_id == received_upload_id
raise StopIteration
def handle_uploads():
handle_upload_put(upload_id='provided by decorator')
handle_upload_put(received_upload_id='provided by decorator')
handle_uploads_thread = Thread(target=handle_uploads)
handle_uploads_thread.start()
time.sleep(1)
test_presigned_url()
test_presigned_url(upload_id)
handle_uploads_thread.join()
......@@ -82,3 +113,13 @@ def test_hash(uploaded_id: str):
assert isinstance(hash, str)
print(hash)
def test_archive(archive_id: str):
with files.write_archive_json(archive_id) as out:
json.dump({'test': 'value'}, out)
result = json.load(files.open_archive_json(archive_id))
assert 'test' in result
assert result['test'] == 'value'
# Copyright 2018 Markus Scheidgen
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an"AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from nomad.parsing import JSONStreamWriter, parser_dict
from io import StringIO
import json
......
# Copyright 2018 Markus Scheidgen
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an"AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Generator
import pytest
from minio import ResponseError
......
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