Tolerate new columns being added to tables during upsert

Up until now, upserting to a table which had a column added not
know to Django would make the query crash. This commit introduces
a more robust mechanism to constructing model instances from
query results that tolerates a column being added at the end
of the table.

Author: Photonios

psqlextra/compiler.py

@@ -3,7 +3,7 @@
 import sys
 
 from collections.abc import Iterable
-from typing import Tuple, Union
+from typing import TYPE_CHECKING, Tuple, Union, cast
 
 import django
 
@@ -12,11 +12,13 @@
 from django.db.models import Expression, Model, Q
 from django.db.models.fields.related import RelatedField
 from django.db.models.sql import compiler as django_compiler
-from django.db.utils import ProgrammingError
 
 from .expressions import HStoreValue
 from .types import ConflictAction
 
+if TYPE_CHECKING:
+    from .sql import PostgresInsertQuery
+
 
 def append_caller_to_sql(sql):
     """Append the caller to SQL queries.
@@ -161,6 +163,8 @@ def as_sql(self, *args, **kwargs):
 class PostgresInsertOnConflictCompiler(django_compiler.SQLInsertCompiler):  # type: ignore [name-defined]
     """Compiler for SQL INSERT statements."""
 
+    query: "PostgresInsertQuery"
+
     def __init__(self, *args, **kwargs):
         """Initializes a new instance of
         :see:PostgresInsertOnConflictCompiler."""
@@ -169,35 +173,14 @@ def __init__(self, *args, **kwargs):
 
     def as_sql(self, return_id=False, *args, **kwargs):
         """Builds the SQL INSERT statement."""
+
         queries = [
             self._rewrite_insert(sql, params, return_id)
             for sql, params in super().as_sql(*args, **kwargs)
         ]
 
         return queries
 
-    def execute_sql(self, return_id=False):
-        # execute all the generate queries
-        with self.connection.cursor() as cursor:
-            rows = []
-            for sql, params in self.as_sql(return_id):
-                cursor.execute(sql, params)
-                try:
-                    rows.extend(cursor.fetchall())
-                except ProgrammingError:
-                    pass
-            description = cursor.description
-
-        # create a mapping between column names and column value
-        return [
-            {
-                column.name: row[column_index]
-                for column_index, column in enumerate(description)
-                if row
-            }
-            for row in rows
-        ]
-
     def _rewrite_insert(self, sql, params, return_id=False):
         """Rewrites a formed SQL INSERT query to include the ON CONFLICT
         clause.
@@ -209,9 +192,9 @@ def _rewrite_insert(self, sql, params, return_id=False):
             params:
                 The parameters passed to the query.
 
-            returning:
-                What to put in the `RETURNING` clause
-                of the resulting query.
+            return_id:
+                Whether to only return the ID or all
+                columns.
 
         Returns:
             A tuple of the rewritten SQL query and new params.
@@ -284,7 +267,7 @@ def _build_set_statement(self) -> Tuple[str, tuple]:
         # the compiler and the queries.
         from .sql import PostgresUpdateQuery
 
-        query = self.query.chain(PostgresUpdateQuery)
+        query = cast(PostgresUpdateQuery, self.query.chain(PostgresUpdateQuery))
         query.add_update_values(self.query.update_values)
 
         sql, params = query.get_compiler(self.connection.alias).as_sql()

psqlextra/introspect/__init__.py

@@ -0,0 +1,8 @@
+from .fields import inspect_model_local_concrete_fields
+from .models import model_from_cursor, models_from_cursor
+
+__all__ = [
+    "models_from_cursor",
+    "model_from_cursor",
+    "inspect_model_local_concrete_fields",
+]

psqlextra/introspect/fields.py

@@ -0,0 +1,21 @@
+from typing import List, Type
+
+from django.db.models import Field, Model
+
+
+def inspect_model_local_concrete_fields(model: Type[Model]) -> List[Field]:
+    """Gets a complete list of local and concrete fields on a model, these are
+    fields that directly map to a database colmn directly on the table backing
+    the model.
+
+    This is similar to Django's `Meta.local_concrete_fields`, which is a
+    private API. This method utilizes only public APIs.
+    """
+
+    local_concrete_fields = []
+
+    for field in model._meta.get_fields(include_parents=False):
+        if isinstance(field, Field) and field.column and not field.many_to_many:
+            local_concrete_fields.append(field)
+
+    return local_concrete_fields

psqlextra/introspect/models.py

