Merge pull request #7 from collerek/extract_joins

Extract joins
This commit is contained in:
collerek
2020-09-17 21:25:11 +07:00
committed by GitHub
5 changed files with 148 additions and 104 deletions

BIN
.coverage

Binary file not shown.

View File

@ -1,6 +1,6 @@
from ormar.fields.base import BaseField
from ormar.fields.foreign_key import ForeignKey
from ormar.fields.many_to_many import ManyToMany
from ormar.fields.many_to_many import ManyToMany, ManyToManyField
from ormar.fields.model_fields import (
BigInteger,
Boolean,
@ -29,5 +29,6 @@ __all__ = [
"Time",
"ForeignKey",
"ManyToMany",
"ManyToManyField",
"BaseField",
]

131
ormar/queryset/join.py Normal file
View File

@ -0,0 +1,131 @@
from typing import List, NamedTuple, TYPE_CHECKING, Tuple, Type
import sqlalchemy
from sqlalchemy import text
from ormar.fields import ManyToManyField # noqa I100
from ormar.relations import AliasManager
if TYPE_CHECKING: # pragma no cover
from ormar import Model
class JoinParameters(NamedTuple):
prev_model: Type["Model"]
previous_alias: str
from_table: str
model_cls: Type["Model"]
class SqlJoin:
def __init__(
self,
used_aliases: List,
select_from: sqlalchemy.sql.select,
order_bys: List,
columns: List,
) -> None:
self.used_aliases = used_aliases
self.select_from = select_from
self.order_bys = order_bys
self.columns = columns
@staticmethod
def relation_manager(model_cls: Type["Model"]) -> AliasManager:
return model_cls.Meta.alias_manager
@staticmethod
def on_clause(
previous_alias: str, alias: str, from_clause: str, to_clause: str,
) -> text:
left_part = f"{alias}_{to_clause}"
right_part = f"{previous_alias + '_' if previous_alias else ''}{from_clause}"
return text(f"{left_part}={right_part}")
def build_join(
self, item: str, join_parameters: JoinParameters
) -> Tuple[List, sqlalchemy.sql.select, List, List]:
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()
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
def _build_join_parameters(
self, part: str, join_params: JoinParameters, is_multi: bool = False
) -> JoinParameters:
if is_multi:
model_cls = join_params.model_cls.Meta.model_fields[part].through
else:
model_cls = join_params.model_cls.Meta.model_fields[part].to
to_table = model_cls.Meta.table.name
alias = model_cls.Meta.alias_manager.resolve_relation_join(
join_params.from_table, to_table
)
if alias not in self.used_aliases:
self._process_join(join_params, is_multi, model_cls, part, alias)
previous_alias = alias
from_table = to_table
prev_model = model_cls
return JoinParameters(prev_model, previous_alias, from_table, model_cls)
def _process_join(
self,
join_params: JoinParameters,
is_multi: bool,
model_cls: Type["Model"],
part: str,
alias: str,
) -> None:
to_table = model_cls.Meta.table.name
to_key, from_key = self.get_to_and_from_keys(
join_params, is_multi, model_cls, part
)
on_clause = self.on_clause(
previous_alias=join_params.previous_alias,
alias=alias,
from_clause=f"{join_params.from_table}.{from_key}",
to_clause=f"{to_table}.{to_key}",
)
target_table = self.relation_manager(model_cls).prefixed_table_name(
alias, to_table
)
self.select_from = sqlalchemy.sql.outerjoin(
self.select_from, target_table, on_clause
)
self.order_bys.append(text(f"{alias}_{to_table}.{model_cls.Meta.pkname}"))
self.columns.extend(
self.relation_manager(model_cls).prefixed_columns(
alias, model_cls.Meta.table
)
)
self.used_aliases.append(alias)
@staticmethod
def get_to_and_from_keys(
join_params: JoinParameters,
is_multi: bool,
model_cls: Type["Model"],
part: str,
) -> Tuple[str, str]:
if join_params.prev_model.Meta.model_fields[part].virtual or is_multi:
to_field = model_cls.resolve_relation_field(
model_cls, join_params.prev_model
)
to_key = to_field.name
from_key = model_cls.Meta.pkname
else:
to_key = model_cls.Meta.pkname
from_key = part
return to_key, from_key

View File

@ -1,24 +1,16 @@
from typing import List, NamedTuple, TYPE_CHECKING, Tuple, Type
from typing import List, TYPE_CHECKING, Tuple, Type
import sqlalchemy
from sqlalchemy import text
import ormar # noqa I100
from ormar.fields.many_to_many import ManyToManyField
from ormar.queryset import FilterQuery, LimitQuery, OffsetQuery, OrderQuery
from ormar.relations.alias_manager import AliasManager
from ormar.queryset.join import JoinParameters, SqlJoin
if TYPE_CHECKING: # pragma no cover
from ormar import Model
class JoinParameters(NamedTuple):
prev_model: Type["Model"]
previous_alias: str
from_table: str
model_cls: Type["Model"]
class Query:
def __init__(
self,
@ -28,7 +20,6 @@ class Query:
limit_count: int,
offset: int,
) -> None:
self.query_offset = offset
self.limit_count = limit_count
self._select_related = select_related
@ -43,10 +34,6 @@ class Query:
self.columns = None
self.order_bys = None
@property
def relation_manager(self) -> AliasManager:
return self.model_cls.Meta.alias_manager
@property
def prefixed_pk_name(self) -> str:
return f"{self.table.name}.{self.model_cls.Meta.pkname}"
@ -63,17 +50,19 @@ class Query:
self.model_cls, "", self.table.name, self.model_cls
)
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()
join_parameters = self._build_join_parameters(
part, join_parameters, is_multi=True
)
part = new_part
join_parameters = self._build_join_parameters(part, join_parameters)
sql_join = SqlJoin(
used_aliases=self.used_aliases,
select_from=self.select_from,
columns=self.columns,
order_bys=self.order_bys,
)
(
self.used_aliases,
self.select_from,
self.columns,
self.order_bys,
) = sql_join.build_join(item, join_parameters)
expr = sqlalchemy.sql.select(self.columns)
expr = expr.select_from(self.select_from)
@ -85,81 +74,6 @@ class Query:
return expr
@staticmethod
def on_clause(
previous_alias: str, alias: str, from_clause: str, to_clause: str,
) -> text:
left_part = f"{alias}_{to_clause}"
right_part = f"{previous_alias + '_' if previous_alias else ''}{from_clause}"
return text(f"{left_part}={right_part}")
def _build_join_parameters(
self, part: str, join_params: JoinParameters, is_multi: bool = False
) -> JoinParameters:
if is_multi:
model_cls = join_params.model_cls.Meta.model_fields[part].through
else:
model_cls = join_params.model_cls.Meta.model_fields[part].to
to_table = model_cls.Meta.table.name
alias = model_cls.Meta.alias_manager.resolve_relation_join(
join_params.from_table, to_table
)
if alias not in self.used_aliases:
self._process_join(join_params, is_multi, model_cls, part, alias)
previous_alias = alias
from_table = to_table
prev_model = model_cls
return JoinParameters(prev_model, previous_alias, from_table, model_cls)
def _process_join(
self,
join_params: JoinParameters,
is_multi: bool,
model_cls: Type["Model"],
part: str,
alias: str,
) -> None:
to_table = model_cls.Meta.table.name
to_key, from_key = self._get_to_and_from_keys(
join_params, is_multi, model_cls, part
)
on_clause = self.on_clause(
previous_alias=join_params.previous_alias,
alias=alias,
from_clause=f"{join_params.from_table}.{from_key}",
to_clause=f"{to_table}.{to_key}",
)
target_table = self.relation_manager.prefixed_table_name(alias, to_table)
self.select_from = sqlalchemy.sql.outerjoin(
self.select_from, target_table, on_clause
)
self.order_bys.append(text(f"{alias}_{to_table}.{model_cls.Meta.pkname}"))
self.columns.extend(
self.relation_manager.prefixed_columns(alias, model_cls.Meta.table)
)
self.used_aliases.append(alias)
def _get_to_and_from_keys(
self,
join_params: JoinParameters,
is_multi: bool,
model_cls: Type["Model"],
part: str,
) -> Tuple[str, str]:
if join_params.prev_model.Meta.model_fields[part].virtual or is_multi:
to_field = model_cls.resolve_relation_field(
model_cls, join_params.prev_model
)
to_key = to_field.name
from_key = model_cls.Meta.pkname
else:
to_key = model_cls.Meta.pkname
from_key = part
return to_key, from_key
def _apply_expression_modifiers(
self, expr: sqlalchemy.sql.select
) -> sqlalchemy.sql.select:

View File

@ -108,13 +108,11 @@ class QuerySet:
async def exists(self) -> bool:
expr = self.build_select_expression()
expr = sqlalchemy.exists(expr).select()
# print(expr.compile(compile_kwargs={"literal_binds": True}))
return await self.database.fetch_val(expr)
async def count(self) -> int:
expr = self.build_select_expression().alias("subquery_for_count")
expr = sqlalchemy.func.count().select().select_from(expr)
# print(expr.compile(compile_kwargs={"literal_binds": True}))
return await self.database.fetch_val(expr)
async def delete(self, **kwargs: Any) -> int: