working order by, refactor models merging, refactors ordering into dict

This commit is contained in:
collerek
2020-11-09 12:45:32 +01:00
parent 9f4bde595f
commit ca38f7853e
7 changed files with 259 additions and 37 deletions

View File

@ -37,10 +37,7 @@ class Model(NewBaseModel):
objects: "QuerySet"
def __repr__(self) -> str: # pragma nocover
attrs_to_include = ["tablename", "columns", "pkname"]
_repr = {k: v for k, v in self.Meta.model_fields.items()}
for atr in attrs_to_include:
_repr[atr] = getattr(self.Meta, atr)
_repr = {k: getattr(self, k) for k, v in self.Meta.model_fields.items()}
return f"{self.__class__.__name__}({str(_repr)})"
@classmethod

View File

@ -1,4 +1,5 @@
import inspect
from collections import OrderedDict
from typing import Dict, List, Sequence, Set, TYPE_CHECKING, Type, TypeVar, Union
import ormar
@ -181,11 +182,18 @@ class ModelTableProxy:
@classmethod
def merge_instances_list(cls, result_rows: Sequence["Model"]) -> Sequence["Model"]:
merged_rows: List["Model"] = []
for index, model in enumerate(result_rows):
if index > 0 and model is not None and model.pk == merged_rows[-1].pk:
merged_rows[-1] = cls.merge_two_instances(model, merged_rows[-1])
else:
grouped_instances = OrderedDict()
for model in result_rows:
grouped_instances.setdefault(model.pk, []).append(model)
for group in grouped_instances.values():
model = group.pop(0)
if group:
for next_model in group:
model = cls.merge_two_instances(next_model, model)
merged_rows.append(model)
return merged_rows
@classmethod

View File

