Issue: Using CTE in a UNION ALL query
Author: 8BitMate
State: closed
Created: 2020-10-30 11:13:27

I'm trying to use a CTE on the top level of a UNION ALL query, but I can't figure out how to do it in peewee. The functionality I'm trying to achieve is an "INSERT INTO ... ON CONFLICT DO NOTHING RETURNING ..." query where it always returns the value(s) even when it already exists.

The peewee code I'm using looks something like this:
```Python
vehicle_insert = pw.CTE("vehicle_insert",
    meta.Vehicle.insert(
        uid=result.vehicle_ref,
        type=vehicle_type,
        provider=transport_provider,
    ).on_conflict_ignore().returning(meta.Vehicle.id),
    columns=("id",)
)

vehicle = (
    meta.Vehicle.select(
        meta.Vehicle.id,
    )
    .where(
        (meta.Vehicle.uid == result.vehicle_ref) &
        (meta.Vehicle.provider == transport_provider)
    )
    .union_all(
        vehicle_insert
    )
    .with_cte(
        vehicle_insert
    )
)
```
When I execute this query, peewee doesn't generate the WITH clause.

I've tried moving the ```with_cte``` function inside the ```union_all``` like this: 
```Python
.union_all(
    vehicle_insert.select(vehicle_insert.c.id).with_cte(vehicle_insert)
)
```
but then I get another error:
```
peewee.NotSupportedError: WITH clause containing a data-modifying statement must be at the top level
```
Is there any way to generate a WITH clause on the top level of a UNION ALL in peewee?

Can you share the SQL for the query you're trying to write? It might be easier for me to help you work backwards from that.

Sure, I'm trying to write something like this:

```SQL
WITH vehicle_insert (id) AS (
    INSERT INTO vehicle (uid, type_id, provider_id)
    VALUES ('Some:ID', 2, 3)
    ON CONFLICT DO NOTHING RETURNING id
)
SELECT id
FROM vehicle
WHERE ((uid = 'Some:ID') AND (provider_id = 3))
UNION ALL
SELECT id
FROM vehicle_insert
```
I have a unique index on ```uid``` and ```provider_id```, so the query will get the id from the vehicle table if a matching row exists, or it will return the result from the returning clause if it doesn't exist.

I also noticed from my example above that the call to ```union_all``` should be:
```Python
union_all(vehicle_insert.select(vehicle_insert.c.id))
```
But it doesn't solve the problem.

I needed to add a small patch to the `CompoundSelectQuery` implementation to properly handle cte's bound to the union itself. I was able to get an *almost* correct query except for the parentheses, but now I've got a working query and the syntax isn't too awkward. The main takeaways are:

1. the cte needs to be declared explicitly (as you are doing in your code), since I don't have `cte()` method support on the data modifying queries yet (see #2152).
2. the cte needs to be bound to the union so that it ends up in the outermost layer
3. we need to also explicitly create the bit that selects from the CTE since peewee's `CTE.select_from()` assumes you want to bind the cte to the select.

Here is the code:

```python

iq = (Vehicle
      .insert(uid=1, type='car', provider='p1')
      .on_conflict_ignore()
      .returning(Vehicle.id))

cte = CTE('vehicle_iq', iq, columns=('id',))

q1 = (Vehicle
      .select(Vehicle.id)
      .where((Vehicle.uid == 1) & (Vehicle.provider == 'p1')))
q2 = Select([cte], [cte.c.id])

u = q1.union_all(q2).with_cte(cte)
```

This patch is needed as well: 92494cd9ed01b12826b2fddb25cde221ea34f155

I'm going to close this for now and will look at reducing some of these clumsier pieces as part of #2152.