diff --git a/specifyweb/backend/stored_queries/execution.py b/specifyweb/backend/stored_queries/execution.py index cc893fb1516..94ec53ad531 100644 --- a/specifyweb/backend/stored_queries/execution.py +++ b/specifyweb/backend/stored_queries/execution.py @@ -52,6 +52,11 @@ SERIES_MAX_ROWS = 10000 +def _get_tree_rank_functions(): + from .tree_rank import prepare_tree_ranks, rank_path, path_key, resolve_page + return prepare_tree_ranks, rank_path, path_key, resolve_page + + class QuerySort: SORT_TYPES = [None, asc, desc] @@ -81,6 +86,8 @@ class BuildQueryProps(NamedTuple): search_synonymy: bool = False implicit_or: bool = True formatter_props: ObjectFormatterProps = DefaultQueryFormatterProps() + defer_tree_display: bool = False + count_only: bool = False def set_group_concat_max_len(connection): @@ -892,12 +899,14 @@ def execute( series=series, search_synonymy=search_synonymy, formatter_props=formatter_props, + defer_tree_display=(count_only or bool(limit)) and not ( + distinct or series or search_synonymy or formatauditobjs), + count_only=count_only, ), ) - log_sqlalchemy_query(query) - if count_only: + log_sqlalchemy_query(query) if series: cat_num_sort_type = 0 for query_field in query_fields: @@ -930,14 +939,20 @@ def execute( order_by_exprs.insert(0, text("collectionobject.`CatalogNumber`")) # query = query.limit(SERIES_MAX_ROWS) + log_sqlalchemy_query(query) return {'results': series_post_query(query, limit=limit, offset=offset, sort_type=cat_num_sort_type)} logger.debug("order by: %s", order_by_exprs) - query = query.order_by(*order_by_exprs).offset(offset) - - if limit: - query = query.limit(limit) + ranks = query.get_execution_options().get("specify_tree_ranks") + if ranks is not None: + resolve_page = _get_tree_rank_functions()[3] + query = resolve_page(query, order_by_exprs, ranks, limit, offset) + else: + query = query.order_by(*order_by_exprs).offset(offset) + if limit: + query = query.limit(limit) + log_sqlalchemy_query(query) results = list( apply_special_post_query_processing( query=query, @@ -1054,8 +1069,9 @@ def apply_where_condition_to_query( where = reduce(sql.and_, implicit_ors) query = query.filter(where) else: - where = reduce(sql.and_, (p for ps in predicates_by_fieldspec.values() for p in ps)) - query = query.filter(where) + predicates = [p for ps in predicates_by_fieldspec.values() for p in ps] + if predicates: + query = query.filter(reduce(sql.and_, predicates)) return query # REFACTOR: Clean up this and the other QueryConstruct functions @@ -1069,17 +1085,47 @@ def add_fields_to_query( query_fields: list[QueryField], series: bool = False, formatauditobjs: bool = False, - use_implicit_ors: bool = True + use_implicit_ors: bool = True, + defer_tree_display: bool = False, + count_only: bool = False, ): + prepare_tree_ranks, rank_path, path_key, _ = _get_tree_rank_functions() + query_fields = list(query_fields) + ordinary_predicates = defaultdict(list) + if any(rank_path(f.fieldspec) is not None for f in query_fields): + for query_field in query_fields: + if (rank_path(query_field.fieldspec) is None + and (query_field.op_num != 8 or query_field.negate)): + query, _, predicate = query_field.add_to_query( + query, formatauditobjs=formatauditobjs, collection=collection, user=user) + if predicate is not None: + ordinary_predicates[query_field.fieldspec].append(predicate) + query = apply_where_condition_to_query(query, ordinary_predicates, use_implicit_ors) + tree_tables = { + rank_path(f.fieldspec)[1].relatedModelName + for f in query_fields if rank_path(f.fieldspec) is not None + } + # Grouped/formatted object queries retain eager resolution. A page may + # defer scalar ranks from one tree table without changing row semantics. + defer_tree_display = defer_tree_display and len(tree_tables) == 1 and all( + not f.fieldspec.is_relationship() and f.fieldspec.get_field() is not None + for f in query_fields + ) + query = prepare_tree_ranks(query, query_fields, defer_tree_display, count_only) + ranks = query.join_cache.get("tree_ranks") order_by_exprs = [] selected_fields = [] predicates_by_fieldspec = defaultdict(list) for query_field in query_fields: + if ranks is not None: + ranks.current_field = query_field + ranks.current_rank = None sort_type = QuerySort.by_id(query_field.sort_type) if series and query_field.fieldspec.get_field() and query_field.fieldspec.get_field().name.lower() == 'catalognumber': _, _, predicate = query_field.add_to_query(query, formatauditobjs=formatauditobjs) - predicates_by_fieldspec[query_field.fieldspec].append(predicate) if predicate is not None else None + if predicate is not None and query_field.fieldspec not in ordinary_predicates: + predicates_by_fieldspec[query_field.fieldspec].append(predicate) continue query, field, predicate = query_field.add_to_query( @@ -1098,12 +1144,17 @@ def add_fields_to_query( formatted_field = query.objectformatter.fieldformat(query_field, field) query = query.add_columns(formatted_field) selected_fields.append(formatted_field) + if (ranks is not None and ranks.current_rank is not None + and rank_path(query_field.fieldspec) is not None + and path_key(query_field.fieldspec) in ranks.deferred): + ranks.displays[len(query.query.column_descriptions) - 1] = ( + query_field, ranks.current_rank, query.objectformatter) if sort_type is not None: order_by_exprs.append(sort_type(field)) - if predicate is not None: + if predicate is not None and query_field.fieldspec not in ordinary_predicates: predicates_by_fieldspec[query_field.fieldspec].append(predicate) query = apply_where_condition_to_query(query, predicates_by_fieldspec, use_implicit_ors) return query, selected_fields, order_by_exprs @@ -1163,6 +1214,10 @@ def build_query( Return all record IDs associated with a row. search_synonymy = if True, search synonym nodes as well, and return all record IDs associated with parent node + + defer_tree_display = internal execution mode: display-only scalar ranks + remain unresolved until resolve_page() selects the requested page. Direct + callers that execute the returned query must leave this disabled. """ model = models.models_by_tableid[tableid] id_field = model._id @@ -1207,7 +1262,10 @@ def build_query( query_fields=query_fields, series=props.series, formatauditobjs=props.formatauditobjs, - use_implicit_ors=props.implicit_or + use_implicit_ors=props.implicit_or, + defer_tree_display=props.defer_tree_display and not ( + props.distinct or props.series or props.search_synonymy or props.formatauditobjs), + count_only=props.count_only, ) if props.series: @@ -1222,7 +1280,18 @@ def build_query( internal_predicate = query.get_internal_filters() query = query.filter(internal_predicate) - return query.query, order_by_exprs + ranks = query.join_cache.get("tree_ranks") + if ranks is not None and props.defer_tree_display and not ( + props.distinct or props.series or props.search_synonymy or props.formatauditobjs): + ranks.row_ids = [id_field, *[ + alias._id for key, alias in query.join_cache.items() + if isinstance(key, tuple) and len(key) == 2 + and isinstance(key[1], str) + and getattr(key[0], key[1]).property.uselist + ]] + return query.query.execution_options( + specify_tree_ranks=ranks + ), order_by_exprs def series_post_query(query, limit=40, offset=0, sort_type=0, co_id_cat_num_pair_col_index=0, is_count=False): """Transform the query results by removing the co_id:catnum pair column diff --git a/specifyweb/backend/stored_queries/models.py b/specifyweb/backend/stored_queries/models.py index 59d148a5d5c..4b3fb22d1fc 100644 --- a/specifyweb/backend/stored_queries/models.py +++ b/specifyweb/backend/stored_queries/models.py @@ -16,6 +16,16 @@ Session = sessionmaker(bind=engine) +@sqlalchemy.event.listens_for(engine, "before_cursor_execute", retval=True) +def tree_query_recursion_limit(conn, cursor, statement, parameters, context, executemany): + if (context.execution_options.get("specify_tree_ranks") is not None + and getattr(conn.dialect, "is_mariadb", False)): + # MariaDB otherwise returns truncated recursive results after 1,000 + # iterations. Scope the override to this statement, not the connection. + statement = "SET STATEMENT max_recursive_iterations=4294967295 FOR " + statement + return statement, parameters + + def make_session_context(session_maker): @contextmanager def _session_context(): diff --git a/specifyweb/backend/stored_queries/query_construct.py b/specifyweb/backend/stored_queries/query_construct.py index 569e3fc0466..d1934d322fa 100644 --- a/specifyweb/backend/stored_queries/query_construct.py +++ b/specifyweb/backend/stored_queries/query_construct.py @@ -11,12 +11,14 @@ logger = logging.getLogger(__name__) + def _safe_filter(query): count = query.count() if count <= 1: return query.first() raise Exception(f"Got more than one matching: {list(query)}") + class QueryConstruct(namedtuple('QueryConstruct', 'collection objectformatter query join_cache tree_rank_count internal_filters')): def __new__(cls, *args, **kwargs): @@ -28,8 +30,12 @@ def __new__(cls, *args, **kwargs): return super().__new__(cls, *args, **kwargs) def handle_tree_field(self, node, table, tree_rank: TreeRankQuery, next_join_path, current_field_spec: QueryFieldSpec): + """Resolve a rank by indexed parent traversal, not reverse range scans. + + Each relationship/rank gets one outer join, shared by its selected fields. + Ancestors are shared only for identical tree node IDs and definitions. + """ query = self - query_before_tree_joins = query if query.collection is None: # Not sure it makes sense to query across collections raise AssertionError( f"No Collection found in Query for {table}", @@ -37,30 +43,9 @@ def handle_tree_field(self, node, table, tree_rank: TreeRankQuery, next_join_pat ) logger.info('handling treefield %s rank: %s field: %s', table, tree_rank.name, next_join_path) - treedefitem_column = table.name + 'TreeDefItemID' treedef_column = table.name + 'TreeDefID' - cache_key = (table, 'TreeRanks') - if cache_key in query.join_cache: - logger.debug("using join cache for %r tree ranks.", node) - ancestors, treedefs = query.join_cache[cache_key] - else: - treedefs = get_treedefs(query.collection, table.name) - - # We need to take the max here. Otherwise, it is possible that the same rank - # name may not occur at the same level across tree defs. - max_depth = max(depth for _, depth in treedefs) - - ancestors = [node] - for _ in range(max_depth - 1): - ancestor = orm.aliased(node) - query = query.outerjoin(ancestor, ancestors[-1].ParentID == ancestor._id) - ancestors.append(ancestor) - - logger.debug("adding to join cache for %r tree ranks.", node) - query = query._replace(join_cache=query.join_cache.copy()) - query.join_cache[cache_key] = (ancestors, treedefs) - + treedefs = get_treedefs(query.collection, table.name) item_model = getattr(spmodels, table.django_name + "treedefitem") # TODO: optimize out the ranks that appear? cache them @@ -77,29 +62,28 @@ def handle_tree_field(self, node, table, tree_rank: TreeRankQuery, next_join_pat tree_rank.name, table.name, ) - return query_before_tree_joins, None, None, table - - treedefitem_params = [treedefitem_id for (_, treedefitem_id) in treedefs_with_ranks] - - def make_tree_field_spec(tree_node): - return current_field_spec._replace( - root_table=table, # rebasing the query - root_sql_table=tree_node, # this is needed to preserve SQL aliased going to next part - join_path=next_join_path, # slicing join path to begin from after the tree - ) - - cases = [] - field = None # just to stop mypy from complaining. - for ancestor in ancestors: - field_spec = make_tree_field_spec(ancestor) - query, orm_field, field, table = field_spec.add_spec_to_query(query) - # Field and table won't matter. Rank acts as fork, and these two will be same across siblings - for treedefitem_param in treedefitem_params: - cases.append((getattr(ancestor, treedefitem_column) == treedefitem_param, orm_field)) - - column = sql.case(cases) + return query, None, None, table + treedefitem_params = tuple(sorted( + treedefitem_id for (_, treedefitem_id) in treedefs_with_ranks + )) defs_to_filter_on = [def_id for (def_id, _) in treedefs_with_ranks] + cache_key = (node, table, treedefitem_params, 'TreeRank') + ranks = query.join_cache.get("tree_ranks") + if ranks is None: + from .tree_rank import TreeRanks + + ranks = TreeRanks() + query.join_cache["tree_ranks"] = ranks + ranks.current_rank = (node, table, treedefitem_params) + if cache_key in query.join_cache: + ancestor = query.join_cache[cache_key] + else: + ancestor = orm.aliased(getattr(models, table.name)) + query = ranks.join(query, node, ancestor, table, treedefitem_params) + query = query._replace(join_cache=query.join_cache.copy()) + query.join_cache[cache_key] = ancestor + # We don't want to include treedef if the rank is not present. new_filters = [ *query.internal_filters, @@ -110,6 +94,12 @@ def make_tree_field_spec(tree_node): ] query = query._replace(internal_filters=new_filters) + field_spec = current_field_spec._replace( + root_table=table, + root_sql_table=ancestor, + join_path=next_join_path, + ) + query, column, field, table = field_spec.add_spec_to_query(query) return query, column, field, table def tables_in_path(self, table, join_path): diff --git a/specifyweb/backend/stored_queries/tests/test_build_query.py b/specifyweb/backend/stored_queries/tests/test_build_query.py index 38d5449ffe0..f8c3ac6bee7 100644 --- a/specifyweb/backend/stored_queries/tests/test_build_query.py +++ b/specifyweb/backend/stored_queries/tests/test_build_query.py @@ -1,7 +1,36 @@ -from specifyweb.specify.models import Taxontreedefitem, datamodel +import re +import warnings +from types import SimpleNamespace + +from django.db import connection +from django.test.utils import CaptureQueriesContext +from sqlalchemy.dialects import mysql +from sqlalchemy.exc import SAWarning +from sqlalchemy.sql.selectable import Join + +from specifyweb.specify.models import ( + Collectingeventattribute, + Collectionobjecttype, + Determination, + Geography, + Geographytreedefitem, + Locality, + Taxon, + Taxontreedef, + Taxontreedefitem, + datamodel, +) from specifyweb.backend.stored_queries import models as models from specifyweb.backend.stored_queries.queryfield import fields_from_json -from specifyweb.backend.stored_queries.execution import QuerySort, BuildQueryProps, build_query_construct_base, add_fields_to_query +from specifyweb.backend.stored_queries.queryfieldspec import TreeRankQuery +from specifyweb.backend.stored_queries.tree_rank import resolve_page +from specifyweb.backend.stored_queries.execution import ( + QuerySort, + BuildQueryProps, + build_query, + build_query_construct_base, + add_fields_to_query, +) from specifyweb.backend.stored_queries.tests.tests import SQLAlchemySetup class TestBuildQuery(SQLAlchemySetup): @@ -100,30 +129,805 @@ def test_no_extra_tree_joins(self): series=props.series, formatauditobjs=props.formatauditobjs ) - # There should be 4 objects in the join cache - # CollectionObject -> Determinations - # Determination -> taxon - # Taxon Ranks for Determination -> taxon - # Determination -> preferredTaxon - self.assertEqual( - len(query.join_cache), - 4 + # Three relationship joins and at most one join per requested rank. + self._assert_join_budget(query.query, 5) + + def _assert_join_budget(self, query, max_joins): + statement = str(query.statement.compile(dialect=mysql.dialect())) + join_count = len(re.findall(r"\bJOIN\b", statement, flags=re.IGNORECASE)) + def joins(source): + return 1 + joins(source.left) + joins(source.right) if isinstance(source, Join) else 0 + final_joins = sum(joins(source) for source in query.statement.get_final_froms()) + self.assertLessEqual(final_joins, max_joins, statement) + # At most six relationship joins, one unpivot and two indexed traversal + # joins in the shared seed CTE. This allowance never grows with depth. + self.assertLessEqual( + join_count, + max_joins + 9, + f"Generated {join_count} JOINs including seed traversal:\n{statement}", + ) + + def test_direct_rank_fields_resolve_without_preparation(self): + self._make_distinct_taxon_relationships() + fields = [ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Species", "Preferred Species"), + ] + for missing_preferred in (False, True): + if missing_preferred: + Determination.objects.filter( + collectionobject=self.collectionobjects[0] + ).update(preferredtaxon=None) + for ordered_fields in (fields, list(reversed(fields))): + with self.subTest( + missing_preferred=missing_preferred, + first_field=ordered_fields[0][0], + ), self.__class__.test_session_context() as session: + query = build_query_construct_base( + session, self.collection, self.specifyuser, + models.CollectionObject, BuildQueryProps(), + ) + self.assertNotIn("tree_ranks", query.join_cache) + query_fields = self._taxon_relationship_query_fields(ordered_fields) + for query_field in query_fields: + query, column, predicate = query_field.add_to_query( + query, collection=self.collection, user=self.specifyuser + ) + query = query.add_columns(column) + if predicate is not None: + query = query.filter(predicate) + query = query.filter(query.get_internal_filters()) + self.assertEqual(len(query.query.statement.get_final_froms()), 1) + with warnings.catch_warnings(): + warnings.simplefilter("error", SAWarning) + rows = list(query.query) + self.assertEqual( + [tuple(row) for row in rows], + [( + self.collectionobjects[0].id, "num-0", + *(None if missing_preferred and "preferredTaxon" in path + else value for path, value in ordered_fields), + )], + ) + + def _make_distinct_taxon_relationships(self, depth=7): + phylum_rank = self.taxontreedef.treedefitems.get(name="Phylum") + order_rank = Taxontreedefitem.objects.create( + rankid=100, + parent=phylum_rank, + treedef=self.taxontreedef, + name="Order", + title="Order", + ) + # Build deep fixtures without invoking the unrelated full-name join generator. + intermediate_ranks = Taxontreedefitem.objects.bulk_create([ + Taxontreedefitem( + rankid=101 + index, + treedef=self.taxontreedef, + name=f"Intermediate {index}", + title=f"Intermediate {index}", + ) + for index in range(depth - 7) + ]) + parent_rank = order_rank + for rank in intermediate_ranks: + rank.parent = parent_rank + parent_rank = rank + family_rank = self.taxontreedef.treedefitems.get(name="Family") + family_rank.parent = parent_rank + Taxontreedefitem.objects.bulk_update([*intermediate_ranks, family_rank], ["parent"]) + ranks = list(self.taxontreedef.treedefitems.order_by("rankid")) + self.assertEqual(len(ranks), depth) + + species = [] + for relationship_index, relationship in enumerate(("Taxon", "Preferred", "Host")): + parent = None + for rank_index, rank in enumerate(ranks): + parent = Taxon( + name=f"{relationship} {rank.name}", + fullname=f"{relationship} {rank.name}", + parent=parent, + definition=self.taxontreedef, + definitionitem=rank, + rankid=rank.rankid, + nodenumber=relationship_index * depth + rank_index + 1, + highestchildnodenumber=(relationship_index + 1) * depth, + ) + parent.save(skip_tree_extras=True) + species.append(parent) + + self._update(species[0], {"acceptedtaxon": species[1]}) + determination = Determination.objects.create( + collectionobject=self.collectionobjects[0], + taxon=species[0], + preferredtaxon=species[1], + iscurrent=True, + ) + determination.refresh_from_db() + self.assertEqual(determination.taxon.name, "Taxon Species") + self.assertEqual(determination.preferredtaxon.name, "Preferred Species") + attribute = Collectingeventattribute.objects.create( + discipline=self.discipline, + hosttaxon=species[2], + ) + self._update( + self.collectingevent, {"collectingeventattribute": attribute} + ) + self._update( + self.collectionobjects[0], {"collectingevent": self.collectingevent} + ) + + def _taxon_relationship_query_fields(self, fields): + return fields_from_json([ + { + "position": position, + "stringid": stringid, + "formatname": None, + "isdisplay": True, + "isnot": False, + "isrelfld": False, + "operstart": 1 if position == 0 else 8, + "sorttype": QuerySort.NONE, + "startvalue": "num-0" if position == 0 else "", + "isstrict": False, + } + for position, (stringid, _) in enumerate([ + ("1.collectionobject.catalogNumber", "num-0"), + *fields, + ]) + ]) + + def _assert_taxon_relationship_values(self, fields, max_joins=None): + query_fields = self._taxon_relationship_query_fields(fields) + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, + self.collection, + self.specifyuser, + datamodel.get_table_strict("CollectionObject").tableId, + query_fields, + ) + # Check before execution so deep-tree regressions cannot hit MySQL's table limit. + if max_joins is not None: + self._assert_join_budget(query, max_joins) + results = list(query) + self.assertEqual( + [tuple(row) for row in results], + [( + self.collectionobjects[0].id, + "num-0", + *(expected for _, expected in fields), + )], + ) + + def _make_second_taxon_tree(self): + treedef = Taxontreedef.objects.create(name="Second taxonomy") + Collectionobjecttype.objects.create( + name="Second type", collection=self.collection, taxontreedef=treedef + ) + parent = None + # Deliberately reuse the first tree's node numbers and rank names. + for rank_index, (name, rankid) in enumerate( + (("Root", 0), ("Genus", 180), ("Species", 220)) + ): + rank = Taxontreedefitem.objects.create( + name=name, + title=name, + rankid=rankid, + treedef=treedef, + parent=parent.definitionitem if parent else None, + ) + parent = Taxon( + name=f"Second {name}", + fullname=f"Second {name}", + parent=parent, + definition=treedef, + definitionitem=rank, + rankid=rankid, + nodenumber=rank_index + 1, + highestchildnodenumber=7, ) - # BUG: This is technically undesirable, as it causes #8650 - # In the underlying query, the Preferred Taxon currently uses the - # JOIN for Determination -> taxon. Specifically, it uses the cached - # JOINs for the tree ranks. - taxon_table = datamodel.get_table_strict("taxon") - cache_key = (taxon_table, 'TreeRanks') - tree_ranks_in_cache = list(filter(lambda cache_key: 'TreeRanks' in cache_key, query.join_cache.keys())) + parent.save(skip_tree_extras=True) + return treedef, parent + + def test_rank_joins_isolate_tree_definitions_with_overlapping_numbers(self): + self._make_distinct_taxon_relationships() + _treedef, second_species = self._make_second_taxon_tree() + Determination.objects.filter( + collectionobject=self.collectionobjects[0] + ).update(preferredtaxon=second_species) + fields = [ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Genus", "Second Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Species", "Second Species"), + ] + for ordered_fields in (fields, list(reversed(fields))): + with self.subTest(first_field=ordered_fields[0][0]): + self._assert_taxon_relationship_values(ordered_fields, max_joins=6) + + def test_explicit_tree_definition_ranks_do_not_share_aliases(self): + self._make_distinct_taxon_relationships() + second_treedef, _second_species = self._make_second_taxon_tree() + fields = [("1,9-determinations,4.taxon.Genus", "Taxon Genus")] + query_fields = self._taxon_relationship_query_fields(fields) + genus_field = query_fields[-1] + for definitions in ( + (self.taxontreedef, second_treedef), + (second_treedef, self.taxontreedef), + ): + with self.subTest(first_definition=definitions[0].name): + restricted_fields = [] + for treedef in definitions: + spec = genus_field.fieldspec + restricted_rank = TreeRankQuery.create( + "Genus", "Taxon", treedef.id, treedef.name + ) + restricted_fields.append(genus_field._replace( + fieldspec=spec._replace(join_path=tuple( + restricted_rank if isinstance(part, TreeRankQuery) else part + for part in spec.join_path + )) + )) + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, self.collection, self.specifyuser, + datamodel.get_table_strict("CollectionObject").tableId, + [query_fields[0], *restricted_fields], + ) + self._assert_join_budget(query, 4) + self.assertEqual( + [tuple(row) for row in query], + [( + self.collectionobjects[0].id, "num-0", + *("Taxon Genus" if treedef == self.taxontreedef else None + for treedef in definitions), + )], + ) + + def test_skipped_rank_keeps_null_without_matching_another_branch(self): + self._make_distinct_taxon_relationships() + species = Taxon.objects.get(name="Taxon Species") + Taxon.objects.filter(id=species.id).update(parent=species.parent.parent) + # The skipped genus is a sibling of this species, not its ancestor. + Taxon.objects.filter(id=species.parent_id).update( + highestchildnodenumber=species.parent.nodenumber + ) + self._assert_taxon_relationship_values([ + ("1,9-determinations,4.taxon.Family", "Taxon Family"), + ("1,9-determinations,4.taxon.Genus", None), + ("1,9-determinations,4.taxon.Species", "Taxon Species"), + ("1,9-determinations,4-preferredTaxon.taxon.Genus", "Preferred Genus"), + ], max_joins=7) + + def test_geography_rank_fields_preserve_values(self): + root_rank = self.geographytreedef.treedefitems.get(rankid=0) + country_rank = Geographytreedefitem.objects.create( + name="Country", rankid=200, treedef=self.geographytreedef, parent=root_rank + ) + state_rank = Geographytreedefitem.objects.create( + name="State", rankid=300, treedef=self.geographytreedef, parent=country_rank + ) + parent = None + for rank, name in ( + (root_rank, "Earth"), (country_rank, "Country A"), (state_rank, "State A") + ): + parent = Geography.objects.create( + name=name, fullname=name, definition=self.geographytreedef, + definitionitem=rank, rankid=rank.rankid, parent=parent, + ) + locality = Locality.objects.create( + localityname="Locality A", discipline=self.discipline, geography=parent + ) + self._update(self.collectingevent, {"locality": locality}) + self._update(self.collectionobjects[0], {"collectingevent": self.collectingevent}) + self._assert_taxon_relationship_values([ + ("1,10,2,3.geography.Country", "Country A"), + ("1,10,2,3.geography.State", "State A"), + ], max_joins=5) + + def test_missing_rank_omits_filter_and_sort_but_preserves_column(self): + self._make_distinct_taxon_relationships() + fields = [ + ("1,9-determinations,4.taxon.Unmapped rank", None), + ("1,9-determinations,4-preferredTaxon.taxon.Species", "Preferred Species"), + ] + query_fields = self._taxon_relationship_query_fields(fields) + query_fields[1] = query_fields[1]._replace( + op_num=1, value="No match", sort_type=QuerySort.ASC + ) + with self.assertLogs( + "specifyweb.backend.stored_queries.query_construct", level="WARNING" + ): + with self.__class__.test_session_context() as session: + query, order_by = build_query( + session, self.collection, self.specifyuser, + datamodel.get_table_strict("CollectionObject").tableId, + query_fields, + ) + self._assert_join_budget(query, 4) + self.assertEqual(order_by, []) + self.assertEqual( + [tuple(row) for row in query], + [(self.collectionobjects[0].id, "num-0", None, "Preferred Species")], + ) + + def test_rank_filters_and_sorting_use_the_selected_relationship(self): + self._make_distinct_taxon_relationships() + taxon_species = Taxon.objects.get(name="Taxon Species") + host_species = Taxon.objects.get(name="Host Species") + Determination.objects.create( + collectionobject=self.collectionobjects[1], + taxon=taxon_species, + preferredtaxon=host_species, + iscurrent=True, + ) + Determination.objects.filter( + collectionobject=self.collectionobjects[1] + ).update(preferredtaxon=host_species) + fields = [ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Genus", "Preferred Genus"), + ] + base_fields = self._taxon_relationship_query_fields(fields) + base_fields[0] = base_fields[0]._replace(op_num=10, value="num-0,num-1") + cases = ( + ("equals", 1, "Preferred Genus", False, QuerySort.ASC, [0]), + ("not-equals", 1, "Preferred Genus", True, QuerySort.ASC, [1]), + ("contains", 11, "Genus", False, QuerySort.ASC, [1, 0]), + ("descending", 8, "", False, QuerySort.DESC, [0, 1]), + ) + for name, op_num, value, negate, sort_type, indexes in cases: + with self.subTest(operation=name): + query_fields = [ + *base_fields[:2], + base_fields[2]._replace( + op_num=op_num, value=value, negate=negate, sort_type=sort_type + ), + ] + with self.__class__.test_session_context() as session: + query, order_by = build_query( + session, self.collection, self.specifyuser, + datamodel.get_table_strict("CollectionObject").tableId, + query_fields, + ) + self._assert_join_budget(query, 5) + self.assertEqual(len(order_by), 1) + self.assertEqual( + [tuple(row) for row in query.order_by(*order_by)], + [( + self.collectionobjects[index].id, f"num-{index}", + "Taxon Genus", "Preferred Genus" if index == 0 else "Host Genus", + ) for index in indexes], + ) + + def test_negated_rank_filter_includes_missing_relationships(self): + self._make_distinct_taxon_relationships() + taxon_species = Taxon.objects.get(name="Taxon Species") + Determination.objects.create( + collectionobject=self.collectionobjects[1], + taxon=taxon_species, + iscurrent=True, + ) + Determination.objects.filter( + collectionobject=self.collectionobjects[1] + ).update(preferredtaxon=None) + query_fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4-preferredTaxon.taxon.Genus", None), + ]) + query_fields[0] = query_fields[0]._replace(op_num=10, value="num-0,num-1") + for op_num, value, negate in ( + (1, "Preferred Genus", True), + (12, "", False), + ): + with self.subTest(op_num=op_num, negate=negate): + fields = [ + query_fields[0], + query_fields[1]._replace(op_num=op_num, value=value, negate=negate), + ] + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, self.collection, self.specifyuser, + datamodel.get_table_strict("CollectionObject").tableId, + fields, + ) + self._assert_join_budget(query, 3) + self.assertEqual( + [tuple(row) for row in query], + [(self.collectionobjects[1].id, "num-1", None)], + ) + + def test_rank_query_construction_does_not_read_tree_nodes(self): + self._make_distinct_taxon_relationships() + query_fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4.taxon.Species", "Taxon Species"), + ("1,9-determinations,4-preferredTaxon.taxon.Genus", "Preferred Genus"), + ]) + with self.__class__.test_session_context() as session: + with CaptureQueriesContext(connection) as captured: + query, _ = build_query( + session, self.collection, self.specifyuser, + datamodel.get_table_strict("CollectionObject").tableId, + query_fields, + ) + self._assert_join_budget(query, 6) + for statement in captured.captured_queries: + self.assertNotRegex( + statement["sql"], + r"(?i)\b(?:FROM|JOIN)\s+`?taxon`?(?=\s|$)", + ) self.assertEqual( - tree_ranks_in_cache, - [cache_key] + [tuple(row) for row in query], + [(self.collectionobjects[0].id, "num-0", + "Taxon Genus", "Taxon Species", "Preferred Genus")], ) - tree_join_information = query.join_cache[cache_key] - tree_ranks = tree_join_information[0] + def test_explicit_tree_definition_ignores_other_trees(self): + self._make_distinct_taxon_relationships() + second_treedef, second_species = self._make_second_taxon_tree() + Taxon.objects.filter(id=second_species.id).update(nodenumber=None) + query_fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ]) + spec = query_fields[-1].fieldspec + restricted_rank = TreeRankQuery.create( + "Genus", "Taxon", self.taxontreedef.id, self.taxontreedef.name + ) + query_fields[-1] = query_fields[-1]._replace(fieldspec=spec._replace( + join_path=tuple( + restricted_rank if isinstance(part, TreeRankQuery) else part + for part in spec.join_path + ) + )) + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, self.collection, self.specifyuser, + datamodel.get_table_strict("CollectionObject").tableId, + query_fields, + ) self.assertEqual( - len(tree_ranks), - len(self.taxontreedef.treedefitems.all()) + [tuple(row) for row in query], + [(self.collectionobjects[0].id, "num-0", "Taxon Genus")], ) + self.assertTrue(Taxon.objects.filter( + id=second_species.id, definition=second_treedef, nodenumber__isnull=True + ).exists()) + + def test_explicit_ranks_share_the_same_relationship(self): + self._make_distinct_taxon_relationships() + self._assert_taxon_relationship_values([ + ("1,9-determinations,4.taxon.Order", "Taxon Order"), + ("1,9-determinations,4.taxon.Species", "Taxon Species"), + ]) + + def test_rank_fields_share_one_join(self): + self._make_distinct_taxon_relationships() + genus = Taxon.objects.get(name="Taxon Genus") + Taxon.objects.filter(id=genus.id).update(author="Genus author") + self._assert_taxon_relationship_values([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4.taxon.Genus Author", "Genus author"), + ("1,9-determinations,4.taxon.Genus ID", genus.id), + ], max_joins=3) + + def test_standalone_root_rank_without_node_numbers(self): + root_rank = self.taxontreedef.treedefitems.get(name="Root") + root = Taxon.objects.create( + name="Standalone root", + fullname="Standalone root", + definition=self.taxontreedef, + definitionitem=root_rank, + rankid=root_rank.rankid, + ) + self.assertIsNone(root.nodenumber) + self.assertIsNone(root.highestchildnodenumber) + Determination.objects.create( + collectionobject=self.collectionobjects[0], + taxon=root, + preferredtaxon=root, + iscurrent=True, + ) + self._assert_taxon_relationship_values([ + ("1,9-determinations,4.taxon.Root", root.name), + ("1,9-determinations,4.taxon.Root ID", root.id), + ("1,9-determinations,4.taxon.Genus", None), + ], max_joins=4) + + def test_missing_taxon_relationship_keeps_null_rank_values(self): + self._make_distinct_taxon_relationships() + Determination.objects.filter( + collectionobject=self.collectionobjects[0] + ).update(preferredtaxon=None) + self._assert_taxon_relationship_values([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Genus", None), + ("1,9-determinations,4-preferredTaxon.taxon.Species", None), + ], max_joins=6) + + def test_any_rank_values_use_distinct_relationships(self): + self._make_distinct_taxon_relationships() + self._assert_taxon_relationship_values([ + ("1,9-determinations,4.taxon.name", "Taxon Species"), + ( + "1,9-determinations,4-preferredTaxon.taxon.name", + "Preferred Species", + ), + ("1,10,92,4-hostTaxon.taxon.name", "Host Species"), + ]) + + def test_explicit_ranks_use_distinct_relationships_in_either_field_order(self): + self._make_distinct_taxon_relationships() + taxon_fields = [ + ("1,9-determinations,4.taxon.Order", "Taxon Order"), + ("1,9-determinations,4.taxon.Species", "Taxon Species"), + ] + preferred_fields = [ + ( + "1,9-determinations,4-preferredTaxon.taxon.Species", + "Preferred Species", + ), + ( + "1,9-determinations,4-preferredTaxon.taxon.Order", + "Preferred Order", + ), + ] + for first, second in ( + (taxon_fields, preferred_fields), + (preferred_fields, taxon_fields), + ): + with self.subTest(first_relationship=first[0][0]): + self._assert_taxon_relationship_values([*first, *second]) + + def test_explicit_ranks_use_three_distinct_relationships(self): + self._make_distinct_taxon_relationships() + fields = [ + ("1,9-determinations,4.taxon.Order", "Taxon Order"), + ( + "1,9-determinations,4-preferredTaxon.taxon.Species", + "Preferred Species", + ), + ("1,10,92,4-hostTaxon.taxon.Species", "Host Species"), + ] + for ordered_fields in (fields, list(reversed(fields))): + with self.subTest(first_relationship=ordered_fields[0][0]): + self._assert_taxon_relationship_values(ordered_fields) + + def _assert_genus_species_relationships( + self, depth, include_host=False, max_joins=None + ): + self._make_distinct_taxon_relationships(depth=depth) + relationships = [ + ("1,9-determinations,4", "Taxon"), + ("1,9-determinations,4-preferredTaxon", "Preferred"), + ] + if include_host: + relationships.append(("1,10,92,4-hostTaxon", "Host")) + fields = [ + (f"{path}.taxon.{rank}", f"{relationship} {rank}") + for path, relationship in relationships + for rank in ("Genus", "Species") + ] + for ordering, ordered_fields in ( + ("taxon-first", fields), + ("reversed", list(reversed(fields))), + ("repeated-ranks", [*fields, fields[0], fields[-1]]), + ): + with self.subTest( + depth=depth, ordering=ordering, relationships=len(relationships) + ): + self._assert_taxon_relationship_values( + ordered_fields, max_joins=max_joins + ) + + def test_genus_species_values_use_distinct_relationships(self): + self._assert_genus_species_relationships(depth=11) + + def test_genus_species_join_budget(self): + # Three relationship joins plus two requested ranks for each relationship. + self._assert_genus_species_relationships(depth=11, max_joins=7) + + def test_genus_species_join_budget_with_sixteen_ranks(self): + self._assert_genus_species_relationships(depth=16, max_joins=7) + + def test_three_relationship_join_budget_with_sixteen_ranks(self): + # Six relationship joins plus two requested ranks for each relationship. + self._assert_genus_species_relationships( + depth=16, include_host=True, max_joins=12 + ) + + def test_three_relationship_join_budget_with_thirty_two_ranks(self): + self._assert_genus_species_relationships( + depth=32, include_host=True, max_joins=12 + ) + + def test_three_relationship_join_budget_with_sixty_four_ranks(self): + self._assert_genus_species_relationships( + depth=64, include_host=True, max_joins=12 + ) + + def test_page_rank_resolution_preserves_relationships_and_duplicates(self): + self._make_distinct_taxon_relationships(depth=32) + determination = Determination.objects.get(collectionobject=self.collectionobjects[0]) + Determination.objects.bulk_create([Determination( + collectionobject=self.collectionobjects[0], + collectionmemberid=determination.collectionmemberid, + taxon=determination.taxon, + preferredtaxon=determination.preferredtaxon, + iscurrent=False, + )]) + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Genus", "Preferred Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Species", "Preferred Species"), + ]) + for offset in (0, 1, 2): + with self.subTest(offset=offset), self.__class__.test_session_context() as session: + query, order = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(defer_tree_display=True)) + ranks = query.get_execution_options()["specify_tree_ranks"] + with warnings.catch_warnings(): + warnings.simplefilter("error", SAWarning) + page = resolve_page( + query, [*order, models.CollectionObject._id], ranks, 1, offset) + if offset == 2: + # The SQLAlchemy test adapter cannot wrap an empty result. + self.assertEqual(page.count(), 0) + continue + rows = list(page) + expected = [(self.collectionobjects[0].id, "num-0", "Taxon Genus", + "Preferred Genus", "Preferred Species")] if offset < 2 else [] + self.assertEqual([tuple(row) for row in rows], expected) + + def test_sorted_page_resolves_only_required_relationship_before_pagination(self): + self._make_distinct_taxon_relationships() + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Genus", "Preferred Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Species", "Preferred Species"), + ]) + fields[2] = fields[2]._replace(sort_type=QuerySort.ASC) + with self.__class__.test_session_context() as session: + query, order = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(defer_tree_display=True)) + ranks = query.get_execution_options()["specify_tree_ranks"] + self.assertEqual(set(ranks.displays), {2, 4}) + with warnings.catch_warnings(): + warnings.simplefilter("error", SAWarning) + rows = list(resolve_page(query, order, ranks, 20, 0)) + self.assertEqual([tuple(row) for row in rows], [( + self.collectionobjects[0].id, "num-0", "Taxon Genus", + "Preferred Genus", "Preferred Species")]) + + def test_display_ranks_do_not_require_ancestor_mapping_for_counts(self): + self._make_distinct_taxon_relationships() + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Species", "Preferred Species"), + ]) + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(defer_tree_display=True)) + self.assertNotIn("WITH RECURSIVE", str(query.statement)) + self.assertEqual(query.count(), 1) + + def test_grouped_rank_queries_use_eager_resolution(self): + self._make_distinct_taxon_relationships() + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Species", "Preferred Species"), + ]) + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(distinct=True, defer_tree_display=True)) + self.assertFalse(query.get_execution_options()["specify_tree_ranks"].displays) + rows = list(query) + self.assertEqual(tuple(rows[0])[1:], ("num-0", "Taxon Genus", "Preferred Species")) + + def test_count_does_not_resolve_display_only_sort_ranks(self): + self._make_distinct_taxon_relationships() + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ]) + fields[1] = fields[1]._replace(sort_type=QuerySort.ASC) + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(defer_tree_display=True, count_only=True)) + self.assertNotIn("WITH RECURSIVE", str(query.statement)) + self.assertEqual(query.count(), 1) + + def test_page_resolution_preserves_null_preferred_taxon(self): + self._make_distinct_taxon_relationships() + Determination.objects.filter( + collectionobject=self.collectionobjects[0]).update(preferredtaxon=None) + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4-preferredTaxon.taxon.Species", None), + ]) + with self.__class__.test_session_context() as session: + query, order = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(defer_tree_display=True)) + page = resolve_page(query, order, query.get_execution_options()["specify_tree_ranks"], 20, 0) + self.assertEqual([tuple(row) for row in page], [( + self.collectionobjects[0].id, "num-0", "Taxon Genus", None)]) + + def test_mixed_tree_tables_use_eager_resolution(self): + self._make_distinct_taxon_relationships() + rank = self.geographytreedef.treedefitems.get(rankid=0) + geography = Geography.objects.create( + name="World", fullname="World", definition=self.geographytreedef, + definitionitem=rank, rankid=rank.rankid) + locality = Locality.objects.create( + localityname="Here", discipline=self.discipline, geography=geography) + self._update(self.collectingevent, {"locality": locality}) + self._update(self.collectionobjects[0], {"collectingevent": self.collectingevent}) + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + (f"1,10,2,3.geography.{rank.name}", "World"), + ]) + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(defer_tree_display=True)) + self.assertFalse(query.get_execution_options()["specify_tree_ranks"].displays) + self.assertEqual([tuple(row) for row in query], [( + self.collectionobjects[0].id, "num-0", "Taxon Genus", "World")]) + + def test_recursion_override_is_statement_scoped(self): + statement = "WITH RECURSIVE ancestors AS (...) SELECT * FROM ancestors" + parameters = (1,) + context = SimpleNamespace(execution_options={"specify_tree_ranks": object()}) + mariadb = SimpleNamespace(dialect=SimpleNamespace(is_mariadb=True)) + mysql_conn = SimpleNamespace(dialect=SimpleNamespace(is_mariadb=False)) + changed, unchanged = models.tree_query_recursion_limit( + mariadb, None, statement, parameters, context, False) + self.assertEqual( + changed, "SET STATEMENT max_recursive_iterations=4294967295 FOR " + statement) + self.assertIs(unchanged, parameters) + self.assertEqual(models.tree_query_recursion_limit( + mysql_conn, None, statement, parameters, context, False), (statement, parameters)) + context.execution_options = {} + self.assertEqual(models.tree_query_recursion_limit( + mariadb, None, statement, parameters, context, False), (statement, parameters)) + + def test_tree_seed_preserves_implicit_or_between_ordinary_filters(self): + self._make_distinct_taxon_relationships() + determination = Determination.objects.get(collectionobject=self.collectionobjects[0]) + Determination.objects.create( + collectionobject=self.collectionobjects[1], taxon=determination.taxon, + preferredtaxon=determination.preferredtaxon, iscurrent=True) + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ]) + fields.insert(1, fields[0]._replace(value="num-1", display=False)) + for deferred in (False, True): + with self.subTest(deferred=deferred), self.__class__.test_session_context() as session: + query, order = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(defer_tree_display=deferred)) + ranks = query.get_execution_options()["specify_tree_ranks"] + result = resolve_page(query, order, ranks, 20, 0) if deferred else query + self.assertCountEqual([tuple(row) for row in result], [ + (self.collectionobjects[index].id, f"num-{index}", "Taxon Genus") + for index in (0, 1)]) + + def test_temporal_rank_field_keeps_its_shared_rank_eager(self): + self._make_distinct_taxon_relationships() + fields = self._taxon_relationship_query_fields([ + ("1,9-determinations,4.taxon.Genus", "Taxon Genus"), + ("1,9-determinations,4.taxon.Genus timestampCreatedNumericYear", None), + ]) + results = [] + for deferred in (False, True): + with self.__class__.test_session_context() as session: + query, _ = build_query( + session, self.collection, self.specifyuser, 1, fields, + BuildQueryProps(defer_tree_display=deferred)) + self.assertFalse(query.get_execution_options()["specify_tree_ranks"].displays) + results.append([tuple(row) for row in query]) + self.assertEqual(results[0], results[1]) diff --git a/specifyweb/backend/stored_queries/tree_rank.py b/specifyweb/backend/stored_queries/tree_rank.py new file mode 100644 index 00000000000..b6e202a4af1 --- /dev/null +++ b/specifyweb/backend/stored_queries/tree_rank.py @@ -0,0 +1,252 @@ +"""Query-scoped, indexed ancestor resolution for tree rank fields. + +Ordinary predicates bound the seed set before rank filters are evaluated. Each +tree shares a recursive parent mapping across its relationship aliases. Counts +can omit display-only ranks; scalar pages resolve them after selecting rows. +Grouping, synonym searches, mixed trees and formatted objects remain eager. +Mappings live only in the SQL statement: there is no tree cache to invalidate. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Protocol + +from sqlalchemy import orm, sql + +from . import models + +if TYPE_CHECKING: + from specifyweb.specify.models_utils.load_datamodel import Table + ObjectFormatter = object + Rank = tuple[object, Table, tuple[int, ...]] + + +class RankField(Protocol): + @property + def fieldspec(self) -> object: + raise NotImplementedError + + +def rank_path(spec): + from specifyweb.specify.datamodel import datamodel, is_tree_table + for index, part in enumerate(spec.join_path): + if ( + hasattr(part, "relatedModelName") + and hasattr(part, "name") + and hasattr(part, "treedef_id") + and is_tree_table(datamodel.get_table_strict(part.relatedModelName)) + ): + return spec.join_path[:index], part, spec.join_path[index + 1:] + return None + + +def path_key(spec): + path, rank, _ = rank_path(spec) + return (spec.root_sql_table, tuple(part.name for part in path), + rank.relatedModelName, rank.name, rank.treedef_id) + + +def ancestor_mapping(seeds, table): + model = getattr(models, table.name) + node = orm.aliased(model) + definition = table.name + "TreeDefID" + item = table.name + "TreeDefItemID" + mapping = sql.select( + node._id.label("seed_id"), node._id.label("ancestor_id"), + node.ParentID.label("parent_id"), + getattr(node, item).label("item_id"), + getattr(node, definition).label("definition_id"), + ).select_from(seeds).join(node, node._id == seeds.c.node_id).cte(recursive=True) + parent = orm.aliased(model) + # UNION terminates even if a damaged tree contains a parent cycle. + return mapping.union(sql.select( + mapping.c.seed_id, parent._id, parent.ParentID, + getattr(parent, item), getattr(parent, definition), + ).select_from(mapping).join( + parent, parent._id == mapping.c.parent_id, + ).where(getattr(parent, definition) == mapping.c.definition_id)) + + +def ancestor_id(mapping, node_id, item_ids, definition_id, distinct=False): + ancestors = mapping.alias() + selected = sql.select(ancestors.c.ancestor_id) + if distinct: + selected = selected.distinct() + return selected.where( + ancestors.c.seed_id == node_id, + ancestors.c.item_id.in_(item_ids), + ancestors.c.definition_id == definition_id, + ).correlate_except(ancestors).scalar_subquery() + + +@dataclass +class TreeRanks: + deferred: set = field(default_factory=set) + mappings: dict = field(default_factory=dict) + displays: dict[int, tuple[RankField, Rank, ObjectFormatter]] = field(default_factory=dict) + current_field: RankField | None = None + current_rank: Rank | None = None + mapped_nodes: set = field(default_factory=set) + row_ids: list = field(default_factory=list) + + def join(self, query, node, ancestor, table, item_ids): + deferred = ( + self.current_field is not None + and rank_path(self.current_field.fieldspec) is not None + and path_key(self.current_field.fieldspec) in self.deferred + ) + self.current_rank = (node, table, item_ids) + if deferred: + return query.outerjoin(ancestor, ancestor._id == node._id) + key = table if node in self.mapped_nodes else (table, node) + if key not in self.mappings: + seeds = query.query.with_entities(node._id.label("node_id")).distinct().subquery() + self.mappings[key] = ancestor_mapping(seeds, table) + mapping = self.mappings[key] + definition = table.name + "TreeDefID" + return query.outerjoin(ancestor, sql.and_( + ancestor._id == ancestor_id( + mapping, node._id, item_ids, getattr(node, definition)), + getattr(ancestor, definition) == getattr(node, definition), + )) + + +def prepare_tree_ranks(query, query_fields, defer_display, count_only=False): + """Seed ancestors from ordinary filters, before adding any rank predicates.""" + ranks = TreeRanks() + tree_fields = [f for f in query_fields if rank_path(f.fieldspec) is not None] + required = { + path_key(f.fieldspec) for f in tree_fields + if (f.sort_type and not count_only) or f.negate or f.op_num != 8 + or len(rank_path(f.fieldspec)[2]) != 1 + or f.fieldspec.is_relationship() or f.fieldspec.is_temporal() + } + if defer_display: + ranks.deferred = { + path_key(f.fieldspec) for f in tree_fields + if path_key(f.fieldspec) not in required + and len(rank_path(f.fieldspec)[2]) == 1 + and not f.fieldspec.is_relationship() + and not f.fieldspec.is_temporal() + } + + nodes_by_table = {} + for query_field in tree_fields: + spec = query_field.fieldspec + path, _, _ = rank_path(spec) + query, node, table, _ = query.build_join( + spec.root_table, spec.root_sql_table, path) + if path_key(spec) not in ranks.deferred: + nodes_by_table.setdefault(table, set()).add(node) + for table, nodes in nodes_by_table.items(): + nodes = list(nodes) + # Unpivot relationship IDs without repeating the candidate query. + slots = sql.union_all(*[ + sql.select(sql.literal(i).label("slot")) for i in range(len(nodes)) + ]).subquery() + candidate = query.query.with_entities(*[ + node._id.label(f"node_{i}") for i, node in enumerate(nodes) + ]).subquery() + seeds = sql.select(sql.case([ + (slots.c.slot == i, candidate.c[f"node_{i}"]) + for i in range(len(nodes)) + ]).label("node_id")).select_from(candidate.join(slots, sql.true())).distinct().subquery() + ranks.mappings[table] = ancestor_mapping(seeds, table) + ranks.mapped_nodes.update(nodes) + query.join_cache["tree_ranks"] = ranks + return query + + +def resolve_page(query, order, ranks, limit, offset): + """Walk display-only ancestors from selected rows, not the whole result. + + Base/to-many IDs break requested-sort ties without collapsing duplicate rows. + """ + order = [*order, *ranks.row_ids] + if not ranks.displays: + selected = query.order_by(*order).offset(offset) + return selected.limit(limit) if limit else selected + + records = list(ranks.displays.values()) + # Eligibility is checked before deferring fields: all belong to one tree. + nodes = list(dict.fromkeys(record[1][0] for record in records)) + columns = [entry["expr"] for entry in query.column_descriptions] + page_columns = [ + (sql.null() if i in ranks.displays else column).label(f"value_{i}") + for i, column in enumerate(columns) + ] + page_columns += [node._id.label(f"node_{i}") for i, node in enumerate(nodes)] + page_columns.append(sql.func.row_number().over(order_by=order).label("page_order")) + selected = query.with_entities(*page_columns).order_by(*order).offset(offset) + if limit: + selected = selected.limit(limit) + page = selected.subquery() + + # Carry page rows through real ancestor recursion. The page is a single seed + # of this CTE, not an independently materialized CTE referenced repeatedly. + table = records[0][1][1] + model = getattr(models, table.name) + slots = sql.union_all(*[ + sql.select(sql.literal(i).label("slot")) for i in range(len(nodes)) + ]).subquery() + seed_id = sql.case([ + (slots.c.slot == i, page.c[f"node_{i}"]) for i in range(len(nodes)) + ]) + node = orm.aliased(model) + definition = table.name + "TreeDefID" + item = table.name + "TreeDefItemID" + walk = sql.select( + *page.c, slots.c.slot, seed_id.label("seed_id"), + node._id.label("ancestor_id"), node.ParentID.label("parent_id"), + getattr(node, item).label("item_id"), + getattr(node, definition).label("definition_id"), + sql.true().label("is_seed"), + ).select_from(page.join(slots, sql.true()).outerjoin( + node, node._id == seed_id)).cte(recursive=True) + parent = orm.aliased(model) + recursive = sql.select( + *[walk.c[column.key] for column in page.c], + walk.c.slot, walk.c.seed_id, parent._id, parent.ParentID, + getattr(parent, item), getattr(parent, definition), sql.false(), + ).select_from(walk).join(parent, parent._id == walk.c.parent_id).where( + getattr(parent, definition) == walk.c.definition_id) + mapping = ranks.mappings.get(table) + if mapping is not None: + recursive = recursive.where(~sql.exists( + sql.select(mapping.c.seed_id).where(mapping.c.seed_id == walk.c.seed_id) + )) + walk = walk.union(recursive) + result = query.session.query(*[ + walk.c[f"value_{i}"] for i in range(len(columns)) + ]).select_from(walk).filter(walk.c.is_seed, walk.c.slot == 0) + aliases = {} + selected_nodes = {} + replacements = {} + for index, (query_field, (original_node, _, item_ids), formatter) in ranks.displays.items(): + node_id = walk.c[f"node_{nodes.index(original_node)}"] + key = (original_node, item_ids) + if key not in aliases: + ancestor = orm.aliased(model) + # Definition comes from the selected node; never use rank name alone. + if original_node not in selected_nodes: + selected_node = orm.aliased(model) + result = result.outerjoin(selected_node, selected_node._id == node_id) + selected_nodes[original_node] = selected_node + selected_node = selected_nodes[original_node] + definition_id = getattr(selected_node, definition) + resolved_id = ancestor_id(walk, node_id, item_ids, definition_id, distinct=True) + if mapping is not None: + resolved_id = sql.func.coalesce( + ancestor_id(mapping, node_id, item_ids, definition_id), resolved_id) + result = result.outerjoin(ancestor, sql.and_( + ancestor._id == resolved_id, + getattr(ancestor, definition) == definition_id, + )) + aliases[key] = ancestor + _, _, remaining = rank_path(query_field.fieldspec) + column = getattr(aliases[key], remaining[0].name) + replacements[index] = formatter.fieldformat(query_field, column) + return result.with_entities(*[ + replacements.get(i, walk.c[f"value_{i}"]) for i in range(len(columns)) + ]).order_by(walk.c.page_order).execution_options(specify_tree_ranks=ranks)