fix(cascade): seed a delete with its keys when its restriction is another query, and keep a % in it

src/datajoint/diagram.py

@@ -20,6 +20,7 @@
 import inspect
 import io
 import logging
+import re
 
 import networkx as nx
 
@@ -154,18 +155,37 @@ def stroke_rule(light_hex, dark_hex, opacity=None):
     return "<style>\n@media (prefers-color-scheme: dark) {\n" + body + "\n}\n</style>"
 
 
+def _has_subquery(condition):
+    """Whether a compiled restriction condition contains a subquery."""
+    if not isinstance(condition, str):
+        return False
+    text = re.sub(r"(\\\")|(\\\')", "", condition)  # escaped quotes
+    text = re.sub(r"'[^']*'|\"[^\"]*\"", "", text)  # quoted text and identifiers
+    return re.search(r"\bselect\b", text, re.IGNORECASE) is not None
+
+
 def _seed_restriction(table_expr):
     """
     Return a traversal seed's restriction and the attributes it names.
 
-    A ``dj.Top`` restriction is an ORDER BY ... LIMIT, not part of the WHERE
-    clause, so a seed restricted by ``Top`` is replaced by the primary keys
-    it selects. Otherwise the traversal would start from every row that
-    matches the seed's other conditions.
+    Two kinds of seed are replaced by the primary keys they select:
+
+    - A ``dj.Top`` restriction is an ORDER BY ... LIMIT, not part of the WHERE
+      clause. Otherwise the traversal would start from every row that matches
+      the seed's other conditions.
+    - A restriction by another query, such as ``Subject & (Session & "quality = 0")``,
+      compiles to a subquery. The traversal applies the seed's conditions again
+      to a fresh copy of the table, which rejects a subquery that names other
+      tables' columns. A delete would also remove the rows the subquery reads,
+      when they are downstream, before it deletes the seed itself.
     """
-    if table_expr._top is not None:
+    if table_expr._top is not None or any(_has_subquery(c) for c in table_expr.restriction):
         return AndList([table_expr.keys()]), set(table_expr.primary_key)
-    return AndList(table_expr.restriction), set(table_expr.restriction_attributes)
+    # The seed's conditions are compiled SQL, with each % already escaped as %%
+    # for the driver. The traversal applies them again through restrict(), which
+    # escapes again, so the escape is undone here.
+    restriction = AndList(c.replace("%%", "%") if isinstance(c, str) else c for c in table_expr.restriction)
+    return restriction, set(table_expr.restriction_attributes)
 
 
 class Diagram(nx.MultiDiGraph):  # noqa: C901

tests/integration/test_cascade_overdelete.py

@@ -192,3 +192,82 @@ class Downstream(dj.Manual):
     assert Master.to_arrays("master_id").tolist() == [2]
     assert len(Master.P()) == 0
     assert Downstream.to_arrays("master_id").tolist() == [2]
+
+
+SUBQUERY_SEEDS = {
+    # seed -> (subjects deleted, sessions deleted)
+    "restricted by a query": (lambda S, T: S & (T & "quality = 0"), {1, 3}, {(1, 1), (1, 2), (3, 1)}),
+    "anti-restricted by a query": (lambda S, T: S - (T & "quality = 0"), {2, 4}, {(2, 1)}),
+    "restricted by a table": (lambda S, T: S & T, {1, 2, 3}, {(1, 1), (1, 2), (2, 1), (3, 1)}),
+}
+
+
+@pytest.mark.parametrize("seed_kind", list(SUBQUERY_SEEDS))
+def test_seed_restricted_by_a_downstream_query(schema_by_backend, seed_kind):
+    """A seed restricted by a query on a downstream table deletes the rows that
+    query selected when the delete started, with all their descendants."""
+
+    @schema_by_backend
+    class Subject(dj.Manual):
+        definition = """
+        subject_id : int32
+        """
+
+    @schema_by_backend
+    class Session(dj.Manual):
+        definition = """
+        -> Subject
+        session_id : int32
+        ---
+        quality : int32
+        """
+
+    Subject.insert([(1,), (2,), (3,), (4,)])
+    Session.insert([(1, 1, 1), (1, 2, 0), (2, 1, 1), (3, 1, 0)])
+    all_sessions = {(1, 1), (1, 2), (2, 1), (3, 1)}
+
+    make_seed, subjects, sessions = SUBQUERY_SEEDS[seed_kind]
+    seed = make_seed(Subject, Session)
+
+    counts = dj.Diagram.cascade(seed).counts()
+    assert counts[Subject.full_table_name] == len(subjects)
+    assert counts[Session.full_table_name] == len(sessions)
+
+    seed.delete(prompt=False)
+    assert set(Subject.to_arrays("subject_id").tolist()) == {1, 2, 3, 4} - subjects
+    assert set(zip(*Session.to_arrays("subject_id", "session_id"))) == all_sessions - sessions
+
+
+def test_seed_condition_with_a_percent_sign(schema_by_backend):
+    """A seed condition containing % selects the same rows in the cascade as
+    on its own."""
+
+    @schema_by_backend
+    class Label(dj.Manual):
+        definition = """
+        label_id : int32
+        ---
+        name : varchar(16)
+        """
+
+    @schema_by_backend
+    class Use(dj.Manual):
+        definition = """
+        -> Label
+        use_id : int32
+        """
+
+    Label.insert([(1, "50%"), (2, "other"), (3, "5%%")])
+    Use.insert([(1, 1), (2, 1), (3, 1)])
+
+    for seed, label in [
+        (Label & "name = '50%'", 1),
+        (Label & {"name": "50%"}, 1),
+        (Label & "name = '5%%'", 3),
+    ]:
+        assert len(seed) == 1
+        assert dj.Diagram.cascade(seed).counts()[Use.full_table_name] == 1
+
+    (Label & "name LIKE '50%'").delete(prompt=False)
+    assert sorted(Label.to_arrays("label_id").tolist()) == [2, 3]
+    assert sorted(Use.to_arrays("label_id").tolist()) == [2, 3]