@ -1,4 +1,4 @@
from typing import List, NamedTuple, TYPE_CHECKING, Tuple, Type
from typing import Dict, List, NamedTuple, TYPE_CHECKING, Tuple, Type
import sqlalchemy
from sqlalchemy import text
@ -18,19 +18,21 @@ class JoinParameters(NamedTuple):
class SqlJoin:
def __init__(
def __init__( # noqa: CFQ002
self,
used_aliases: List,
select_from: sqlalchemy.sql.select,
order_bys: List[sqlalchemy.sql.elements.TextClause],
columns: List[sqlalchemy.Column],
fields: List,
order_columns: List,
sorted_orders: Dict,
) -> None:
self.used_aliases = used_aliases
self.select_from = select_from
self.order_bys = order_bys
self.columns = columns
self.fields = fields
self.order_columns = order_columns
self.sorted_orders = sorted_orders
@staticmethod
def relation_manager(model_cls: Type["Model"]) -> AliasManager:
@ -46,20 +48,26 @@ class SqlJoin:
def build_join(
self, item: str, join_parameters: JoinParameters
) -> Tuple[List, sqlalchemy.sql.select, List, List]:
) -> Tuple[List, sqlalchemy.sql.select, List, Dict]:
for part in item.split("__"):
if issubclass(
join_parameters.model_cls.Meta.model_fields[part], ManyToManyField
):
_fields = join_parameters.model_cls.Meta.model_fields
new_part = _fields[part].to.get_name()
self._switch_many_to_many_order_columns(part, new_part)
join_parameters = self._build_join_parameters(
part, join_parameters, is_multi=True
)
part = new_part
join_parameters = self._build_join_parameters(part, join_parameters)
return self.used_aliases, self.select_from, self.columns, self.order_bys
return (
self.used_aliases,
self.select_from,
self.columns,
self.sorted_orders,
)
def _build_join_parameters(
self, part: str, join_params: JoinParameters, is_multi: bool = False
@ -108,7 +116,9 @@ class SqlJoin:
)
pkname_alias = model_cls.get_column_alias(model_cls.Meta.pkname)
self.order_bys.append(text(f"{alias}_{to_table}.{pkname_alias}"))
if not is_multi:
self.get_order_bys(alias, to_table, pkname_alias, part)
self_related_fields = model_cls.own_table_columns(
model_cls, self.fields, nested=True,
)
@ -119,6 +129,40 @@ class SqlJoin:
)
self.used_aliases.append(alias)
def _switch_many_to_many_order_columns(self, part: str, new_part: str) -> None:
if self.order_columns:
split_order_columns = [
x.split("__") for x in self.order_columns if "__" in x
]
for condition in split_order_columns:
if condition[-2] == part or condition[-2][1:] == part:
condition[-2] = condition[-2].replace(part, new_part)
self.order_columns = [x for x in self.order_columns if "__" not in x] + [
"__".join(x) for x in split_order_columns
]
@staticmethod
def _check_if_condition_apply(condition: List, part: str) -> bool:
return len(condition) >= 2 and (
condition[-2] == part or condition[-2][1:] == part
)
def get_order_bys(
self, alias: str, to_table: str, pkname_alias: str, part: str
) -> None:
if self.order_columns:
split_order_columns = [
x.split("__") for x in self.order_columns if "__" in x
]
for condition in split_order_columns:
if self._check_if_condition_apply(condition, part):
direction = f"{'desc' if condition[0][0] == '-' else ''}"
order = text(f"{alias}_{to_table}.{condition[-1]} {direction}")
self.sorted_orders["__".join(condition)] = order
else:
order = text(f"{alias}_{to_table}.{pkname_alias}")
self.sorted_orders[pkname_alias] = order
@staticmethod
def get_to_and_from_keys(
join_params: JoinParameters,

View File

@ -1,14 +1,15 @@
from typing import List
from typing import Dict
import sqlalchemy
class OrderQuery:
def __init__(self, order_bys: List) -> None:
self.order_bys = order_bys
def __init__(self, sorted_orders: Dict) -> None:
self.sorted_orders = sorted_orders
def apply(self, expr: sqlalchemy.sql.select) -> sqlalchemy.sql.select:
if self.order_bys:
for order in self.order_bys:
if self.sorted_orders:
for order in list(self.sorted_orders.values()):
if order is not None:
expr = expr.order_by(order)
return expr

View File

@ -1,3 +1,4 @@
from collections import OrderedDict
from typing import List, Optional, TYPE_CHECKING, Tuple, Type
import sqlalchemy
@ -38,21 +39,31 @@ class Query:
self.select_from: List[str] = []
self.columns = [sqlalchemy.Column]
self.order_columns = order_bys
self.order_bys: List[sqlalchemy.sql.elements.TextClause] = []
self.sorted_orders = OrderedDict()
self._init_sorted_orders()
def _init_sorted_orders(self) -> None:
for clause in self.order_columns:
self.sorted_orders[clause] = None
@property
def prefixed_pk_name(self) -> str:
pkname_alias = self.model_cls.get_column_alias(self.model_cls.Meta.pkname)
return f"{self.table.name}.{pkname_alias}"
def apply_order_bys_for_primary_model(self):
def apply_order_bys_for_primary_model(self) -> None:
if self.order_columns:
return [
text(f"{x[1:]} desc") if x.startswith("-") else text(x)
for x in self.order_columns
if "__" not in x
]
return [text(self.prefixed_pk_name)]
for clause in self.order_columns:
if "__" not in clause:
clause = (
text(f"{clause[1:]} desc")
if clause.startswith("-")
else text(clause)
)
self.sorted_orders[clause] = clause
else:
order = text(self.prefixed_pk_name)
self.sorted_orders[self.prefixed_pk_name] = order
def build_select_expression(self) -> Tuple[sqlalchemy.sql.select, List[str]]:
self_related_fields = self.model_cls.own_table_columns(
@ -61,7 +72,7 @@ class Query:
self.columns = self.model_cls.Meta.alias_manager.prefixed_columns(
"", self.table, self_related_fields
)
self.order_bys = self.apply_order_bys_for_primary_model()
self.apply_order_bys_for_primary_model()
self.select_from = self.table
self._select_related.sort(key=lambda item: (item, -len(item)))
@ -75,15 +86,16 @@ class Query:
used_aliases=self.used_aliases,
select_from=self.select_from,
columns=self.columns,
order_bys=self.order_bys,
fields=self.fields,
order_columns=self.order_columns,
sorted_orders=self.sorted_orders,
)
(
self.used_aliases,
self.select_from,
self.columns,
self.order_bys,
self.sorted_orders,
) = sql_join.build_join(item, join_parameters)
expr = sqlalchemy.sql.select(self.columns)
@ -105,7 +117,7 @@ class Query:
)
expr = LimitQuery(limit_count=self.limit_count).apply(expr)
expr = OffsetQuery(query_offset=self.query_offset).apply(expr)
expr = OrderQuery(order_bys=self.order_bys).apply(expr)
expr = OrderQuery(sorted_orders=self.sorted_orders).apply(expr)
return expr
def _reset_query_parameters(self) -> None:

View File

@ -114,7 +114,7 @@ class QuerySet:
order_bys=self.order_bys,
)
exp = qry.build_select_expression()
print(exp.compile(compile_kwargs={"literal_binds": True}))
# print(exp.compile(compile_kwargs={"literal_binds": True}))
return exp
def filter(self, _exclude: bool = False, **kwargs: Any) -> "QuerySet": # noqa: A003

View File

