diff --git a/docs/src/piccolo/query_clauses/index.rst b/docs/src/piccolo/query_clauses/index.rst index 889ab0cc1..5a5ffb705 100644 --- a/docs/src/piccolo/query_clauses/index.rst +++ b/docs/src/piccolo/query_clauses/index.rst @@ -31,6 +31,7 @@ by modifying the return values. ./on_conflict ./output ./returning + ./with_ .. toctree:: :maxdepth: 1 diff --git a/docs/src/piccolo/query_clauses/with_.rst b/docs/src/piccolo/query_clauses/with_.rst new file mode 100644 index 000000000..1df854d0d --- /dev/null +++ b/docs/src/piccolo/query_clauses/with_.rst @@ -0,0 +1,167 @@ +.. _with_: + +``with_`` +========= + +You can use the ``with_`` clause with the following queries: + +* :ref:`Select` +* :ref:`Update` +* :ref:`Delete` +* :ref:`Insert` + +It adds a ``WITH`` clause (a Common Table Expression, or CTE) to the start of +the query. A CTE is a named sub query, which the rest of the query can refer +to. + +.. note:: Postgres, CockroachDB and SQLite 3.8.3 or above. + +------------------------------------------------------------------------------- + +Basic usage +----------- + +Create a :class:`CTE ` with a name and a body. The body can +be a query, or a :class:`QueryString `. + +.. code-block:: python + + from piccolo.query import CTE + + popular = CTE( + "popular", + Band.select(Band.id).where(Band.popularity > 1000), + ) + +Then pass it to ``with_``. The name of the CTE can now be used in the query, +for example in a sub select: + +.. code-block:: python + + >>> from piccolo.querystring import QueryString + >>> await Concert.select(Concert.id).where( + ... Concert.band_1.is_in(QueryString('SELECT "id" FROM "popular"')) + ... ).with_(popular) + +Which is equivalent to: + +.. code-block:: sql + + WITH "popular" AS (SELECT ... FROM "band" WHERE "band"."popularity" > 1000) + SELECT ... FROM "concert" WHERE "concert"."band_1" IN (SELECT "id" FROM "popular") + +You can pass in several CTEs, and they're added in the order given. Each one +must have a different name. Later CTEs can refer to earlier ones. + +.. code-block:: python + + await Band.select().with_(first_cte, second_cte) + +The ``with_`` clause works in the same way with updates, deletes and inserts: + +.. code-block:: python + + await Band.update({Band.popularity: 0}).where( + Band.id.is_in(QueryString('SELECT "id" FROM "popular"')) + ).with_(popular) + +------------------------------------------------------------------------------- + +column_names +------------ + +By default the columns of the CTE are named by its body. You can name them +yourself using ``column_names``. If the body is a ``Select`` query, it must +select the same number of columns. + +.. code-block:: python + + popular = CTE( + "popular", + Band.select(Band.id, Band.name).where(Band.popularity > 1000), + column_names=["band_id", "band_name"], + ) + +``column_names`` and the CTE name must only contain letters, digits and +underscores, and must not start with a digit. Otherwise a +:class:`CTEError ` is raised. + +To refer to the columns, call ``ref`` on the CTE. This returns a +:class:`CTERef `, which has a ``CTEColumn`` for each of +the column names: + +.. code-block:: python + + >>> ref = popular.ref() + >>> str(ref.band_id) + '"popular"."band_id"' + >>> ref.select_all().compile_string()[0] + 'SELECT "popular"."band_id", "popular"."band_name" FROM "popular"' + +------------------------------------------------------------------------------- + +Recursive CTEs +-------------- + +A recursive CTE has a base query, and a recursive query which refers to the CTE +itself. They're combined using ``UNION ALL``. ``column_names`` is required. +For example, to find everyone who reports to a particular employee, directly +or indirectly: + +.. code-block:: python + + class Employee(Table): + name = Varchar() + manager_id = Integer(null=True) + + reports = CTE.recursive( + "reports", + base=QueryString( + 'SELECT "id", 0 FROM "employee" WHERE "id" = {}', 2 + ), + recursive=QueryString( + 'SELECT "employee"."id", "reports"."depth" + 1 ' + 'FROM "employee" JOIN "reports" ' + 'ON "employee"."manager_id" = "reports"."id"' + ), + column_names=["id", "depth"], + ) + + >>> await Employee.select(Employee.name).where( + ... Employee.id.is_in(QueryString('SELECT "id" FROM "reports"')) + ... ).with_(reports) + [{'name': 'vp'}, {'name': 'lead'}, {'name': 'engineer'}] + +If any of the CTEs passed to ``with_`` is recursive, the query starts with +``WITH RECURSIVE``. + +------------------------------------------------------------------------------- + +materialized +------------ + +Postgres, CockroachDB and SQLite 3.35 or above let you control whether the result of a CTE is +computed once and stored, using ``MATERIALIZED`` or ``NOT MATERIALIZED``. + +.. code-block:: python + + CTE("popular", query, materialized=True) # AS MATERIALIZED (...) + CTE("popular", query, materialized=False) # AS NOT MATERIALIZED (...) + CTE("popular", query) # database decides + +------------------------------------------------------------------------------- + +Source +------ + +.. currentmodule:: piccolo.query + +.. autoclass:: CTE + :members: recursive, ref + +.. autoclass:: CTERef + :members: column, columns, from_clause, select_all + +.. autoclass:: CTEColumn + +.. autoclass:: CTEError diff --git a/docs/src/piccolo/query_types/delete.rst b/docs/src/piccolo/query_types/delete.rst index 54401f9f3..7d8c86bad 100644 --- a/docs/src/piccolo/query_types/delete.rst +++ b/docs/src/piccolo/query_types/delete.rst @@ -42,3 +42,8 @@ where ~~~~~ See :ref:`Where` + +with\_ +~~~~~~ + +See :ref:`with_`. diff --git a/docs/src/piccolo/query_types/insert.rst b/docs/src/piccolo/query_types/insert.rst index 6648d6d0b..c20ad601e 100644 --- a/docs/src/piccolo/query_types/insert.rst +++ b/docs/src/piccolo/query_types/insert.rst @@ -47,3 +47,8 @@ returning ~~~~~~~~~ See :ref:`Returning`. + +with\_ +~~~~~~ + +See :ref:`with_`. diff --git a/docs/src/piccolo/query_types/select.rst b/docs/src/piccolo/query_types/select.rst index bce3c250d..2fab0f7d1 100644 --- a/docs/src/piccolo/query_types/select.rst +++ b/docs/src/piccolo/query_types/select.rst @@ -409,3 +409,8 @@ where ~~~~~ See :ref:`where`. + +with\_ +~~~~~~ + +See :ref:`with_`. diff --git a/docs/src/piccolo/query_types/update.rst b/docs/src/piccolo/query_types/update.rst index acd0e1bb8..1e822d3e5 100644 --- a/docs/src/piccolo/query_types/update.rst +++ b/docs/src/piccolo/query_types/update.rst @@ -228,3 +228,8 @@ where ~~~~~ See :ref:`Where`. + +with\_ +~~~~~~ + +See :ref:`with_`. diff --git a/piccolo/query/__init__.py b/piccolo/query/__init__.py index 846714795..1febec429 100644 --- a/piccolo/query/__init__.py +++ b/piccolo/query/__init__.py @@ -17,12 +17,17 @@ TableExists, Update, ) +from .methods.cte import CTE, CTEColumn, CTEError, CTERef from .methods.select import SelectRaw # for backwards compatibility from .mixins import OrderByRaw # for backwards compatibility __all__ = [ "Alter", "Avg", + "CTE", + "CTEColumn", + "CTEError", + "CTERef", "Count", "Create", "CreateIndex", diff --git a/piccolo/query/methods/cte.py b/piccolo/query/methods/cte.py new file mode 100644 index 000000000..c6b21e8a3 --- /dev/null +++ b/piccolo/query/methods/cte.py @@ -0,0 +1,332 @@ +from __future__ import annotations + +import re +from collections.abc import Iterator, Sequence +from typing import TYPE_CHECKING, Any, Optional, Union + +from piccolo.querystring import QueryString, Selectable + +if TYPE_CHECKING: + from piccolo.query.base import Query + + +_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +class CTEError(Exception): + """ + Raised when a ``CTE`` is malformed, or is used incorrectly. + """ + + pass + + +def _validate_identifier(value: str, kind: str) -> None: + if not isinstance(value, str): + raise CTEError(f"{kind} must be a string.") + if not _NAME_PATTERN.match(value): + raise CTEError(f"{kind} {value!r} is not a valid SQL identifier.") + + +def _count_select_columns(query: Any) -> Optional[int]: + from piccolo.query.methods.select import Select + + if isinstance(query, Select): + cols = query.columns_delegate.selected_columns + if cols: + return len(cols) + return None + + +class CTEColumn(Selectable): + """ + A column of a ``CTE``, which is rendered as ``"cte_name"."column_name"``. + Usually created via a :class:`CTERef `. + """ + + __slots__ = ("cte_name", "column_name") + + def __init__(self, cte_name: str, column_name: str) -> None: + _validate_identifier(cte_name, "CTE name") + _validate_identifier(column_name, "CTE column name") + self.cte_name = cte_name + self.column_name = column_name + self._alias = None + + def get_select_string( + self, engine_type: str, with_alias: bool = True + ) -> QueryString: + return QueryString(f'"{self.cte_name}"."{self.column_name}"') + + def __str__(self) -> str: + return f'"{self.cte_name}"."{self.column_name}"' + + +class CTERef: + """ + Describes the columns of a ``CTE``, so they can be referenced in other + parts of a query. Get one using :meth:`CTE.ref`. + """ + + __slots__ = ("cte_name", "column_names") + + def __init__(self, cte_name: str, column_names: Sequence[str]) -> None: + _validate_identifier(cte_name, "CTE name") + if not column_names: + raise CTEError("CTERef requires at least one column to reference.") + for cn in column_names: + _validate_identifier(cn, "CTE column name") + self.cte_name = cte_name + self.column_names = tuple(column_names) + + def column(self, name: str) -> CTEColumn: + """ + Returns the named column. Raises a ``CTEError`` if the CTE doesn't + declare it. + """ + if name not in self.column_names: + raise CTEError( + f"CTE {self.cte_name!r} has no column {name!r}; " + f"declared columns: {list(self.column_names)}" + ) + return CTEColumn(self.cte_name, name) + + def columns(self) -> tuple[CTEColumn, ...]: + """ + Returns all of the columns, in the order they were declared. + """ + return tuple( + CTEColumn(self.cte_name, name) for name in self.column_names + ) + + def __getattr__(self, name: str) -> CTEColumn: + if name.startswith("_"): + raise AttributeError(name) + if name not in self.column_names: + raise AttributeError( + f"CTE {self.cte_name!r} has no column {name!r}" + ) + return CTEColumn(self.cte_name, name) + + def __len__(self) -> int: + return len(self.column_names) + + def __iter__(self) -> Iterator[CTEColumn]: + for name in self.column_names: + yield CTEColumn(self.cte_name, name) + + def __contains__(self, name: object) -> bool: + return isinstance(name, str) and name in self.column_names + + def from_clause(self) -> QueryString: + """ + Returns the quoted CTE name, for use in a ``FROM`` clause. + """ + return QueryString(f'"{self.cte_name}"') + + def select_all(self) -> QueryString: + """ + Returns a ``SELECT`` query which fetches every column of the CTE. + """ + cols_sql = ", ".join( + f'"{self.cte_name}"."{c}"' for c in self.column_names + ) + return QueryString(f'SELECT {cols_sql} FROM "{self.cte_name}"') + + +class CTE: + """ + A Common Table Expression, which can be added to a ``Select``, ``Update``, + ``Delete`` or ``Insert`` query using ``with_``. + + :param name: + The name used to refer to the CTE in the rest of the query. + :param body: + A query, or a :class:`QueryString `. + :param recursive_body: + If provided, this is combined with ``body`` using ``UNION ALL``, to + make a recursive CTE. It should select from the CTE itself, and + ``column_names`` must also be provided. + :param column_names: + Names for the columns of the CTE. If ``body`` is a ``Select`` query, + it must select the same number of columns. + :param materialized: + ``True`` adds a ``MATERIALIZED`` hint, ``False`` adds + ``NOT MATERIALIZED``, and ``None`` leaves it up to the database. + + """ + + __slots__ = ( + "name", + "body", + "recursive_body", + "column_names", + "_recursive", + "_materialized", + ) + + def __init__( + self, + name: str, + body: Union["Query", QueryString], + recursive_body: Optional[Union["Query", QueryString]] = None, + column_names: Optional[Sequence[str]] = None, + materialized: Optional[bool] = None, + ) -> None: + _validate_identifier(name, "CTE name") + if column_names is not None: + if len(column_names) == 0: + raise CTEError( + "column_names must contain at least one entry, " + "or be omitted entirely." + ) + seen: set[str] = set() + for cn in column_names: + _validate_identifier(cn, "CTE column name") + if cn in seen: + raise CTEError( + f"duplicate column name {cn!r} in column_names." + ) + seen.add(cn) + if materialized is not None and not isinstance(materialized, bool): + raise CTEError( + "materialized must be True, False, or None (default)." + ) + self.name = name + self.body = body + self.recursive_body = recursive_body + self.column_names = ( + tuple(column_names) if column_names is not None else None + ) + self._recursive = recursive_body is not None + self._materialized = materialized + if self._recursive and column_names is None: + raise CTEError( + "A recursive CTE must declare column_names so the " + "recursive reference is well-typed." + ) + if column_names is not None: + body_count = _count_select_columns(body) + if body_count is not None and body_count != len(column_names): + raise CTEError( + f"CTE {name!r}: body selects {body_count} columns " + f"but column_names declares {len(column_names)}." + ) + if self._recursive: + rec_count = _count_select_columns(self.recursive_body) + if rec_count is not None and rec_count != len(column_names): + raise CTEError( + f"CTE {name!r}: recursive body selects " + f"{rec_count} columns but column_names " + f"declares {len(column_names)}." + ) + + @classmethod + def recursive( + cls, + name: str, + base: Union["Query", QueryString], + recursive: Union["Query", QueryString], + column_names: Sequence[str], + materialized: Optional[bool] = None, + ) -> "CTE": + """ + A shortcut for creating a recursive CTE. + """ + return cls( + name=name, + body=base, + recursive_body=recursive, + column_names=column_names, + materialized=materialized, + ) + + @property + def is_recursive(self) -> bool: + """ + ``True`` if a ``recursive_body`` was provided. + """ + return self._recursive + + @property + def materialized(self) -> Optional[bool]: + return self._materialized + + def ref(self) -> CTERef: + """ + Returns a :class:`CTERef `, which can be used to + reference the columns of the CTE. ``column_names`` must have been + provided. + """ + if self.column_names is None: + raise CTEError( + f"CTE {self.name!r} has no declared column_names; " + "cannot build a reference for unnamed columns." + ) + return CTERef(self.name, self.column_names) + + def _body_querystring( + self, body: Union["Query", QueryString] + ) -> QueryString: + from piccolo.query.base import Query as _Query + + if isinstance(body, QueryString): + return body + if isinstance(body, _Query): + qss = body.default_querystrings + if len(qss) != 1: + raise CTEError( + "CTE body must compile to exactly one statement; " + f"got {len(qss)}." + ) + return qss[0] + raise CTEError( + "CTE body must be a Query or QueryString instance, got " + f"{type(body).__name__}." + ) + + def get_inner_querystring(self) -> QueryString: + """ + Returns the definition of the CTE, as it appears within the ``WITH`` + clause. + """ + body_qs = self._body_querystring(self.body) + if self._recursive: + assert self.recursive_body is not None + recursive_qs = self._body_querystring(self.recursive_body) + inner = QueryString( + "{} UNION ALL {}", + body_qs, + recursive_qs, + ) + else: + inner = body_qs + cols = "" + if self.column_names is not None: + cols = " (" + ", ".join(f'"{c}"' for c in self.column_names) + ")" + if self._materialized is None: + prefix = f'"{self.name}"{cols} AS ' + return QueryString("{}({})", QueryString(prefix), inner) + materialized_word = ( + "MATERIALIZED" if self._materialized else "NOT MATERIALIZED" + ) + prefix = f'"{self.name}"{cols} AS {materialized_word} ' + return QueryString("{}({})", QueryString(prefix), inner) + + +def build_with_clause_querystring( + ctes: Sequence[CTE], +) -> Optional[QueryString]: + if not ctes: + return None + any_recursive = any(c.is_recursive for c in ctes) + keyword = "WITH RECURSIVE " if any_recursive else "WITH " + template_parts = [] + args: list[QueryString] = [] + for index, cte in enumerate(ctes): + if index > 0: + template_parts.append(", ") + template_parts.append("{}") + args.append(cte.get_inner_querystring()) + template = keyword + "".join(template_parts) + " " + return QueryString(template, *args) diff --git a/piccolo/query/methods/delete.py b/piccolo/query/methods/delete.py index ea7d56a7e..c09431e74 100644 --- a/piccolo/query/methods/delete.py +++ b/piccolo/query/methods/delete.py @@ -5,11 +5,12 @@ from piccolo.custom_types import Combinable from piccolo.query.base import Query -from piccolo.query.mixins import ReturningDelegate, WhereDelegate +from piccolo.query.mixins import ReturningDelegate, WhereDelegate, WithDelegate from piccolo.querystring import QueryString if TYPE_CHECKING: # pragma: no cover from piccolo.columns import Column + from piccolo.query.methods.cte import CTE from piccolo.table import Table @@ -23,6 +24,7 @@ class Delete(Query): "force", "returning_delegate", "where_delegate", + "with_delegate", ) def __init__(self, table: type[Table], force: bool = False, **kwargs): @@ -30,6 +32,7 @@ def __init__(self, table: type[Table], force: bool = False, **kwargs): self.force = force self.returning_delegate = ReturningDelegate() self.where_delegate = WhereDelegate() + self.with_delegate = WithDelegate() def where(self: Self, *where: Union[Combinable, QueryString]) -> Self: self.where_delegate.where(*where) @@ -39,6 +42,10 @@ def returning(self: Self, *columns: Column) -> Self: self.returning_delegate.returning(columns) return self + def with_(self: Self, *ctes: CTE) -> Self: + self.with_delegate.with_(*ctes) + return self + def _validate(self): """ Don't let a deletion happen unless it has a where clause, or is @@ -72,6 +79,10 @@ def default_querystrings(self) -> Sequence[QueryString]: self.returning_delegate._returning.querystring, ) + with_qs = self.with_delegate.querystring + if with_qs is not None: + querystring = QueryString("{}{}", with_qs, querystring) + return [querystring] diff --git a/piccolo/query/methods/insert.py b/piccolo/query/methods/insert.py index f9bce9516..a2b01788f 100644 --- a/piccolo/query/methods/insert.py +++ b/piccolo/query/methods/insert.py @@ -18,18 +18,25 @@ OnConflictAction, OnConflictDelegate, ReturningDelegate, + WithDelegate, ) from piccolo.querystring import QueryString if TYPE_CHECKING: # pragma: no cover from piccolo.columns.base import Column + from piccolo.query.methods.cte import CTE from piccolo.table import Table class Insert( Generic[TableInstance], Query[TableInstance, list[dict[str, Any]]] ): - __slots__ = ("add_delegate", "on_conflict_delegate", "returning_delegate") + __slots__ = ( + "add_delegate", + "on_conflict_delegate", + "returning_delegate", + "with_delegate", + ) def __init__( self, table: type[TableInstance], *instances: TableInstance, **kwargs @@ -38,8 +45,13 @@ def __init__( self.add_delegate = AddDelegate() self.returning_delegate = ReturningDelegate() self.on_conflict_delegate = OnConflictDelegate() + self.with_delegate = WithDelegate() self.add(*instances) + def with_(self: Self, *ctes: CTE) -> Self: + self.with_delegate.with_(*ctes) + return self + ########################################################################### # Clauses @@ -136,15 +148,23 @@ def default_querystrings(self) -> Sequence[QueryString]: ): returning = self.returning_delegate._returning if returning: - return [ - QueryString( - "{}{}", - querystring, - returning.querystring, - query_type="insert", - table=self.table, - ) - ] + querystring = QueryString( + "{}{}", + querystring, + returning.querystring, + query_type="insert", + table=self.table, + ) + + with_qs = self.with_delegate.querystring + if with_qs is not None: + querystring = QueryString( + "{}{}", + with_qs, + querystring, + query_type="insert", + table=self.table, + ) return [querystring] diff --git a/piccolo/query/methods/select.py b/piccolo/query/methods/select.py index b982c31d2..35e781161 100644 --- a/piccolo/query/methods/select.py +++ b/piccolo/query/methods/select.py @@ -34,6 +34,7 @@ OrderByDelegate, OutputDelegate, WhereDelegate, + WithDelegate, ) from piccolo.query.proxy import Proxy from piccolo.querystring import QueryString @@ -43,6 +44,7 @@ if TYPE_CHECKING: # pragma: no cover from piccolo.custom_types import Combinable + from piccolo.query.methods.cte import CTE from piccolo.table import Table # noqa # Here to avoid breaking changes - will be removed in the future. @@ -160,6 +162,7 @@ class Select(Query[TableInstance, list[dict[str, Any]]]): "where_delegate", "having_delegate", "lock_rows_delegate", + "with_delegate", ) def __init__( @@ -186,9 +189,14 @@ def __init__( self.where_delegate = WhereDelegate() self.having_delegate = WhereDelegate() self.lock_rows_delegate = LockRowsDelegate() + self.with_delegate = WithDelegate() self.columns(*columns_list) + def with_(self: Self, *ctes: CTE) -> Self: + self.with_delegate.with_(*ctes) + return self + def columns(self: Self, *columns: Union[Selectable, str]) -> Self: _columns = self.table._process_column_args(*columns) self.columns_delegate.columns(*_columns) @@ -683,6 +691,10 @@ def default_querystrings(self) -> Sequence[QueryString]: querystring = QueryString(query, *args) + with_qs = self.with_delegate.querystring + if with_qs is not None: + querystring = QueryString("{}{}", with_qs, querystring) + return [querystring] async def run( diff --git a/piccolo/query/methods/update.py b/piccolo/query/methods/update.py index 97715b7d9..20f8635fa 100644 --- a/piccolo/query/methods/update.py +++ b/piccolo/query/methods/update.py @@ -9,11 +9,13 @@ ReturningDelegate, ValuesDelegate, WhereDelegate, + WithDelegate, ) from piccolo.querystring import QueryString if TYPE_CHECKING: # pragma: no cover from piccolo.columns import Column + from piccolo.query.methods.cte import CTE class UpdateError(Exception): @@ -26,6 +28,7 @@ class Update(Query[TableInstance, list[Any]]): "returning_delegate", "values_delegate", "where_delegate", + "with_delegate", ) def __init__( @@ -36,6 +39,11 @@ def __init__( self.returning_delegate = ReturningDelegate() self.values_delegate = ValuesDelegate(table=table) self.where_delegate = WhereDelegate() + self.with_delegate = WithDelegate() + + def with_(self, *ctes: CTE) -> Update: + self.with_delegate.with_(*ctes) + return self ########################################################################### # Clauses @@ -115,4 +123,8 @@ def default_querystrings(self) -> Sequence[QueryString]: self.returning_delegate._returning.querystring, ) + with_qs = self.with_delegate.querystring + if with_qs is not None: + querystring = QueryString("{}{}", with_qs, querystring) + return [querystring] diff --git a/piccolo/query/mixins.py b/piccolo/query/mixins.py index 6d973d60f..44a091534 100644 --- a/piccolo/query/mixins.py +++ b/piccolo/query/mixins.py @@ -17,6 +17,7 @@ from piccolo.utils.sql_values import convert_to_sql_value if TYPE_CHECKING: # pragma: no cover + from piccolo.query.methods.cte import CTE from piccolo.querystring import Selectable from piccolo.table import Table # noqa @@ -891,3 +892,37 @@ def lock_rows( raise ValueError("Unrecognised `lock_strength` value.") self._lock_rows = LockRows(lock_strength_, nowait, skip_locked, of) + + +@dataclass +class WithDelegate: + """ + Stores the ``CTE`` objects which are used to build a ``WITH`` clause. + """ + + _ctes: list[CTE] = field(default_factory=list) + + def with_(self, *ctes: CTE) -> None: + from piccolo.query.methods.cte import CTE, CTEError + + if len(ctes) == 0: + raise CTEError("`with_` requires at least one CTE.") + + names = {cte.name for cte in self._ctes} + for cte in ctes: + if not isinstance(cte, CTE): + raise CTEError( + f"`with_` arguments must be CTE instances, got " + f"{type(cte).__name__}." + ) + if cte.name in names: + raise CTEError(f"Duplicate CTE name {cte.name!r}.") + names.add(cte.name) + + self._ctes.extend(ctes) + + @property + def querystring(self) -> Optional[QueryString]: + from piccolo.query.methods.cte import build_with_clause_querystring + + return build_with_clause_querystring(self._ctes) diff --git a/tests/query/test_cte.py b/tests/query/test_cte.py new file mode 100644 index 000000000..e180f0d52 --- /dev/null +++ b/tests/query/test_cte.py @@ -0,0 +1,393 @@ +from unittest import TestCase + +from piccolo.columns import Integer, Varchar +from piccolo.query import CTE, CTEColumn, CTEError, CTERef +from piccolo.querystring import QueryString +from piccolo.table import Table, create_db_tables_sync, drop_db_tables_sync + + +class Employee(Table): + name = Varchar() + manager_id = Integer(null=True) + + +def compile_query(query) -> str: + sql, _ = query.querystrings[0].compile_string("sqlite") + return sql + + +def compile_cte(cte: CTE) -> str: + sql, _ = cte.get_inner_querystring().compile_string("sqlite") + return sql + + +def counter(name: str = "counter", limit: int = 3) -> CTE: + return CTE.recursive( + name, + base=QueryString("SELECT 1"), + recursive=QueryString( + f'SELECT "n" + 1 FROM "{name}" WHERE "n" < {limit}' + ), + column_names=["n"], + ) + + +class TestCTEConstruction(TestCase): + def test_properties(self): + cte = CTE("numbers", QueryString("SELECT 1")) + self.assertEqual(cte.name, "numbers") + self.assertFalse(cte.is_recursive) + self.assertIsNone(cte.materialized) + + self.assertTrue(counter().is_recursive) + self.assertEqual(counter().column_names, ("n",)) + + def test_query_body(self): + cte = CTE( + "managers", + Employee.select(Employee.id, Employee.name).where( + Employee.manager_id.is_null() + ), + column_names=["id", "name"], + ) + self.assertEqual( + compile_cte(cte), + '"managers" ("id", "name") AS (SELECT ALL "employee"."id" AS ' + '"id", "employee"."name" AS "name" FROM "employee" WHERE ' + '"employee"."manager_id" IS NULL)', + ) + + def test_invalid_names(self): + for name in ("", "1abc", "has space", 'quo"te', "drop;table", 5): + with self.subTest(name=name): + with self.assertRaises(CTEError): + CTE(name, QueryString("SELECT 1")) + + for column_name in ("", "1abc", "a-b", None): + with self.subTest(column_name=column_name): + with self.assertRaises(CTEError): + CTE( + "c", + QueryString("SELECT 1"), + column_names=[column_name], + ) + + def test_invalid_column_names(self): + with self.assertRaises(CTEError): + CTE("c", QueryString("SELECT 1"), column_names=[]) + + with self.assertRaises(CTEError): + CTE("c", QueryString("SELECT 1, 2"), column_names=["a", "a"]) + + def test_invalid_materialized(self): + with self.assertRaises(CTEError): + CTE("c", QueryString("SELECT 1"), materialized="yes") + + def test_column_count_mismatch(self): + with self.assertRaises(CTEError): + CTE( + "c", + Employee.select(Employee.id, Employee.name), + column_names=["id"], + ) + + with self.assertRaises(CTEError): + CTE.recursive( + "c", + base=Employee.select(Employee.id, Employee.name), + recursive=Employee.select(Employee.id), + column_names=["id", "name"], + ) + + def test_recursive_requires_column_names(self): + with self.assertRaises(CTEError): + CTE( + "c", + QueryString("SELECT 1"), + recursive_body=QueryString("SELECT 2"), + ) + + def test_invalid_body_type(self): + cte = CTE("c", 123) + with self.assertRaises(CTEError): + cte.get_inner_querystring() + + +class TestCTEInnerSQL(TestCase): + def test_non_recursive(self): + cte = CTE("numbers", QueryString("SELECT 42 AS n")) + self.assertEqual(compile_cte(cte), '"numbers" AS (SELECT 42 AS n)') + + def test_recursive(self): + self.assertEqual( + compile_cte(counter()), + '"counter" ("n") AS (SELECT 1 UNION ALL SELECT "n" + 1 FROM ' + '"counter" WHERE "n" < 3)', + ) + + def test_column_order(self): + cte = CTE( + "c", + QueryString("SELECT 1, 2, 3"), + column_names=["z", "a", "m"], + ) + self.assertEqual( + compile_cte(cte), '"c" ("z", "a", "m") AS (SELECT 1, 2, 3)' + ) + + def test_materialized(self): + for materialized, keyword in ( + (True, "MATERIALIZED "), + (False, "NOT MATERIALIZED "), + (None, ""), + ): + with self.subTest(materialized=materialized): + cte = CTE( + "c", + QueryString("SELECT 1"), + column_names=["v"], + materialized=materialized, + ) + self.assertIs(cte.materialized, materialized) + self.assertEqual( + compile_cte(cte), + f'"c" ("v") AS {keyword}(SELECT 1)', + ) + + +class TestCTERef(TestCase): + def setUp(self): + self.ref = CTE( + "tree", + QueryString("SELECT 1, 2, 3"), + column_names=["id", "parent", "depth"], + ).ref() + + def test_requires_column_names(self): + with self.assertRaises(CTEError): + CTE("c", QueryString("SELECT 1")).ref() + + def test_column(self): + self.assertIsInstance(self.ref, CTERef) + column = self.ref.column("parent") + self.assertIsInstance(column, CTEColumn) + self.assertEqual(str(column), '"tree"."parent"') + + with self.assertRaises(CTEError): + self.ref.column("missing") + + def test_attribute_access(self): + self.assertEqual(str(self.ref.depth), '"tree"."depth"') + + with self.assertRaises(AttributeError): + self.ref.missing + + def test_collection_protocol(self): + self.assertEqual(len(self.ref), 3) + self.assertIn("id", self.ref) + self.assertNotIn("missing", self.ref) + self.assertNotIn(42, self.ref) + self.assertEqual( + [str(column) for column in self.ref], + ['"tree"."id"', '"tree"."parent"', '"tree"."depth"'], + ) + self.assertEqual( + [str(column) for column in self.ref.columns()], + ['"tree"."id"', '"tree"."parent"', '"tree"."depth"'], + ) + + def test_from_clause(self): + sql, _ = self.ref.from_clause().compile_string("sqlite") + self.assertEqual(sql, '"tree"') + + def test_select_all(self): + sql, _ = self.ref.select_all().compile_string("sqlite") + self.assertEqual( + sql, + 'SELECT "tree"."id", "tree"."parent", "tree"."depth" ' + 'FROM "tree"', + ) + + def test_column_get_select_string(self): + sql, _ = self.ref.id.get_select_string("sqlite").compile_string( + "sqlite" + ) + self.assertEqual(sql, '"tree"."id"') + + +class TestWith(TestCase): + def test_all_query_types(self): + cte = CTE("numbers", QueryString("SELECT 1")) + + queries = { + "SELECT": Employee.select(), + "UPDATE": Employee.update({Employee.name: "x"}, force=True), + "DELETE": Employee.delete(force=True), + "INSERT": Employee.insert(Employee(name="x")), + } + for keyword, query in queries.items(): + with self.subTest(keyword=keyword): + self.assertIs(query.with_(cte), query) + sql = compile_query(query) + self.assertTrue( + sql.startswith('WITH "numbers" AS (SELECT 1) ' + keyword) + ) + + def test_multiple_ctes_keep_order(self): + query = Employee.select().with_( + CTE("b", QueryString("SELECT 2")), + CTE("a", QueryString("SELECT 1")), + ) + self.assertTrue( + compile_query(query).startswith( + 'WITH "b" AS (SELECT 2), "a" AS (SELECT 1) SELECT' + ) + ) + + def test_with_called_twice(self): + query = ( + Employee.select() + .with_(CTE("a", QueryString("SELECT 1"))) + .with_(CTE("b", QueryString("SELECT 2"))) + ) + self.assertTrue( + compile_query(query).startswith( + 'WITH "a" AS (SELECT 1), "b" AS (SELECT 2) SELECT' + ) + ) + + def test_recursive_keyword(self): + plain = CTE("plain", QueryString("SELECT 1")) + + for ctes in ((counter(),), (plain, counter())): + with self.subTest(count=len(ctes)): + sql = compile_query(Employee.select().with_(*ctes)) + self.assertTrue(sql.startswith("WITH RECURSIVE ")) + + sql = compile_query(Employee.delete(force=True).with_(counter())) + self.assertTrue(sql.startswith("WITH RECURSIVE ")) + + sql = compile_query(Employee.select().with_(plain)) + self.assertFalse(sql.startswith("WITH RECURSIVE ")) + + def test_invalid_arguments(self): + query = Employee.select() + + with self.assertRaises(CTEError): + query.with_() + + with self.assertRaises(CTEError): + query.with_("numbers") + + with self.assertRaises(CTEError): + query.with_( + CTE("dup", QueryString("SELECT 1")), + CTE("dup", QueryString("SELECT 2")), + ) + + query.with_(CTE("dup", QueryString("SELECT 1"))) + with self.assertRaises(CTEError): + query.with_(CTE("dup", QueryString("SELECT 2"))) + + +class TestWithExecution(TestCase): + def setUp(self): + create_db_tables_sync(Employee) + Employee.insert( + *( + Employee( + { + Employee.id: id_, + Employee.name: name, + Employee.manager_id: manager_id, + } + ) + for id_, name, manager_id in ( + (1, "ceo", None), + (2, "vp", 1), + (3, "lead", 2), + (4, "engineer", 3), + (5, "contractor", None), + ) + ) + ).run_sync() + + def tearDown(self): + drop_db_tables_sync(Employee) + + @staticmethod + def reports() -> CTE: + return CTE.recursive( + "reports", + base=QueryString( + 'SELECT "id", 0 FROM "employee" WHERE "id" = {}', 2 + ), + recursive=QueryString( + 'SELECT "employee"."id", "reports"."depth" + 1 ' + 'FROM "employee" JOIN "reports" ' + 'ON "employee"."manager_id" = "reports"."id"' + ), + column_names=["id", "depth"], + ) + + def test_select(self): + response = ( + Employee.select(Employee.name) + .where( + Employee.id.is_in(QueryString('SELECT "id" FROM "reports"')) + ) + .order_by(Employee.id) + .with_(self.reports()) + .run_sync() + ) + self.assertEqual( + response, + [{"name": "vp"}, {"name": "lead"}, {"name": "engineer"}], + ) + + def test_select_multiple_ctes(self): + ref = self.reports().ref() + deep = CTE( + "deep", + QueryString( + f"SELECT {ref.id} FROM {ref.from_clause()} WHERE " + f"{ref.depth} > 0" + ), + column_names=["id"], + ) + response = ( + Employee.select(Employee.name) + .where(Employee.id.is_in(QueryString('SELECT "id" FROM "deep"'))) + .order_by(Employee.id) + .with_(self.reports(), deep) + .run_sync() + ) + self.assertEqual(response, [{"name": "lead"}, {"name": "engineer"}]) + + def test_update(self): + Employee.update({Employee.name: "moved"}).where( + Employee.id.is_in(QueryString('SELECT "id" FROM "reports"')) + ).with_(self.reports()).run_sync() + + names = Employee.select(Employee.name).order_by(Employee.id).run_sync() + self.assertEqual( + [row["name"] for row in names], + ["ceo", "moved", "moved", "moved", "contractor"], + ) + + def test_delete(self): + Employee.delete().where( + Employee.id.is_in(QueryString('SELECT "id" FROM "reports"')) + ).with_(self.reports()).run_sync() + + names = Employee.select(Employee.name).order_by(Employee.id).run_sync() + self.assertEqual([row["name"] for row in names], ["ceo", "contractor"]) + + def test_insert(self): + Employee.insert( + Employee({Employee.id: 6, Employee.name: "intern"}) + ).with_(CTE("numbers", QueryString("SELECT 1"))).run_sync() + + self.assertEqual( + Employee.count().where(Employee.name == "intern").run_sync(), 1 + )