def get_update_sql(
self,
update_fields: Optional[Iterable[str]],
arithmetic_or_function: Optional[Dict[str, Union[ArithmeticExpression,
Function]]],
) -> str:
"""
Generates the SQL for updating a model depending on provided update_fields.
Result is cached for performance.
"""
key = ",".join(update_fields) if update_fields else ""
if key in self.update_cache:
return self.update_cache[key]
arithmetic_or_function = arithmetic_or_function or {}
table = self.model._meta.basetable
query = self.db.query_class.update(table)
count = 0
for field in update_fields or self.model._meta.fields_db_projection.keys(
):
db_column = self.model._meta.fields_db_projection[field]
field_object = self.model._meta.fields_map[field]
if not field_object.pk:
if field not in arithmetic_or_function.keys():
query = query.set(db_column, self.parameter(count))
count += 1
else:
value = F.resolver_arithmetic_expression(
self.model, arithmetic_or_function.get(field))[0]
query = query.set(db_column, value)
query = query.where(
table[self.model._meta.db_pk_column] == self.parameter(count))
sql = self.update_cache[key] = query.get_sql()
return sql
def resolve(self, model: "Type[Model]", table: Table) -> dict:
"""
Used to resolve the Function statement for SQL generation.
:param model: Model the function is applied on to.
:param table: ``pypika.Table`` to keep track of the virtual SQL table
(to allow self referential joins)
:return: Dict with keys ``"joins"`` and ``"fields"``
"""
if isinstance(self.field, str):
function = self._resolve_field_for_model(model, table, self.field)
function["field"] = self._get_function_field(function["field"], *self.default_values)
return function
field, field_object = F.resolver_arithmetic_expression(model, self.field)
if self.populate_field_object:
self.field_object = field_object
return {"joins": [], "field": self._get_function_field(field, *self.default_values)}
def _make_query(self) -> None:
table = self.model._meta.basetable
self.query = self._db.query_class.update(table)
self.resolve_filters(
model=self.model,
q_objects=self.q_objects,
annotations=self.annotations,
custom_filters=self.custom_filters,
)
# Need to get executor to get correct column_map
executor = self._db.executor_class(model=self.model, db=self._db)
for key, value in self.update_kwargs.items():
field_object = self.model._meta.fields_map.get(key)
if not field_object:
raise FieldError(
f"Unknown keyword argument {key} for model {self.model}")
if field_object.pk:
raise IntegrityError(
f"Field {key} is PK and can not be updated")
if isinstance(field_object,
(ForeignKeyFieldInstance, OneToOneFieldInstance)):
fk_field: str = field_object.source_field # type: ignore
db_field = self.model._meta.fields_map[fk_field].source_field
value = executor.column_map[fk_field](getattr(
value, field_object.to_field_instance.model_field_name),
None)
else:
try:
db_field = self.model._meta.fields_db_projection[key]
except KeyError:
raise FieldError(
f"Field {key} is virtual and can not be updated")
if isinstance(value, Term):
value = F.resolver_arithmetic_expression(
self.model, value)[0]
elif isinstance(value, Function):
value = value.resolve(self.model, table)["field"]
else:
value = executor.column_map[key](value, None)
self.query = self.query.set(db_field, value)