Array many to many patches for django 3.1.4

# VERY IMPORTANT
# this code in core.lib.django_postgres_extensions.array_m2m_patches
# needs to be update in every django version update in order to avoid
# issues, our own addition to django's patched functions are made
# between `OUR CUSTOM CODE` comments, so the rest of the original
# django code must be update regularly
import copy

from django.core import exceptions
from django.core.exceptions import FieldError
from django.db import transaction
from django.db.models.constants import LOOKUP_SEP
from django.db.models.sql.constants import CURSOR
from django.db.models.sql.subqueries import UpdateQuery as BaseUpdateQuery


def prefetch_one_level(instances, prefetcher, lookup, level):
    """
    Helper function for prefetch_related_objects().

    Run prefetches on all instances using the prefetcher object,
    assigning results to relevant caches in instance.

    Return the prefetched objects along with any additional prefetches that
    must be done due to prefetch_related lookups found from default managers.
    """
    # prefetcher must have a method get_prefetch_queryset() which takes a list
    # of instances, and returns a tuple:

    # (queryset of instances of self.model that are related to passed in instances,
    #  callable that gets value to be matched for returned instances,
    #  callable that gets value to be matched for passed in instances,
    #  boolean that is True for singly related objects,
    #  cache or field name to assign to,
    #  boolean that is True when the previous argument is a cache name vs a field name).

    # The 'values to be matched' must be hashable as they will be used
    # in a dictionary.

    rel_qs, rel_obj_attr, instance_attr, single, cache_name, is_descriptor = (
        prefetcher.get_prefetch_queryset(instances, lookup.get_current_queryset(level)))
    # We have to handle the possibility that the QuerySet we just got back
    # contains some prefetch_related lookups. We don't want to trigger the
    # prefetch_related functionality by evaluating the query. Rather, we need
    # to merge in the prefetch_related lookups.
    # Copy the lookups in case it is a Prefetch object which could be reused
    # later (happens in nested prefetch_related).
    additional_lookups = [
        copy.copy(additional_lookup) for additional_lookup
        in getattr(rel_qs, '_prefetch_related_lookups', ())
    ]
    if additional_lookups:
        # Don't need to clone because the manager should have given us a fresh
        # instance, so we access an internal instead of using public interface
        # for performance reasons.
        rel_qs._prefetch_related_lookups = ()

    all_related_objects = list(rel_qs)
    # OUR CUSTOM CODE
    is_multi_reference = getattr(rel_qs, 'is_multi_reference', False)
    # END OF OUR CUSTOM CODE

    rel_obj_cache = {}
    # OUR CUSTOM CODE
    if not is_multi_reference:
        # END OF OUR CUSTOM CODE
        for rel_obj in all_related_objects:
            rel_attr_val = rel_obj_attr(rel_obj)
            rel_obj_cache.setdefault(rel_attr_val, []).append(rel_obj)

    to_attr, as_attr = lookup.get_current_to_attr(level)
    # Make sure `to_attr` does not conflict with a field.
    if as_attr and instances:
        # We assume that objects retrieved are homogeneous (which is the premise
        # of prefetch_related), so what applies to first object applies to all.
        model = instances[0].__class__
        try:
            model._meta.get_field(to_attr)
        except exceptions.FieldDoesNotExist:
            pass
        else:
            msg = 'to_attr={} conflicts with a field on the {} model.'
            raise ValueError(msg.format(to_attr, model.__name__))

    # Whether or not we're prefetching the last part of the lookup.
    leaf = len(lookup.prefetch_through.split(LOOKUP_SEP)) - 1 == level

    for obj in instances:
        instance_attr_val = instance_attr(obj)
        # OUR CUSTOM CODE
        if is_multi_reference:
            vals = [rel_obj for rel_obj in all_related_objects if rel_obj_attr(rel_obj, instance_attr_val)]
        else:
            # END OF OUR CUSTOM CODE
            vals = rel_obj_cache.get(instance_attr_val, [])

        if single:
            val = vals[0] if vals else None
            if as_attr:
                # A to_attr has been given for the prefetch.
                setattr(obj, to_attr, val)
            elif is_descriptor:
                # cache_name points to a field name in obj.
                # This field is a descriptor for a related object.
                setattr(obj, cache_name, val)
            else:
                # No to_attr has been given for this prefetch operation and the
                # cache_name does not point to a descriptor. Store the value of
                # the field in the object's field cache.
                obj._state.fields_cache[cache_name] = val
        else:
            if as_attr:
                setattr(obj, to_attr, vals)
            else:
                manager = getattr(obj, to_attr)
                if leaf and lookup.queryset is not None:
                    qs = manager._apply_rel_filters(lookup.queryset)
                else:
                    qs = manager.get_queryset()
                qs._result_cache = vals
                # We don't want the individual qs doing prefetch_related now,
                # since we have merged this into the current work.
                qs._prefetch_done = True
                obj._prefetched_objects_cache[cache_name] = qs
    return all_related_objects, additional_lookups


