issue while using cte with union querysets

When using cte with a queryset that has been union before, it is losing the WITH part on the SQL.

# example 
non_staff_users = User.objects.filter(is_staff=False).only("id", "email")
non_staff_users_cte = With(non_staff_users, name="non_staff_users_cte")

non_staff_users_cte_queryset = non_staff_users_cte.queryset().only("id", "email")

non_staff_gmail_users = non_staff_users_cte_queryset.filter(
    email__endswith="@gmail.com"
).annotate(
    email_provider=Value("gmail", output_field=TextField()),
)

non_staff_yahoo_users = non_staff_users_cte_queryset.filter(
    email__endswith="@yahoo.com"
).annotate(
    email_provider=Value("yahoo", output_field=TextField()),
)

non_staff_gmail_and_yahoo_users = non_staff_gmail_users.union(
    non_staff_yahoo_users, all=True)

non_staff_gmail_users_with_cte = non_staff_gmail_users.with_cte(
    non_staff_users_cte)
print(non_staff_gmail_users_with_cte.query)
# WITH RECURSIVE "non_staff_users_cte" AS (SELECT "auth_user"."id", "auth_user"."email" FROM "auth_user" WHERE NOT "auth_user"."is_staff") SELECT "non_staff_users_cte"."id", "non_staff_users_cte"."email", gmail AS "email_provider" FROM "non_staff_users_cte" WHERE "non_staff_users_cte"."email"::text LIKE %@gmail.com

non_staff_yahoo_users_with_cte = non_staff_yahoo_users.with_cte(
    non_staff_users_cte)
print(non_staff_yahoo_users_with_cte.query)
# WITH RECURSIVE "non_staff_users_cte" AS (SELECT "auth_user"."id", "auth_user"."email" FROM "auth_user" WHERE NOT "auth_user"."is_staff") SELECT "non_staff_users_cte"."id", "non_staff_users_cte"."email", yahoo AS "email_provider" FROM "non_staff_users_cte" WHERE "non_staff_users_cte"."email"::text LIKE %@yahoo.com

non_staff_gmail_and_yahoo_users_with_cte = non_staff_gmail_and_yahoo_users.with_cte(
    non_staff_users_cte)
print(non_staff_gmail_and_yahoo_users_with_cte.query)
# Incorrectly Printed Result
# (SELECT "non_staff_users_cte"."id" AS "col1", "non_staff_users_cte"."email" AS "col2", gmail AS "email_provider" FROM "non_staff_users_cte" WHERE "non_staff_users_cte"."email"::text LIKE %@gmail.com) UNION ALL (SELECT "non_staff_users_cte"."id" AS "col1", "non_staff_users_cte"."email" AS "col2", yahoo AS "email_provider" FROM "non_staff_users_cte" WHERE "non_staff_users_cte"."email"::text LIKE %@yahoo.com)

# Expected Result
# WITH RECURSIVE "non_staff_users_cte" AS (SELECT "auth_user"."id", "auth_user"."email" FROM "auth_user" WHERE NOT "auth_user"."is_staff") (SELECT "non_staff_users_cte"."id" AS "col1", "non_staff_users_cte"."email" AS "col2", gmail AS "email_provider" FROM "non_staff_users_cte" WHERE "non_staff_users_cte"."email"::text LIKE %@gmail.com) UNION ALL (SELECT "non_staff_users_cte"."id" AS "col1", "non_staff_users_cte"."email" AS "col2", yahoo AS "email_provider" FROM "non_staff_users_cte" WHERE "non_staff_users_cte"."email"::text LIKE %@yahoo.com)

if you attemt to call

non_staff_gmail_and_yahoo_users_with_cte.first()

it is going to throw an error as:

Unhandled ProgrammingError: relation "non_staff_users_cte" does not exist
LINE 1: ...mail" AS "col2", 'gmail' AS "email_provider" FROM "non_staff...
                                                             ^

Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/django/db/backends/utils.py", line 89, in _execute
    return self.cursor.execute(sql, params)
psycopg2.errors.UndefinedTable: relation "non_staff_users_cte" does not exist
LINE 1: ...mail" AS "col2", 'gmail' AS "email_provider" FROM "non_staff...

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/lib/python3.9/code.py", line 90, in runcode
    exec(code, self.locals)
  File "<console>", line 1, in <module>
  File "/usr/local/lib/python3.9/site-packages/django/db/models/query.py", line 1057, in first
    for obj in queryset[:1]:
  File "/usr/local/lib/python3.9/site-packages/django/db/models/query.py", line 398, in __iter__
    self._fetch_all()
  File "/usr/local/lib/python3.9/site-packages/django/db/models/query.py", line 1881, in _fetch_all
    self._result_cache = list(self._iterable_class(self))
  File "/usr/local/lib/python3.9/site-packages/django/db/models/query.py", line 91, in __iter__
    results = compiler.execute_sql(
  File "/usr/local/lib/python3.9/site-packages/django/db/models/sql/compiler.py", line 1562, in execute_sql
    cursor.execute(sql, params)
  File "/usr/local/lib/python3.9/site-packages/django/db/backends/utils.py", line 102, in execute
    return super().execute(sql, params)
  File "/usr/local/lib/python3.9/site-packages/django/db/backends/utils.py", line 67, in execute
    return self._execute_with_wrappers(
  File "/usr/local/lib/python3.9/site-packages/django/db/backends/utils.py", line 80, in _execute_with_wrappers
    return executor(sql, params, many, context)
  File "/usr/local/lib/python3.9/site-packages/django/db/backends/utils.py", line 89, in _execute
    return self.cursor.execute(sql, params)
  File "/usr/local/lib/python3.9/site-packages/django/db/utils.py", line 91, in __exit__
    raise dj_exc_value.with_traceback(traceback) from exc_value
  File "/usr/local/lib/python3.9/site-packages/django/db/backends/utils.py", line 89, in _execute
    return self.cursor.execute(sql, params)
django.db.utils.ProgrammingError: relation "non_staff_users_cte" does not exist
LINE 1: ...mail" AS "col2", 'gmail' AS "email_provider" FROM "non_staff...

dimagi / django-cte

issue while using cte with union querysets #85