def join_records(
select: sqlalchemy.sql.Select,
location_table: sqlalchemy.schema.Table
) -> sqlalchemy.sql.FromClause:
# mypy needs to be sure
assert isinstance(records_table, ByNameOpaqueTableStorage)
return select.select_from(
records_table._table.join(
location_table,
onclause=records_table._table.columns.dataset_id ==
location_table.columns.dataset_id,
)).where(location_table.columns.datastore_name ==
self.datastoreName)
def prepare_query(self, q: sa.sql.Select) -> sa.sql.Select:
""" Prepare the statement for loading: add columns to select, add filter condition
Args:
q: SELECT statement prepared by QueryExecutor.statement().
It has no columns yet, but has a select_from(self.target_model), unaliased.
NOTE: we never alias the target model: the one we're loading. It would've made things too complicated.
"""
# Use SelectInLoader
# self.query_info: primary key columns, the IN expression, etc
# self._parent_alias: used with JOINed relationships where our table has to be joined to an alias of the parent table
query_info = self.loader._query_info
parent_alias = self.loader._parent_alias if query_info.load_with_join else NotImplemented
effective_entity = self.target_model
# [ADDED] Adapt pk_cols
# [o] pk_cols = query_info.pk_cols
# [o] in_expr = query_info.in_expr
pk_cols = query_info.pk_cols
in_expr = query_info.in_expr
# [o] if not query_info.load_with_join:
if not query_info.load_with_join:
# [o] if effective_entity.is_aliased_class:
# [o] pk_cols = [ effective_entity._adapt_element(col) for col in pk_cols ]
# [o] in_expr = effective_entity._adapt_element(in_expr)
adapter = SimpleColumnsAdapter(self.target_model)
pk_cols = adapter.replace_many(pk_cols)
in_expr = adapter.replace(in_expr)
# [o] bundle_ent = orm_util.Bundle("pk", *pk_cols)
# [o] entity_sql = effective_entity.__clause_element__()
# [o] q = Select._create_raw_select(
# [o] _raw_columns=[bundle_sql, entity_sql],
# [o] _label_style=LABEL_STYLE_TABLENAME_PLUS_COL,
# [CUSTOMIZED]
if not query_info.load_with_join:
q = add_columns(q, pk_cols) # [CUSTOMIZED]
else:
# NOTE: we cannot always add our FK columns: when `load_with_join` is used, these columns
# may actually refer to columns from a M2M table with conflicting names!
# Example:
# SELECT articles.id, tags.id
# FROM articles JOIN ... JOIN tags
# So we have to rename them. We use "table.column" aliases because this horrible "." makes it clear
# it's not just another column
# label_prefix = self.source_model.__table__.name + '.'
self.fk_label_prefix = self.source_model.__tablename__ + '.' # type: ignore[union-attr]
q = add_columns(q, [ # [CUSTOMIZED]
col.label(self.fk_label_prefix + col.key)
for col in pk_cols
])
# Effective entity
# This is the class that we select from
# [o] if not query_info.load_with_join:
# [o] q = q.select_from(effective_entity)
# [o] else:
# [o] q = q.select_from(self._parent_alias).join(...)
# [CUSTOMIZED]
if not query_info.load_with_join:
q = q.select_from(self.target_model)
else:
if SA_13:
q = q.select_from(
sa.orm.join(parent_alias, self.target_model, onclause=getattr(parent_alias, self.key).of_type(self.target_model))
)
else:
q = q.select_from(parent_alias).join(
getattr(parent_alias, self.key).of_type(self.target_model)
)
# [o] q = q.filter(in_expr.in_(sql.bindparam("primary_keys")))
if SA_13:
q = q.where(in_expr.in_(sa.sql.bindparam("primary_keys", expanding=True)))
else:
q = q.filter(in_expr.in_(sa.sql.bindparam("primary_keys")))
return q