def as_sql(self, compiler, connection):
    """
    Generate the full
       LEFT OUTER JOIN sometable ON sometable.somecol = othertable.othercol, params
    clause for this join.
    """
    join_conditions = []
    params = []
    qn = compiler.quote_name_unless_alias
    qn2 = connection.ops.quote_name

    # Add a join condition for each pair of joining columns.
    for lhs_col, rhs_col in self.join_cols:
        # OUR CUSTOM CODE
        if hasattr(self.join_field, 'get_join_on'):
            join_condition = self.join_field.get_join_on(qn(self.parent_alias), qn2(lhs_col), qn(self.table_alias),
                                                         qn2(rhs_col))
            join_conditions.append(join_condition)
        else:
            # END OF OUR CUSTOM CODE
            join_conditions.append('%s.%s = %s.%s' % (
                qn(self.parent_alias),
                qn2(lhs_col),
                qn(self.table_alias),
                qn2(rhs_col),
            ))

    # Add a single condition inside parentheses for whatever
    # get_extra_restriction() returns.
    extra_cond = self.join_field.get_extra_restriction(
        compiler.query.where_class, self.table_alias, self.parent_alias)
    if extra_cond:
        extra_sql, extra_params = compiler.compile(extra_cond)
        join_conditions.append('(%s)' % extra_sql)
        params.extend(extra_params)
    if self.filtered_relation:
        extra_sql, extra_params = compiler.compile(self.filtered_relation)
        if extra_sql:
            join_conditions.append('(%s)' % extra_sql)
            params.extend(extra_params)
    if not join_conditions:
        # This might be a rel on the other end of an actual declared field.
        declared_field = getattr(self.join_field, 'field', self.join_field)
        raise ValueError(
            "Join generated an empty ON clause. %s did not yield either "
            "joining columns or extra restrictions." % declared_field.__class__
        )
    on_clause_sql = ' AND '.join(join_conditions)
    alias_str = '' if self.table_alias == self.table_name else (' %s' % self.table_alias)
    sql = '%s %s%s ON (%s)' % (self.join_type, qn(self.table_name), alias_str, on_clause_sql)
    return sql, params


def update(self, **kwargs):
    """
    Update all elements in the current QuerySet, setting all the given
    fields to the appropriate values.
    """
    self._not_support_combined_queries('update')
    assert not self.query.is_sliced, \
        "Cannot update a query once a slice has been taken."
    self._for_write = True
    # OUR CUSTOM CODE is UpdateQuery
    query = self.query.chain(UpdateQuery)
    query.add_update_values(kwargs)
    # Clear any annotations so that they won't be present in subqueries.
    query.annotations = {}
    with transaction.mark_for_rollback_on_error(using=self.db):
        rows = query.get_compiler(self.db).execute_sql(CURSOR)
    self._result_cache = None
    return rows


update.alters_data = True


def _update(self, values):
    """
    A version of update() that accepts field objects instead of field names.
    Used primarily for model saving and not intended for use by general
    code (it requires too much poking around at model internals to be
    useful at that level).
    """
    assert not self.query.is_sliced, \
        "Cannot update a query once a slice has been taken."
    query = self.query.chain(UpdateQuery)
    # OUR CUSTOM CODE
    values = [value for value in values if not getattr(value[0], 'many_to_many_array', False)]
    # END OF OUR CUSTOM CODE

    query.add_update_fields(values)
    # Clear any annotations so that they won't be present in subqueries.
    query.annotations = {}
    self._result_cache = None
    return query.get_compiler(self.db).execute_sql(CURSOR)


_update.alters_data = True
_update.queryset_only = False


class UpdateQuery(BaseUpdateQuery):
    def add_update_values(self, values):
        """
        Convert a dictionary of field name to value mappings into an update
        query. This is the entry point for the public update() method on
        querysets.
        """
        values_seq = []
        for name, val in values.items():
            # OUR CUSTOM CODE
            if '__' in name:
                indexes = name.split('__')
                field_name = indexes.pop(0)
                field = self.get_meta().get_field(field_name)
                val = field.get_update_type(indexes, val)
                model = field.model
            else:
                # END OF OUR CUSTOM CODE
                field = self.get_meta().get_field(name)
                direct = not (field.auto_created and not field.concrete) or not field.concrete
                model = field.model._meta.concrete_model
                if not direct or (field.is_relation and field.many_to_many):
                    raise FieldError(
                        'Cannot update model field %r (only non-relations and '
                        'foreign keys permitted).' % field
                    )
                if model is not self.get_meta().concrete_model:
                    self.add_related_update(model, field, val)
                    continue
            values_seq.append((field, model, val))
        return self.add_update_fields(values_seq)