Commit ebf8da30 authored by Kenzo-Hugo Hillion's avatar Kenzo-Hugo Hillion
Browse files

refactor and add limit of 100000 fasta sequences

parent 387346c5
Pipeline #32468 passed with stages
in 3 minutes and 14 seconds
import hashlib
from io import StringIO
from django.http import HttpResponse
from django.core.cache import cache
from django.http import HttpResponse, JsonResponse
from drf_yasg.utils import swagger_auto_schema
from marshmallow.exceptions import ValidationError
from rest_framework.response import Response
......@@ -21,6 +23,28 @@ class GeneViewSet(BulkViewSet):
query_params_parser = GeneQueryParams
lookup_field = 'gene_id'
def _get_queryset_count(self, queryset):
hash_object = hashlib.md5(str(queryset.query).encode('utf-8'))
redis_key = hash_object.hexdigest()
if redis_key in cache:
return cache.get(redis_key)
else:
return queryset.count()
def _build_fasta_response(self):
queryset = self.filter_queryset(self.get_queryset())
count = self._get_queryset_count(queryset)
if count >= 100000:
return JsonResponse({'error': 'Too many genes, can obtain only up to 100000 fasta seq.'})
fasta_file = StringIO()
for gene in queryset.iterator():
fasta_file.write(gene.fasta)
# generate the file
response = HttpResponse(fasta_file.getvalue(), content_type='text/fasta')
filename = 'metagenedb_sequences.fasta'
response['Content-Disposition'] = 'attachment; filename=%s' % filename
return response
@swagger_auto_schema(
tags=['Genes'],
)
......@@ -34,15 +58,7 @@ class GeneViewSet(BulkViewSet):
})
return Response(error_message, status=HTTP_422_UNPROCESSABLE_ENTITY)
if query_params.get('fasta', False) is True:
queryset = self.filter_queryset(self.get_queryset())
fasta_file = StringIO()
for gene in queryset.iterator():
fasta_file.write(gene.fasta)
# generate the file
response = HttpResponse(fasta_file.getvalue(), content_type='text/fasta')
filename = 'metagenedb_sequences.fasta'
response['Content-Disposition'] = 'attachment; filename=%s' % filename
return response
return self._build_fasta_response()
return super().list(request, *args, **kwargs)
@swagger_auto_schema(
......
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