def _validate_returning_term(self, term: Term) -> None:
for field in term.fields_():
if not any([self._insert_table, self._update_table, self._delete_from]):
raise QueryException("Returning can't be used in this query")
if (
field.table not in {self._insert_table, self._update_table}
and term not in self._from
):
raise QueryException("You can't return from other tables")
def do_update(self, update_field: Union[str, Field], update_value: Any) -> "PostgreQueryBuilder":
if self._on_conflict_do_nothing:
raise QueryException("Can not have two conflict handlers")
if isinstance(update_field, str):
field = self._conflict_field_str(update_field)
elif isinstance(update_field, Field):
field = update_field
else:
raise QueryException("Unsupported update_field")
self._on_conflict_do_updates.append((field, ValueWrapper(update_value)))
def on_conflict(self, *target_fields: Union[str, Term]) -> "PostgreQueryBuilder":
if not self._insert_table:
raise QueryException("On conflict only applies to insert query")
self._on_conflict = True
for target_field in target_fields:
if isinstance(target_field, str):
self._on_conflict_fields.append(self._conflict_field_str(target_field))
elif isinstance(target_field, Term):
self._on_conflict_fields.append(target_field)
def top(self, value: Union[str, int]) -> "MSSQLQueryBuilder":
"""
Implements support for simple TOP clauses.
Does not include support for PERCENT or WITH TIES.
https://docs.microsoft.com/en-us/sql/t-sql/queries/top-transact-sql?view=sql-server-2017
"""
try:
self._top = int(value)
except ValueError:
raise QueryException("TOP value must be an integer")
def returning(self, *terms: Any) -> "PostgreQueryBuilder":
for term in terms:
if isinstance(term, Field):
self._return_field(term)
elif isinstance(term, str):
self._return_field_str(term)
elif isinstance(term, ArithmeticExpression):
self._return_other(term)
elif isinstance(term, Function):
raise QueryException("Aggregate functions are not allowed in returning")
else:
self._return_other(self.wrap_constant(term, self._wrapper_cls))
def _on_conflict_sql(self, **kwargs: Any) -> str:
if not self._on_conflict_do_nothing and len(self._on_conflict_do_updates) == 0:
if not self._on_conflict_fields:
return ""
raise QueryException("No handler defined for on conflict")
if self._on_conflict_do_updates and not self._on_conflict_fields:
raise QueryException("Can not have fieldless on conflict do update")
conflict_query = " ON CONFLICT"
if self._on_conflict_fields:
fields = [f.get_sql(with_alias=True, **kwargs)
for f in self._on_conflict_fields]
conflict_query += " (" + ', '.join(fields) + ")"
if self._on_conflict_wheres:
conflict_query += " WHERE {where}".format(
where=self._on_conflict_wheres.get_sql(subquery=True, **kwargs)
)
return conflict_query
def where(self, criterion: Criterion) -> "PostgreQueryBuilder":
if not self._on_conflict:
return super().where(criterion)
if isinstance(criterion, EmptyCriterion):
return
if self._on_conflict_do_nothing:
raise QueryException('DO NOTHING doest not support WHERE')
if self._on_conflict_fields and self._on_conflict_do_updates:
if self._on_conflict_do_update_wheres:
self._on_conflict_do_update_wheres &= criterion
else:
self._on_conflict_do_update_wheres = criterion
elif self._on_conflict_fields:
if self._on_conflict_wheres:
self._on_conflict_wheres &= criterion
else:
self._on_conflict_wheres = criterion
else:
raise QueryException('Can not have fieldless ON CONFLICT WHERE')
def _return_field_str(self, term: Union[str, Field]) -> None:
if term == "*":
self._set_returns_for_star()
self._returns.append(Star())
return
if self._insert_table:
self._return_field(Field(term, table=self._insert_table))
elif self._update_table:
self._return_field(Field(term, table=self._update_table))
elif self._delete_from:
self._return_field(Field(term, table=self._from[0]))
else:
raise QueryException("Returning can't be used in this query")
def on_duplicate_key_ignore(self) -> "MySQLQueryBuilder":
if self._duplicate_updates:
raise QueryException("Can not have two conflict handlers")
self._ignore_duplicates = True
def on_duplicate_key_update(self, field: Union[Field, str], value: Any) -> "MySQLQueryBuilder":
if self._ignore_duplicates:
raise QueryException("Can not have two conflict handlers")
field = Field(field) if not isinstance(field, Field) else field
self._duplicate_updates.append((field, ValueWrapper(value)))
def do_nothing(self) -> "PostgreQueryBuilder":
if len(self._on_conflict_do_updates) > 0:
raise QueryException("Can not have two conflict handlers")
self._on_conflict_do_nothing = True