GlitchTip/psqlextra/query.py

from collections import OrderedDict
from itertools import chain
from typing import Dict, Iterable, List, Optional, Tuple, Union

from django.core.exceptions import SuspiciousOperation
from django.db import connections, models, router
from django.db.models import Expression, Q
from django.db.models.fields import NOT_PROVIDED

from .sql import PostgresInsertQuery, PostgresQuery
from .types import ConflictAction

ConflictTarget = List[Union[str, Tuple[str]]]


class PostgresQuerySet(models.QuerySet):
    """Adds support for PostgreSQL specifics."""

    def __init__(self, model=None, query=None, using=None, hints=None):
        """Initializes a new instance of :see:PostgresQuerySet."""

        super().__init__(model, query, using, hints)

        self.query = query or PostgresQuery(self.model)

        self.conflict_target = None
        self.conflict_action = None
        self.conflict_update_condition = None
        self.index_predicate = None

    def annotate(self, **annotations):
        """Custom version of the standard annotate function that allows using
        field names as annotated fields.

        Normally, the annotate function doesn't allow you to use the
        name of an existing field on the model as the alias name. This
        version of the function does allow that.

        This is done by temporarily renaming the fields in order to avoid the
        check for conflicts that the base class does.
        We rename all fields instead of the ones that already exist because
        the annotations are stored in an OrderedDict. Renaming only the
        conflicts will mess up the order.
        """
        fields = {field.name: field for field in self.model._meta.get_fields()}

        new_annotations = OrderedDict()

        renames = {}

        for name, value in annotations.items():
            if name in fields:
                new_name = "%s_new" % name
                new_annotations[new_name] = value
                renames[new_name] = name
            else:
                new_annotations[name] = value

        # run the base class's annotate function
        result = super().annotate(**new_annotations)

        # rename the annotations back to as specified
        result.rename_annotations(**renames)
        return result

    def rename_annotations(self, **annotations):
        """Renames the aliases for the specified annotations:

            .annotate(myfield=F('somestuf__myfield'))
            .rename_annotations(myfield='field')

        Arguments:
            annotations:
                The annotations to rename. Mapping the
                old name to the new name.
        """

        self.query.rename_annotations(annotations)
        return self

    def on_conflict(
        self,
        fields: ConflictTarget,
        action: ConflictAction,
        index_predicate: Optional[Union[Expression, Q, str]] = None,
        update_condition: Optional[Union[Expression, Q, str]] = None,
    ):
        """Sets the action to take when conflicts arise when attempting to
        insert/create a new row.

        Arguments:
            fields:
                The fields the conflicts can occur in.

            action:
                The action to take when the conflict occurs.

            index_predicate:
                The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking
                conflicts)

            update_condition:
                Only update if this SQL expression evaluates to true.
        """

        self.conflict_target = fields
        self.conflict_action = action
        self.conflict_update_condition = update_condition
        self.index_predicate = index_predicate

        return self

    def bulk_insert(
        self,
        rows: List[dict],
        return_model: bool = False,
        using: Optional[str] = None,
    ):
        """Creates multiple new records in the database.

        This allows specifying custom conflict behavior using .on_conflict().
        If no special behavior was specified, this uses the normal Django create(..)

        Arguments:
            rows:
                An iterable of dictionaries, where each dictionary
                describes the fields to insert.

            return_model (default: False):
                If model instances should be returned rather than
                just dicts.

            using:
                Name of the database connection to use for
                this query.

        Returns:
            A list of either the dicts of the rows inserted, including the pk or
            the models of the rows inserted with defaults for any fields not specified
        """

        def is_empty(r):
            return all([False for _ in r])

        if not rows or is_empty(rows):
            return []

        if not self.conflict_target and not self.conflict_action:
            # no special action required, use the standard Django bulk_create(..)
            return super().bulk_create(
                [self.model(**fields) for fields in rows]
            )

        deduped_rows = rows

        # when we do a ConflictAction.NOTHING, we are actually
        # doing a ON CONFLICT DO UPDATE with a trick to avoid
        # touching conflicting rows... however, ON CONFLICT UPDATE
        # barfs when you specify the exact same row twice:
        #
        # > "cannot affect row a second time"
        #
        # we filter out the duplicates here to make sure we maintain
        # the same behaviour as the real ON CONFLICT DO NOTHING
        if self.conflict_action == ConflictAction.NOTHING:
            deduped_rows = []
            for row in rows:
                if row in deduped_rows:
                    continue

                deduped_rows.append(row)

        compiler = self._build_insert_compiler(deduped_rows, using=using)
        objs = compiler.execute_sql(return_id=not return_model)
        if return_model:
            return [
                self._create_model_instance(dict(row, **obj), compiler.using)
                for row, obj in zip(deduped_rows, objs)
            ]

        return [dict(row, **obj) for row, obj in zip(deduped_rows, objs)]

    def insert(self, using: Optional[str] = None, **fields):
        """Creates a new record in the database.

        This allows specifying custom conflict behavior using .on_conflict().
        If no special behavior was specified, this uses the normal Django create(..)

        Arguments:
            fields:
                The fields of the row to create.

            using:
                The name of the database connection
                to use for this query.

        Returns:
            The primary key of the record that was created.
        """

        if self.conflict_target or self.conflict_action:
            compiler = self._build_insert_compiler([fields], using=using)
            rows = compiler.execute_sql(return_id=True)

            _, pk_db_column = self.model._meta.pk.get_attname_column()
            if not rows or len(rows) == 0:
                return None

            return rows[0][pk_db_column]

        # no special action required, use the standard Django create(..)
        return super().create(**fields).pk

    def insert_and_get(self, using: Optional[str] = None, **fields):
        """Creates a new record in the database and then gets the entire row.

        This allows specifying custom conflict behavior using .on_conflict().
        If no special behavior was specified, this uses the normal Django create(..)

        Arguments:
            fields:
                The fields of the row to create.

            using:
                The name of the database connection
                to use for this query.

        Returns:
            The model instance representing the row that was created.
        """

        if not self.conflict_target and not self.conflict_action:
            # no special action required, use the standard Django create(..)
            return super().create(**fields)

        compiler = self._build_insert_compiler([fields], using=using)
        rows = compiler.execute_sql(return_id=False)

        if not rows:
            return None

        columns = rows[0]

        # get a list of columns that are officially part of the model and
        # preserve the fact that the attribute name
        # might be different than the database column name
        model_columns = {}
        for field in self.model._meta.local_concrete_fields:
            model_columns[field.column] = field.attname

        # strip out any columns/fields returned by the db that
        # are not present in the model
        model_init_fields = {}
        for column_name, column_value in columns.items():
            try:
                model_init_fields[model_columns[column_name]] = column_value
            except KeyError:
                pass

        return self._create_model_instance(model_init_fields, compiler.using)

    def upsert(
        self,
        conflict_target: ConflictTarget,
        fields: dict,
        index_predicate: Optional[Union[Expression, Q, str]] = None,
        using: Optional[str] = None,
        update_condition: Optional[Union[Expression, Q, str]] = None,
    ) -> int:
        """Creates a new record or updates the existing one with the specified
        data.

        Arguments:
            conflict_target:
                Fields to pass into the ON CONFLICT clause.

            fields:
                Fields to insert/update.

            index_predicate:
                The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking
                conflicts)

            using:
                The name of the database connection to
                use for this query.

            update_condition:
                Only update if this SQL expression evaluates to true.

        Returns:
            The primary key of the row that was created/updated.
        """

        self.on_conflict(
            conflict_target,
            ConflictAction.UPDATE,
            index_predicate=index_predicate,
            update_condition=update_condition,
        )
        return self.insert(**fields, using=using)

    def upsert_and_get(
        self,
        conflict_target: ConflictTarget,
        fields: dict,
        index_predicate: Optional[Union[Expression, Q, str]] = None,
        using: Optional[str] = None,
        update_condition: Optional[Union[Expression, Q, str]] = None,
    ):
        """Creates a new record or updates the existing one with the specified
        data and then gets the row.

        Arguments:
            conflict_target:
                Fields to pass into the ON CONFLICT clause.

            fields:
                Fields to insert/update.

            index_predicate:
                The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking
                conflicts)

            using:
                The name of the database connection to
                use for this query.

            update_condition:
                Only update if this SQL expression evaluates to true.

        Returns:
            The model instance representing the row
            that was created/updated.
        """

        self.on_conflict(
            conflict_target,
            ConflictAction.UPDATE,
            index_predicate=index_predicate,
            update_condition=update_condition,
        )
        return self.insert_and_get(**fields, using=using)

    def bulk_upsert(
        self,
        conflict_target: ConflictTarget,
        rows: Iterable[Dict],
        index_predicate: Optional[Union[Expression, Q, str]] = None,
        return_model: bool = False,
        using: Optional[str] = None,
        update_condition: Optional[Union[Expression, Q, str]] = None,
    ):
        """Creates a set of new records or updates the existing ones with the
        specified data.

        Arguments:
            conflict_target:
                Fields to pass into the ON CONFLICT clause.

            rows:
                Rows to upsert.

            index_predicate:
                The index predicate to satisfy an arbiter partial index (i.e. what partial index to use for checking
                conflicts)

            return_model (default: False):
                If model instances should be returned rather than
                just dicts.

            using:
                The name of the database connection to use
                for this query.

            update_condition:
                Only update if this SQL expression evaluates to true.

        Returns:
            A list of either the dicts of the rows upserted, including the pk or
            the models of the rows upserted
        """

        self.on_conflict(
            conflict_target,
            ConflictAction.UPDATE,
            index_predicate=index_predicate,
            update_condition=update_condition,
        )
        return self.bulk_insert(rows, return_model, using=using)

    def _create_model_instance(
        self, field_values: dict, using: str, apply_converters: bool = True
    ):
        """Creates a new instance of the model with the specified field.

        Use this after the row was inserted into the database. The new
        instance will marked as "saved".
        """

        converted_field_values = field_values.copy()

        if apply_converters:
            connection = connections[using]

            for field in self.model._meta.local_concrete_fields:
                if field.attname not in converted_field_values:
                    continue

                # converters can be defined on the field, or by
                # the database back-end we're using
                field_column = field.get_col(self.model._meta.db_table)
                converters = field.get_db_converters(
                    connection
                ) + connection.ops.get_db_converters(field_column)

                for converter in converters:
                    converted_field_values[field.attname] = converter(
                        converted_field_values[field.attname],
                        field_column,
                        connection,
                    )

        instance = self.model(**converted_field_values)
        instance._state.db = using
        instance._state.adding = False

        return instance

    def _build_insert_compiler(
        self, rows: Iterable[Dict], using: Optional[str] = None
    ):
        """Builds the SQL compiler for a insert query.

        Arguments:
            rows:
                An iterable of dictionaries, where each entry
                describes a record to insert.

            using:
                The name of the database connection to use
                for this query.

        Returns:
            The SQL compiler for the insert.
        """

        # ask the db router which connection to use
        using = (
            using or self._db or router.db_for_write(self.model, **self._hints)
        )

        # create model objects, we also have to detect cases
        # such as:
        #   [dict(first_name='swen'), dict(fist_name='swen', last_name='kooij')]
        # we need to be certain that each row specifies the exact same
        # amount of fields/columns
        objs = []
        rows_iter = iter(rows)
        first_row = next(rows_iter)
        field_count = len(first_row)
        for index, row in enumerate(chain([first_row], rows_iter)):
            if field_count != len(row):
                raise SuspiciousOperation(
                    (
                        "In bulk upserts, you cannot have rows with different field "
                        "configurations. Row {0} has a different field config than "
                        "the first row."
                    ).format(index)
                )

            objs.append(
                self._create_model_instance(row, using, apply_converters=False)
            )

        # get the fields to be used during update/insert
        insert_fields, update_fields = self._get_upsert_fields(first_row)

        # build a normal insert query
        query = PostgresInsertQuery(self.model)
        query.conflict_action = self.conflict_action
        query.conflict_target = self.conflict_target
        query.conflict_update_condition = self.conflict_update_condition
        query.index_predicate = self.index_predicate
        query.values(objs, insert_fields, update_fields)

        compiler = query.get_compiler(using)
        return compiler

    def _is_magical_field(self, model_instance, field, is_insert: bool):
        """Verifies whether this field is gonna modify something on its own.

        "Magical" means that a field modifies the field value
        during the pre_save.

        Arguments:
            model_instance:
                The model instance the field is defined on.

            field:
                The field to get of whether the field is
                magical.

            is_insert:
                Pretend whether this is an insert?

        Returns:
            True when this field modifies something.
        """

        # does this field modify someting upon insert?
        old_value = getattr(model_instance, field.name, None)
        field.pre_save(model_instance, is_insert)
        new_value = getattr(model_instance, field.name, None)

        return old_value != new_value

    def _get_upsert_fields(self, kwargs):
        """Gets the fields to use in an upsert.

        This some nice magic. We'll split the fields into
        a group of "insert fields" and "update fields":

        INSERT INTO bla ("val1", "val2") ON CONFLICT DO UPDATE SET val1 = EXCLUDED.val1

                         ^^^^^^^^^^^^^^                            ^^^^^^^^^^^^^^^^^^^^
                         insert_fields                                 update_fields

        Often, fields appear in both lists. But, for example,
        a :see:DateTime field with `auto_now_add=True` set, will
        only appear in "insert_fields", since it won't be set
        on existing rows.

        Other than that, the user specificies a list of fields
        in the upsert() call. That migt not be all fields. The
        user could decide to leave out optional fields. If we
        end up doing an update, we don't want to overwrite
        those non-specified fields.

        We cannot just take the list of fields the user
        specifies, because as mentioned, some fields
        make modifications to the model on their own.

        We'll have to detect which fields make modifications
        and include them in the list of insert/update fields.
        """

        model_instance = self.model(**kwargs)
        insert_fields = []
        update_fields = []

        for field in model_instance._meta.local_concrete_fields:
            has_default = field.default != NOT_PROVIDED
            if field.name in kwargs or field.column in kwargs:
                insert_fields.append(field)
                update_fields.append(field)
                continue
            elif has_default:
                insert_fields.append(field)
                continue

            # special handling for 'pk' which always refers to
            # the primary key, so if we the user specifies `pk`
            # instead of a concrete field, we have to handle that
            if field.primary_key is True and "pk" in kwargs:
                insert_fields.append(field)
                update_fields.append(field)
                continue

            if self._is_magical_field(model_instance, field, is_insert=True):
                insert_fields.append(field)

            if self._is_magical_field(model_instance, field, is_insert=False):
                update_fields.append(field)

        return insert_fields, update_fields