def decorate_queryset(
cls,
feature_type: FeatureType,
queryset: models.QuerySet,
output_crs: CRS,
**params,
):
"""Apply presentation layer logic to the queryset."""
# Avoid fetching relations, fetch these within the same query,
related = [
# All complex elements directly reference a relation,
# which can be prefetched directly.
Prefetch(
obj_path,
queryset=cls.get_prefetch_queryset(feature_type, xsd_elements,
fields, output_crs),
) for obj_path, xsd_elements, fields in cls._get_prefetch_summary(
feature_type.xsd_type)
]
if related:
queryset = queryset.prefetch_related(*related)
# Also limit the queryset to the actual fields that are shown.
# No need to request more data
fields = [f.orm_field for f in feature_type.xsd_type.elements]
return queryset.only("pk", *fields)
def annotate_datasets_for_view(datasets: QuerySet,
user: Optional[User] = None,
n: int = None):
datasets = datasets.only('id', 'title', 'description')
annotated_fields = ['id', 'title', 'description', 'task_type']
if user:
if user.is_anonymous:
annotated = datasets.annotate(
is_user_authorised_to_contribute=F('is_public'))
annotated_fields += ['is_user_authorised_to_contribute']
else:
annotated = datasets.annotate(is_user_authorised_admin=Case(
When(Q(user=user) | Q(admins__id=user.id), then=True),
default=False,
output_field=BooleanField()))
annotated = annotated.annotate(
is_user_authorised_to_contribute=Case(
When(Q(is_public=False) & Q(is_user_authorised_admin=False)
& ~Q(contributors__id=user.id),
then=False),
default=True,
output_field=BooleanField()))
annotated_fields += [
'is_user_authorised_admin', 'is_user_authorised_to_contribute'
]
if n:
annotated = annotated.order_by('-created_at')
return annotated.values(*annotated_fields)[:n]
return annotated.values(*annotated_fields)
def optimize_query(query: QuerySet, meta: FieldMeta) -> QuerySet:
model = query.model
only_cols = get_only_cols(model, meta.sections)
select_related_cols, prefetchs = get_related_cols(model, meta.sub_fields)
if only_cols:
query = query.only(*only_cols)
if select_related_cols:
query = query.select_related(*select_related_cols)
if prefetchs:
query = query.prefetch_related(*prefetchs)
return query
def filter_queryset(self, request: Request, queryset: QuerySet,
view: GenericViewSet):
if (not issubclass(view.get_serializer_class(),
FlexFieldsSerializerMixin)
or request.method != "GET"):
return queryset
auto_remove_fields_from_query = getattr(
view, "auto_remove_fields_from_query", True)
auto_select_related_on_query = getattr(view,
"auto_select_related_on_query",
True)
required_query_fields = list(getattr(view, "required_query_fields",
[]))
serializer = view.get_serializer( # type: FlexFieldsSerializerMixin
context=view.get_serializer_context())
serializer.apply_flex_fields()
model_fields = [
self._get_field(field.source, queryset.model)
for field in serializer.fields.values()
if self._get_field(field.source, queryset.model)
]
nested_model_fields = [
self._get_field(field.source, queryset.model)
for field in serializer.fields.values()
if self._get_field(field.source, queryset.model)
and field.field_name in serializer.expanded_fields
]
if auto_remove_fields_from_query:
queryset = queryset.only(*(required_query_fields + [
model_field.name for model_field in model_fields
if not model_field.is_relation or model_field.many_to_one
]))
if auto_select_related_on_query and nested_model_fields:
queryset = queryset.select_related(
*(model_field.name for model_field in nested_model_fields
if model_field.is_relation and model_field.many_to_one))
queryset = queryset.prefetch_related(
*(model_field.name for model_field in nested_model_fields
if model_field.is_relation and not model_field.many_to_one))
return queryset
def switch_team_if_needed_and_possible(self, request: HttpRequest,
target_queryset: QuerySet):
user = cast(User, request.user)
current_team = user.team
if current_team is not None and not target_queryset.filter(
team=current_team).exists():
actual_item = target_queryset.only("team").select_related(
"team").first()
if actual_item is not None:
actual_item_team: Team = actual_item.team
if actual_item_team.get_effective_membership_level(
user.id) is not None:
user.current_team = actual_item_team
user.current_organization_id = actual_item_team.organization_id
user.save()
# Information for POSTHOG_APP_CONTEXT
setattr(request, "switched_team", current_team.id)
def update_returning_pk(qs: QuerySet, updates: dict) -> Set[Any]:
"""
Updates QuerySet items returning primary key values.
This method should not depend on database engine, though can have optimization performances for some engines.
:param qs: QuerySet to update
:param updates: Update items as passed to QuerySet.update(**updates) method
:return: A set of primary keys
"""
qs._for_write = True
if django_pg_returning_available(qs.db) and hasattr(
qs, 'update_returning'):
pk_name = qs.model._meta.pk.name
qs = qs.only(pk_name).update_returning(**updates)
pks = set(qs.values_list(pk_name, flat=True))
else:
with transaction.atomic(using=qs.db):
pks = set(qs.select_for_update().values_list('pk', flat=True))
QuerySet.update(qs, **updates)
return pks
def decorate_queryset(
cls,
feature_type: FeatureType,
queryset: models.QuerySet,
output_crs: CRS,
**params,
) -> models.QuerySet:
"""Apply presentation layer logic to the queryset."""
# Avoid fetching relations, fetch these within the same query,
related = cls._get_prefetch_related(feature_type, output_crs)
if related:
queryset = queryset.prefetch_related(*related)
# Also limit the queryset to the actual fields that are shown.
# No need to request more data
fields = [
f.orm_field for f in feature_type.xsd_type.elements if
not f.is_many or f.is_array # exclude M2M, but include ArrayField
]
return queryset.only("pk", *fields)
