def apply_inner_select_joins(
self,
query: Query,
select_columns: List[str] = None,
aliases_mapping: Dict[str, AliasedClass] = None,
) -> Query:
"""
Add select load options to query. The goal
is to only SQL select what is requested and join all the necessary
models when dotted notation is used. Inner implies non dotted columns
and many to one and one to one
:param query:
:param select_columns:
:return:
"""
if not select_columns:
return query
joined_models = list()
for column in select_columns:
if is_column_dotted(column):
root_relation = get_column_root_relation(column)
leaf_column = get_column_leaf(column)
if self.is_relation_many_to_one(
root_relation
) or self.is_relation_one_to_one(root_relation):
if root_relation not in joined_models:
query = self._query_join_relation(
query, root_relation, aliases_mapping=aliases_mapping
)
query = query.add_entity(
self.get_alias_mapping(root_relation, aliases_mapping)
)
# Add relation FK to avoid N+1 performance issue
query = self._apply_relation_fks_select_options(
query, root_relation
)
joined_models.append(root_relation)
related_model_ = self.get_alias_mapping(
root_relation, aliases_mapping
)
relation = getattr(self.obj, root_relation)
# The Zen of eager loading :(
# https://docs.sqlalchemy.org/en/13/orm/loading_relationships.html
query = query.options(
contains_eager(relation.of_type(related_model_)).load_only(
leaf_column
)
)
query = query.options(Load(related_model_).load_only(leaf_column))
else:
query = self._apply_normal_col_select_option(query, column)
return query
def get_count(query: Query) -> int:
# Reference: https://gist.github.com/hest/8798884#gistcomment-2301279
disable_group_by = False
if len(query._entities) > 1:
raise Exception(f"only one entity is supported for get_count, got: {query}")
entity = query._entities[0]
if hasattr(entity, "column"):
col = entity.column
if query._group_by and query._distinct:
raise NotImplementedError
if query._group_by or query._distinct:
col = distinct(col)
if query._group_by:
disable_group_by = True
count_func = func.count(col)
else:
count_func = func.count()
if query._group_by and not disable_group_by:
count_func = count_func.over(None)
count_q = (
query.options(lazyload("*"))
.statement.with_only_columns([count_func])
.order_by(None)
)
if disable_group_by:
count_q = count_q.group_by(None)
return query.session.execute(count_q).scalar()
def _apply_relation_fks_select_options(self, query: Query, relation_name) -> Query:
relation = getattr(self.obj, relation_name)
if hasattr(relation, "property"):
local_cols = getattr(self.obj, relation_name).property.local_columns
for local_fk in local_cols:
query = query.options(Load(self.obj).load_only(local_fk.name))
return query
return query
def apply_outer_select_joins(
self, query: Query, select_columns: List[str] = None
) -> Query:
if not select_columns:
return query
for column in select_columns:
if is_column_dotted(column):
root_relation = get_column_root_relation(column)
leaf_column = get_column_leaf(column)
if self.is_relation_many_to_many(
root_relation
) or self.is_relation_one_to_many(root_relation):
query = query.options(
Load(self.obj).joinedload(root_relation).load_only(leaf_column)
)
else:
related_model = self.get_related_model(root_relation)
query = query.options(Load(related_model).load_only(leaf_column))
else:
query = self._apply_normal_col_select_option(query, column)
return query
def build_entity_query(self):
"""
Builds a :class:`sqla:sqlalchemy.orm.query.Query` object for this
entity (an instance of :class:`sir.schema.searchentities.SearchEntity`)
that eagerly loads the values of all search fields.
:rtype: :class:`sqla:sqlalchemy.orm.query.Query`
"""
root_model = self.model
query = Query(root_model)
paths = [field.paths for field in self.fields]
if (config.CFG.getboolean("sir", "wscompat")
and self.extrapaths is not None):
paths.extend([self.extrapaths])
merged_paths = merge_paths(paths)
for field_paths in paths:
for path in field_paths:
current_merged_path = merged_paths
model = root_model
load = Load(model)
split_path = path.split(".")
for pathelem in split_path:
current_merged_path = current_merged_path[pathelem]
column = getattr(model, pathelem)
# __tablename__s in annotation paths
if (not isinstance(column, InstrumentedAttribute)
and not isinstance(column, CompositeProperty)):
break
prop = column.property
if isinstance(prop, RelationshipProperty):
pk = column.mapper.primary_key[0].name
if prop.direction == ONETOMANY:
load = load.subqueryload(pathelem)
elif prop.direction == MANYTOONE:
load = load.joinedload(pathelem)
else:
load = load.defaultload(pathelem)
required_columns = current_merged_path.keys()
required_columns.append(pk)
# Get the mapper class of the current element of the
# path so the next iteration can access it.
model = prop.mapper.class_
# For composite properties, load the columns they
# consist of because eagerly loading a composite
# property doesn't load automatically load them.
composite_columns = filter(
partial(is_composite_column, model),
required_columns)
for composite_column in composite_columns:
composite_parts = (c.name for c in getattr(
model, composite_column).property.columns)
logger.debug("Loading %s instead of %s on %s",
composite_parts, composite_column,
model)
required_columns.remove(composite_column)
required_columns.extend(composite_parts)
logger.debug("Loading only %s on %s", required_columns,
model)
load = defer_everything_but(class_mapper(model), load,
*required_columns)
query = query.options(load)
if self.extraquery is not None:
query = self.extraquery(query)
return query
def _apply_normal_col_select_option(self, query: Query,
column: str) -> Query:
if not self.is_relation(column) and not self.is_property_or_function(
column):
return query.options(Load(self.obj).load_only(column))
return query
