Commit 6bc63e7a authored by Markus Scheidgen's avatar Markus Scheidgen
Browse files

Added code filter uption to admin uploads CLI command.

parent f0faf255
Pipeline #53695 passed with stages
in 19 minutes and 35 seconds
......@@ -18,6 +18,7 @@ from tabulate import tabulate
from mongoengine import Q
from pymongo import UpdateOne
import threading
import elasticsearch_dsl as es
from nomad import processing as proc, config, infrastructure, utils, search, files, coe_repo
from .admin import admin
......@@ -28,8 +29,9 @@ from .admin import admin
@click.option('--staging', help='Select only uploads in staging', is_flag=True)
@click.option('--processing', help='Select only processing uploads', is_flag=True)
@click.option('--outdated', help='Select published uploads with older nomad version', is_flag=True)
@click.option('--code', multiple=True, type=str, help='Select only uploads with calcs of given codes')
@click.pass_context
def uploads(ctx, user: str, staging: bool, processing: bool, outdated: bool):
def uploads(ctx, user: str, staging: bool, processing: bool, outdated: bool, code: List[str]):
infrastructure.setup_mongo()
infrastructure.setup_elastic()
......@@ -47,6 +49,20 @@ def uploads(ctx, user: str, staging: bool, processing: bool, outdated: bool):
{'metadata.nomad_version': {'$ne': config.version}})
query &= Q(upload_id__in=uploads)
if code is not None and len(code) > 0:
code_queries = [es.Q('match', code_name=code_name) for code_name in code]
code_query = es.Q('bool', should=code_queries, minimum_should_match=1)
code_search = es.Search(index=config.elastic.index_name)
code_search = code_search.query(code_query)
code_search.aggs.bucket('uploads', es.A(
'terms', field='upload_id', size=10000, min_doc_count=1))
uploads = [
upload['key']
for upload in code_search.execute().aggs['uploads']['buckets']]
query &= Q(upload_id__in=uploads)
ctx.obj.query = query
ctx.obj.uploads = proc.Upload.objects(query)
......
......@@ -61,6 +61,21 @@ class TestAdmin:
@pytest.mark.usefixtures('reset_config', 'no_warn')
class TestAdminUploads:
@pytest.mark.parametrize('codes, count', [
(['VASP'], 1),
(['doesNotExist'], 0),
(['VASP', 'doesNotExist'], 1)])
def test_uploads_code(self, published, codes, count):
codes_args = []
for code in codes:
codes_args.append('--code')
codes_args.append(code)
result = click.testing.CliRunner().invoke(
cli, ['admin', 'uploads'] + codes_args + ['ls'], catch_exceptions=False, obj=utils.POPO())
assert result.exit_code == 0
assert '%d uploads selected' % count in result.stdout
def test_ls(self, published):
upload_id = published.upload_id
......
Markdown is supported
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