Add CYCLE clause support for recursive CTEs

django_cte/cte.py

@@ -6,6 +6,7 @@
 from django.db.models.query import Q, QuerySet, ValuesIterable
 from django.db.models.sql.datastructures import BaseTable
 
+from .cycle import CycleConfig
 from .jitmixin import jit_mixin
 from .join import QJoin, INNER
 from .meta import CTEColumnRef, CTEColumns
@@ -43,24 +44,33 @@ class CTE:
     eventually be added.
     :param materialized: Optional parameter (default: False) which enforce
     using of MATERIALIZED statement for supporting databases.
+    :param cycle: Optional parameter (default: None) enabling cycle
+    detection for recursive CTEs. Either a sequence of CTE column names to
+    track for cycles, or a dict with 'columns', 'set', 'to', 'default',
+    'using' and 'using_output_field' keys. See `CycleConfig`.
     """
 
-    def __init__(self, queryset, name="cte", materialized=False):
+    def __init__(self, queryset, name="cte", materialized=False, cycle=None):
         self._set_queryset(queryset)
         self.name = name
         self.col = CTEColumns(self)
         self.materialized = materialized
+        self.cycle = CycleConfig.parse(cycle)
 
     def __getstate__(self):
-        return (self.query, self.name, self.materialized, self._iterable_class)
+        return (self.query, self.name, self.materialized, self._iterable_class, self.cycle)
 
     def __setstate__(self, state):
         if len(state) == 3:
             # Keep compatibility with the previous serialization method
             self.query, self.name, self.materialized = state
             self._iterable_class = ValuesIterable
-        else:
+            self.cycle = None
+        elif len(state) == 4:
             self.query, self.name, self.materialized, self._iterable_class = state
+            self.cycle = None
+        else:
+            self.query, self.name, self.materialized, self._iterable_class, self.cycle = state
         self.col = CTEColumns(self)
 
     def __repr__(self):
@@ -71,7 +81,7 @@ def _set_queryset(self, queryset):
         self._iterable_class = getattr(queryset, "_iterable_class", ValuesIterable)
 
     @classmethod
-    def recursive(cls, make_cte_queryset, name="cte", materialized=False):
+    def recursive(cls, make_cte_queryset, name="cte", materialized=False, cycle=None):
         """Recursive Common Table Expression
 
         :param make_cte_queryset: Function taking a single argument (a
@@ -80,9 +90,10 @@ def recursive(cls, make_cte_queryset, name="cte", materialized=False):
         statement unioned with a recursive statement.
         :param name: See `name` parameter of `__init__`.
         :param materialized: See `materialized` parameter of `__init__`.
+        :param cycle: See `cycle` parameter of `__init__`.
         :returns: The fully constructed recursive cte object.
         """
-        cte = cls(None, name, materialized)
+        cte = cls(None, name, materialized, cycle)
         cte._set_queryset(make_cte_queryset(cte))
         return cte
 
@@ -169,6 +180,10 @@ def queryset(self):
 
     def _resolve_ref(self, column):
         name = column.name
+        if self.cycle is not None and name in self.cycle.generated_columns:
+            # columns generated by the CYCLE clause are not in the CTE query
+            return CTEColumnRef(name, self.name, self.cycle.generated_columns[name])
+
         ref = self.query.resolve_ref(name)
         if ref is column or column in ref.get_source_expressions():
             raise ValueError(f"Circular reference: {column} = {ref}")

django_cte/cycle.py

@@ -0,0 +1,97 @@
+from django.db.models import BooleanField, TextField
+
+DICT_KEYS = frozenset([
+    "columns", "set", "to", "default", "using", "using_output_field",
+])
+
+
+class CycleConfig:
+    """CYCLE clause configuration of a recursive CTE
+
+    The `generated_columns` attribute maps the names of the columns added
+    by the clause to their output fields.
+
+    :param columns: Sequence of CTE column names to track for cycles.
+    :param mark_column: Name of the generated cycle mark column
+    (default: "is_cycle").
+    :param cycle_value: SQL literal assigned to the mark column when a
+    cycle is detected (default: "true"). Interpolated into the query as
+    written, so a string value must include its own quotes.
+    :param default_value: SQL literal assigned to the mark column when no
+    cycle is detected (default: "false"). Interpolated as written.
+    :param path_column: Name of the generated path column (default:
+    "path").
+    :param path_output_field: Output field of the path column (default:
+    `TextField()`). PostgreSQL generates the column as `ARRAY[RECORD]`,
+    and `RECORD` is a pseudo-type for unspecified row types, which
+    psycopg2 does not adapt to a list because it considers `RECORD`
+    unknown. Pass an `ArrayField` and configure list adaptation to get
+    anything other than text.
+
+    See:
+    * https://www.psycopg.org/docs/usage.html#lists-adaptation
+    * https://www.psycopg.org/docs/extensions.html#cast-array-unknown
+    * https://www.postgresql.org/docs/current/datatype-pseudo.html#DATATYPE-PSEUDO
+    """
+
+    def __init__(self, columns, mark_column="is_cycle", cycle_value="true",
+                 default_value="false", path_column="path",
+                 path_output_field=None):
+        if not columns:
+            raise ValueError("CYCLE requires at least one column")
+        self.columns = tuple(columns)
+        self.mark_column = mark_column
+        self.cycle_value = cycle_value
+        self.default_value = default_value
+        self.path_column = path_column
+        self.path_output_field = path_output_field or TextField()
+        self.generated_columns = {
+            self.mark_column: BooleanField(),
+            self.path_column: self.path_output_field,
+        }
+
+    @classmethod
+    def parse(cls, cycle):
+        """Get a config from the `cycle` argument of `CTE`
+
+        :param cycle: A `CycleConfig`, a sequence of column names, a dict
+        of `CycleConfig` keyword arguments by their public key names, or
+        None.
+        :returns: A `CycleConfig` or None.
+        """
+        if cycle is None or isinstance(cycle, cls):
+            return cycle
+        if isinstance(cycle, (list, tuple)):
+            return cls(cycle)
+        if isinstance(cycle, dict):
+            unknown = set(cycle) - DICT_KEYS
+            if unknown:
+                raise ValueError(
+                    f"Unknown cycle option(s): {', '.join(sorted(unknown))}. "
+                    f"Valid options are: {', '.join(sorted(DICT_KEYS))}"
+                )
+            return cls(
+                cycle.get("columns", ()),
+                mark_column=cycle.get("set", "is_cycle"),
+                cycle_value=cycle.get("to", "true"),
+                default_value=cycle.get("default", "false"),
+                path_column=cycle.get("using", "path"),
+                path_output_field=cycle.get("using_output_field"),
+            )
+        raise ValueError(
+            "cycle must be a sequence of column names or a dict, "
+            f"got {type(cycle).__name__}"
+        )
+
+    def as_sql(self, qn):
+        """Get the CYCLE clause SQL
+
+        :param qn: Name quoting function.
+        :returns: The CYCLE clause SQL string.
+        """
+        return (
+            f"CYCLE {', '.join(qn(c) for c in self.columns)} "
+            f"SET {qn(self.mark_column)} "
+            f"TO {self.cycle_value} DEFAULT {self.default_value} "
+            f"USING {qn(self.path_column)}"
+        )

django_cte/query.py

@@ -92,7 +92,7 @@ def generate_cte_sql(connection, query, as_sql):
             # like, col_count and klass_info.
             as_sql()
             raise
-        template = get_cte_query_template(cte)
+        template = get_cte_query_template(cte, qn)
         ctes.append(template.format(name=qn(cte.name), query=cte_sql))
         params.extend(cte_params)
 
@@ -129,10 +129,16 @@ def generate_cte_sql(connection, query, as_sql):
     return " ".join(sql), tuple(params)
 
 
-def get_cte_query_template(cte):
+def get_cte_query_template(cte, qn):
+    template = "{name} AS"
     if cte.materialized:
-        return "{name} AS MATERIALIZED ({query})"
-    return "{name} AS ({query})"
+        template += " MATERIALIZED"
+    template += " ({query})"
+
+    if cte.cycle is not None:
+        template += " " + cte.cycle.as_sql(qn)
+
+    return template
 
 
 def _ignore_with_col_aliases(cte_query):

docs/index.md

@@ -155,6 +155,126 @@ ORDER BY "path" ASC
 ```
 
 
+## Cycle Detection in Recursive CTEs
+
+Recursive queries can loop indefinitely if the data contains cycles. Some
+databases support a `CYCLE` clause, which stops the recursion when a row
+repeats. Support was added in PostgreSQL 14. On earlier versions, and on
+databases without it such as SQLite, cycle detection must be implemented in
+application logic.
+
+Pass the `cycle` parameter to `CTE()` or `CTE.recursive()`, naming the CTE
+column(s) that identify a row:
+
+```py
+def make_regions_cte(cte):
+    return Region.objects.filter(
+        parent__isnull=True
+    ).values("name", "parent_id").union(
+        cte.join(Region, parent=cte.col.name).values("name", "parent_id"),
+        all=True,
+    )
+
+cte = CTE.recursive(make_regions_cte, cycle=["name"])
+```
+
+This generates a `CYCLE` clause with default settings:
+
+```sql
+WITH RECURSIVE "cte" AS (
+    ...
+) CYCLE "name" SET "is_cycle" TO true DEFAULT false USING "path"
+```
+
+The clause adds two columns to the CTE: a mark column, `is_cycle`, which is
+true on the row where a cycle was found, and a path column, `path`, holding
+the rows visited on the way to it. Reference them like any other CTE column,
+with `cte.col`:
+
+```py
+regions = with_cte(
+    cte,
+    select=cte.join(Region, name=cte.col.name)
+    .annotate(is_cycle=cte.col.is_cycle, path=cte.col.path)
+)
+```
+
+Neither name may collide with a column the CTE query already selects.
+PostgreSQL rejects such a query outright, whether or not the column is
+referenced. Rename them with the `set` and `using` keys of the dict form,
+which also controls the values assigned to the mark column:
+
+```py
+cte = CTE.recursive(
+    make_regions_cte,
+    cycle={
+        "columns": ["name", "parent_id"],  # columns to track
+        "set": "cycle_detected",           # mark column name
+        "to": "1",                         # SQL literal when cycle detected
+        "default": "0",                    # SQL literal when no cycle
+        "using": "cycle_path",             # path column name
+    }
+)
+```
+
+This generates:
+
+```sql
+WITH RECURSIVE "cte" AS (
+    ...
+) CYCLE "name", "parent_id" SET "cycle_detected" TO 1 DEFAULT 0 USING "cycle_path"
+```
+
+Column names are quoted, but `to` and `default` are SQL literals written into
+the query verbatim, so a string value must carry its own quotes:
+`"to": "'yes'"`. The mark column takes its type from these two literals.
+
+### Working with the USING Column
+
+PostgreSQL generates the path column as `ARRAY[RECORD]`, and Django has no
+field for that type. Its `output_field` therefore defaults to `TextField`,
+which applies no conversion and returns whatever the driver produced.
+
+The `using_output_field` key changes the field Django attaches to the column.
+It does not change the returned value, only which lookups are allowed.
+`ArrayField` adds `__len`. Other array lookups compare against the declared
+element type and fail, because the elements are anonymous records:
+
+```py
+from django.contrib.postgres.fields import ArrayField
+from django.db.models import TextField
+
+cte = CTE.recursive(
+    make_regions_cte,
+    cycle={
+        "columns": ["name"],
+        "using_output_field": ArrayField(TextField()),
+    }
+)
+```
+
+How the value arrives in Python depends on the driver. psycopg2 returns the
+array as a string, quoting a record only when its text needs escaping, such
+as when a value contains a comma or a space:
+
+```py
+'{(sun),(earth),(moon)}'
+```
+
+psycopg 3 adapts it to a list of tuples:
+
+```py
+[('sun',), ('earth',), ('moon',)]
+```
+
+Using `ArrayField` with psycopg2 needs additional list adaptation, because
+psycopg2 does not convert `ARRAY[RECORD]` to a list: it considers `RECORD` an
+unknown type. For more details, see:
+
+- [psycopg2 Lists Adaptation](https://www.psycopg.org/docs/usage.html#lists-adaptation)
+- [psycopg2 Cast Array Unknown](https://www.psycopg.org/docs/extensions.html#cast-array-unknown)
+
+
 ## Named Common Table Expressions
 
 It is possible to add more than one CTE to a query. To do this, each CTE must