@@ -0,0 +1,170 @@
+from typing import (
+    Any,
+    Dict,
+    Generator,
+    Iterable,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    cast,
+)
+
+from django.core.exceptions import FieldDoesNotExist
+from django.db import connection, models
+from django.db.models import Field, Model
+from django.db.models.expressions import Expression
+
+from .fields import inspect_model_local_concrete_fields
+
+TModel = TypeVar("TModel", bound=models.Model)
+
+
+def _construct_model(
+    model: Type[TModel],
+    columns: Iterable[str],
+    values: Iterable[Any],
+    *,
+    apply_converters: bool = True
+) -> TModel:
+    fields_by_name_and_column = {}
+    for field in inspect_model_local_concrete_fields(model):
+        fields_by_name_and_column[field.attname] = field
+
+        if field.db_column:
+            fields_by_name_and_column[field.db_column] = field
+
+    indexable_columns = list(columns)
+
+    row = {}
+
+    for index, value in enumerate(values):
+        column = indexable_columns[index]
+        try:
+            field = cast(Field, model._meta.get_field(column))
+        except FieldDoesNotExist:
+            field = fields_by_name_and_column[column]
+
+        field_column_expression = field.get_col(model._meta.db_table)
+
+        if apply_converters:
+            converters = cast(Expression, field).get_db_converters(
+                connection
+            ) + connection.ops.get_db_converters(field_column_expression)
+
+            converted_value = value
+            for converter in converters:
+                converted_value = converter(
+                    converted_value,
+                    field_column_expression,
+                    connection,
+                )
+        else:
+            converted_value = value
+
+        row[field.attname] = converted_value
+
+    instance = model(**row)
+    instance._state.adding = False
+    instance._state.db = connection.alias
+
+    return instance
+
+
+def models_from_cursor(
+    model: Type[TModel], cursor, *, related_fields: List[str] = []
+) -> Generator[TModel, None, None]:
+    """Fetches all rows from a cursor and converts the values into model
+    instances.
+
+    This is roughly what Django does internally when you do queries. This
+    goes further than `Model.from_db` as it also applies converters to make
+    sure that values are converted into their Python equivalent.
+
+    Use this when you've outgrown the ORM and you are writing performant
+    queries yourself and you need to map the results back into ORM objects.
+
+    Arguments:
+        model:
+            Model to construct.
+
+        cursor:
+            Cursor to read the rows from.
+
+        related_fields:
+            List of ForeignKey/OneToOneField names that were joined
+            into the raw query. Use this to achieve the same thing
+            that Django's `.select_related()` does.
+
+            Field names should be specified in the order that they
+            are SELECT'd in.
+    """
+
+    columns = [col[0] for col in cursor.description]
+    field_offset = len(inspect_model_local_concrete_fields(model))
+
+    rows = cursor.fetchmany()
+
+    while rows:
+        for values in rows:
+            instance = _construct_model(
+                model, columns[:field_offset], values[:field_offset]
+            )
+
+            for index, related_field_name in enumerate(related_fields):
+                related_model = model._meta.get_field(
+                    related_field_name
+                ).related_model
+                if not related_model:
+                    continue
+
+                related_field_count = len(
+                    inspect_model_local_concrete_fields(related_model)
+                )
+
+                # autopep8: off
+                related_columns = columns[
+                    field_offset : field_offset + related_field_count  # noqa
+                ]
+                related_values = values[
+                    field_offset : field_offset + related_field_count  # noqa
+                ]
+                # autopep8: one
+
+                if (
+                    not related_columns
+                    or not related_values
+                    or all([value is None for value in related_values])
+                ):
+                    continue
+
+                related_instance = _construct_model(
+                    cast(Type[Model], related_model),
+                    related_columns,
+                    related_values,
+                )
+                instance._state.fields_cache[related_field_name] = related_instance  # type: ignore
+
+                field_offset += len(
+                    inspect_model_local_concrete_fields(related_model)
+                )
+
+            yield instance
+
+        rows = cursor.fetchmany()
+
+
+def model_from_cursor(
+    model: Type[TModel], cursor, *, related_fields: List[str] = []
+) -> Optional[TModel]:
+    return next(
+        models_from_cursor(model, cursor, related_fields=related_fields), None
+    )
+
+
+def model_from_dict(
+    model: Type[TModel], row: Dict[str, Any], *, apply_converters: bool = True
+) -> TModel:
+    return _construct_model(
+        model, row.keys(), row.values(), apply_converters=apply_converters
+    )