@ -1,3 +1,5 @@
from typing import List, Optional
import databases
import pytest
import sqlalchemy
@ -41,6 +43,45 @@ class Toy(ormar.Model):
owner: Owner = ormar.ForeignKey(Owner)
class Factory(ormar.Model):
class Meta:
tablename = "factories"
metadata = metadata
database = database
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100)
class Car(ormar.Model):
class Meta:
tablename = "cars"
metadata = metadata
database = database
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100)
factory: Optional[Factory] = ormar.ForeignKey(Factory)
class UsersCar(ormar.Model):
class Meta:
tablename = "cars_x_users"
metadata = metadata
database = database
class User(ormar.Model):
class Meta:
tablename = "users"
metadata = metadata
database = database
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100)
cars: List[Car] = ormar.ManyToMany(Car, through=UsersCar)
@pytest.fixture(autouse=True, scope="module")
def create_test_database():
engine = sqlalchemy.create_engine(DATABASE_URL)
@ -93,10 +134,10 @@ async def test_sort_order_on_related_model():
hermes = await Owner.objects.create(name="Hermes")
zeus = await Owner.objects.create(name="Zeus")
await Toy.objects.create(name="Toy 1", owner=zeus)
await Toy.objects.create(name="Toy 4", owner=zeus)
await Toy.objects.create(name="Toy 5", owner=hermes)
await Toy.objects.create(name="Toy 2", owner=aphrodite)
await Toy.objects.create(name="Toy 4", owner=zeus)
await Toy.objects.create(name="Toy 1", owner=zeus)
await Toy.objects.create(name="Toy 3", owner=aphrodite)
await Toy.objects.create(name="Toy 6", owner=hermes)
@ -108,10 +149,129 @@ async def test_sort_order_on_related_model():
assert toys[1].owner == aphrodite
toys = await Toy.objects.select_related("owner").order_by("owner__name").all()
assert toys[0].owner.name == toys[1].owner.name == "Aphrodite"
assert toys[2].owner.name == toys[3].owner.name == "Hermes"
assert toys[4].owner.name == toys[5].owner.name == "Zeus"
owner = (
await Owner.objects.select_related("toys")
.order_by("toys__name")
.filter(name="Zeus")
.get()
)
assert owner.toys[0].name == "Toy 1"
assert owner.toys[1].name == "Toy 4"
owner = (
await Owner.objects.select_related("toys")
.order_by("-toys__name")
.filter(name="Zeus")
.get()
)
assert owner.toys[0].name == "Toy 4"
assert owner.toys[1].name == "Toy 1"
owners = (
await Owner.objects.select_related("toys")
.order_by("-toys__name")
.filter(name__in=["Zeus", "Hermes"])
.all()
)
assert owners[0].toys[0].name == "Toy 6"
assert owners[0].toys[1].name == "Toy 5"
assert owners[0].name == "Hermes"
assert owners[1].toys[0].name == "Toy 4"
assert owners[1].toys[1].name == "Toy 1"
assert owners[1].name == "Zeus"
await Toy.objects.create(name="Toy 7", owner=zeus)
owners = (
await Owner.objects.select_related("toys")
.order_by("-toys__name")
.filter(name__in=["Zeus", "Hermes"])
.all()
)
assert owners[0].toys[0].name == "Toy 7"
assert owners[0].toys[1].name == "Toy 4"
assert owners[0].toys[2].name == "Toy 1"
assert owners[0].name == "Zeus"
assert owners[1].toys[0].name == "Toy 6"
assert owners[1].toys[1].name == "Toy 5"
assert owners[1].name == "Hermes"
@pytest.mark.asyncio
async def test_sort_order_on_many_to_many():
async with database:
factory1 = await Factory.objects.create(name="Factory 1")
factory2 = await Factory.objects.create(name="Factory 2")
car1 = await Car.objects.create(name="Buggy", factory=factory1)
car2 = await Car.objects.create(name="VW", factory=factory2)
car3 = await Car.objects.create(name="Ferrari", factory=factory1)
car4 = await Car.objects.create(name="Volvo", factory=factory2)
car5 = await Car.objects.create(name="Skoda", factory=factory1)
car6 = await Car.objects.create(name="Seat", factory=factory2)
user1 = await User.objects.create(name="Mark")
user2 = await User.objects.create(name="Julie")
await user1.cars.add(car1)
await user1.cars.add(car3)
await user1.cars.add(car4)
await user1.cars.add(car5)
await user2.cars.add(car1)
await user2.cars.add(car2)
await user2.cars.add(car5)
await user2.cars.add(car6)
user = (
await User.objects.select_related("cars")
.filter(name="Mark")
.order_by("cars__name")
.get()
)
assert user.cars[0].name == "Buggy"
assert user.cars[1].name == "Ferrari"
assert user.cars[2].name == "Skoda"
assert user.cars[3].name == "Volvo"
user = (
await User.objects.select_related("cars")
.filter(name="Mark")
.order_by("-cars__name")
.get()
)
assert user.cars[3].name == "Buggy"
assert user.cars[2].name == "Ferrari"
assert user.cars[1].name == "Skoda"
assert user.cars[0].name == "Volvo"
users = await User.objects.select_related("cars").order_by("-cars__name").all()
assert users[0].name == "Mark"
assert users[1].cars[0].name == "VW"
assert users[1].cars[1].name == "Skoda"
assert users[1].cars[2].name == "Seat"
assert users[1].cars[3].name == "Buggy"
users = (
await User.objects.select_related(["cars", "cars__factory"])
.order_by(["-cars__factory__name", "cars__name"])
.all()
)
assert users[0].name == "Julie"
assert users[0].cars[0].name == "Seat"
assert users[0].cars[1].name == "VW"
assert users[0].cars[2].name == "Buggy"
assert users[0].cars[3].name == "Skoda"
assert users[1].name == "Mark"
assert users[1].cars[0].name == "Volvo"
assert users[1].cars[1].name == "Buggy"
assert users[1].cars[2].name == "Ferrari"
assert users[1].cars[3].name == "Skoda"