Restricting by a union, inserting from a union, or nesting a union with secondary attributes raises TypeError #1658

What happens

A union (A + B) works when it is fetched, counted, joined to another table, or restricted itself. Three other uses raise an error instead of running the query:

• restricting a table by a union, with & or -
• inserting the rows of a union into a table with Table.insert(union)
• a union whose left operand is a union, when the operands have secondary attributes, for example (Parent + ParentInfo) + ParentNote

Each one raises:

TypeError: Union.make_sql() takes 1 positional argument but 2 were given

Reproduction

On master (3c912c4). The results are the same on MySQL 8.0 and PostgreSQL 15.

import datajoint as dj

schema = dj.Schema("union_repro")

@schema
class Parent(dj.Manual):
    definition = """
    parent_id : int32
    ---
    label : varchar(8)
    """

@schema
class ParentInfo(dj.Manual):
    definition = """
    -> Parent
    ---
    info : varchar(8)
    """

@schema
class ParentNote(dj.Manual):
    definition = """
    -> Parent
    ---
    note : varchar(8)
    """

@schema
class Child(dj.Manual):
    definition = """
    -> Parent
    child_id : int32
    ---
    color : varchar(8)
    """

@schema
class Target(dj.Manual):
    definition = """
    parent_id : int32
    """

Parent.insert([
    {"parent_id": 1, "label": "a"},
    {"parent_id": 2, "label": "b"},
    {"parent_id": 3, "label": "a"},
    {"parent_id": 4, "label": "b"},
])
ParentInfo.insert([{"parent_id": 2, "info": "x"}])
ParentNote.insert([{"parent_id": 3, "note": "n"}])
Child.insert([
    {"parent_id": 1, "child_id": 1, "color": "red"},
    {"parent_id": 2, "child_id": 1, "color": "blue"},
    {"parent_id": 4, "child_id": 1, "color": "red"},
])

keys = (Parent & {"label": "a"}).proj() + ParentInfo.proj()
keys.keys()   # parents 1, 2, 3 -- the union on its own is correct

These already work: a nested union of primary keys only, such as (ParentInfo.proj() + ParentNote.proj()) + Parent.proj(); the same nested union with the inner union on the right, ParentNote + (Parent + ParentInfo); a join with a union; and a restriction applied to a union.

Cause

Union.make_sql(self) (src/datajoint/expression.py:1432) takes no fields argument. QueryExpression.make_sql (expression.py:147) and Aggregation.make_sql (expression.py:1355) both take fields=None. Three callers pass fields to a union:

• Restriction. make_condition (src/datajoint/condition.py:455) builds (attrs) in (subquery) from condition.make_sql(common_attributes). Two lines earlier (condition.py:445-446) it wraps an Aggregation in a subquery, but it does not wrap a Union.
• Insert from a query. Table.insert (src/datajoint/table.py:910) builds INSERT INTO ... (fields) {rows.make_sql(fields)}.
• Nested union with secondary attributes. In this branch Union.make_sql computes arg2 - arg1 (expression.py:1445). When arg1 is a union, that restriction goes through make_condition and fails as in the first case. The primary-key-only branch checks isinstance(arg1, Union) and calls make_sql() without arguments (expression.py:1439-1440). That check came from the fix for Cannot form the union of more than two tables #926. The other callers were never changed.

A join, a projection, or a restriction applied to a union first puts the union in a subquery (expression.py:268, 408-414, 613), and that path calls make_sql() with no arguments. This is why those operations work.

Expected

All six queries in the table should return the expected rows. Union.make_sql should take fields=None like the other make_sql methods. When fields is given, it should select those columns from the union, used as a subquery: SELECT <fields> FROM (<union>) AS .... Accepting the argument and ignoring it is not enough:

• In a restriction, the subquery must return only the shared attributes. Without that, Child & (ParentInfo + ParentNote) would put all of the union's columns inside IN (...).
• In an insert, the selected columns must match the insert's column list.

Another option is for both callers to put a union in a subquery with make_subquery(), as condition.py already does for an aggregation. Changing Union.make_sql fixes all three cases in one place.

#1657 (open) changes the same code for dj.Top: the subquery check at condition.py:445 and Union.make_sql. With #1657, a union limited by dj.Top would be put in a subquery when used in a restriction. A union without dj.Top would still raise this error in every case above. The two changes touch the same lines, so this fix should build on #1657 or be merged together with it.