Add CTE prefixes

Added support for prefixes to the :class:`.CTE` construct, to allow
support for Postgresql 12 "MATERIALIZED" and "NOT MATERIALIZED" phrases.
Pull request courtesy Marat Sharafutdinov.

Fixes: #5040
Closes: #5043
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/5043
Pull-request-sha: d1b9059a0b6dae8dc2479ac670999b4af07908e0

Change-Id: I2e9cb5d7f85961ec98ee51965de5b3ec4a97be2f

diff --git a/doc/build/changelog/unreleased_13/5040.rst b/doc/build/changelog/unreleased_13/5040.rst
new file mode 100644 (file)
index 0000000..b1d45f0
--- /dev/null
+++ b/doc/build/changelog/unreleased_13/5040.rst
@@ -0,0 +1,11 @@
+.. change::
+    :tags: usecase, postgresql
+    :tickets: 5040
+
+    Added support for prefixes to the :class:`.CTE` construct, to allow
+    support for Postgresql 12 "MATERIALIZED" and "NOT MATERIALIZED" phrases.
+    Pull request courtesy Marat Sharafutdinov.
+
+    .. seealso::
+
+        :meth:`.HasCTE.cte`

diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index 4ec3b93ea7e189b4f4275fe7e08c88a812c9dbfb..807b01c243da99cf0b0d09778253a79e280d5ac9 100644 (file)
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -1851,7 +1851,8 @@ class SQLCompiler(Compiled):
                     kwargs["positional_names"] = self.cte_positional[cte] = []
 
                 assert kwargs.get("subquery", False) is False
-                text += " AS \n(%s)" % (
+                text += " AS %s\n(%s)" % (
+                    self._generate_prefixes(cte, cte._prefixes, **kwargs),
                     cte.element._compiler_dispatch(
                         self, asfrom=True, **kwargs
                     ),

diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index bece0b3c58bc370fd80fafcf44cbc3c52458c8bf..ed7a6c2b982ac421fde95aeae10b26323a13400f 100644 (file)
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -1451,7 +1451,7 @@ class TableSample(AliasedReturnsRows):
             return functions.func.system(self.sampling)
 
 
-class CTE(Generative, HasSuffixes, AliasedReturnsRows):
+class CTE(Generative, HasPrefixes, HasSuffixes, AliasedReturnsRows):
     """Represent a Common Table Expression.
 
     The :class:`.CTE` object is obtained using the
@@ -1469,6 +1469,7 @@ class CTE(Generative, HasSuffixes, AliasedReturnsRows):
             ("_restates", InternalTraversal.dp_clauseelement_unordered_set),
             ("recursive", InternalTraversal.dp_boolean),
         ]
+        + HasPrefixes._traverse_internals
         + HasSuffixes._traverse_internals
     )
 
@@ -1490,11 +1491,14 @@ class CTE(Generative, HasSuffixes, AliasedReturnsRows):
         recursive=False,
         _cte_alias=None,
         _restates=frozenset(),
+        _prefixes=None,
         _suffixes=None,
     ):
         self.recursive = recursive
         self._cte_alias = _cte_alias
         self._restates = _restates
+        if _prefixes:
+            self._prefixes = _prefixes
         if _suffixes:
             self._suffixes = _suffixes
         super(CTE, self)._init(selectable, name=name)
@@ -1526,6 +1530,7 @@ class CTE(Generative, HasSuffixes, AliasedReturnsRows):
             name=name,
             recursive=self.recursive,
             _cte_alias=self,
+            _prefixes=self._prefixes,
             _suffixes=self._suffixes,
         )
 
@@ -1535,6 +1540,7 @@ class CTE(Generative, HasSuffixes, AliasedReturnsRows):
             name=self.name,
             recursive=self.recursive,
             _restates=self._restates.union([self]),
+            _prefixes=self._prefixes,
             _suffixes=self._suffixes,
         )
 
@@ -1544,6 +1550,7 @@ class CTE(Generative, HasSuffixes, AliasedReturnsRows):
             name=self.name,
             recursive=self.recursive,
             _restates=self._restates.union([self]),
+            _prefixes=self._prefixes,
             _suffixes=self._suffixes,
         )
 
@@ -1570,13 +1577,20 @@ class HasCTE(roles.HasCTERole):
         when combined with RETURNING, as well as a consumer of
         CTE rows.
 
+        .. versionchanged:: 1.1 Added support for UPDATE/INSERT/DELETE as
+           CTE, CTEs added to UPDATE/INSERT/DELETE.
+
         SQLAlchemy detects :class:`.CTE` objects, which are treated
         similarly to :class:`.Alias` objects, as special elements
         to be delivered to the FROM clause of the statement as well
         as to a WITH clause at the top of the statement.
 
-        .. versionchanged:: 1.1 Added support for UPDATE/INSERT/DELETE as
-           CTE, CTEs added to UPDATE/INSERT/DELETE.
+        For special prefixes such as PostgreSQL "MATERIALIZED" and
+        "NOT MATERIALIZED", the :meth:`.CTE.prefix_with` method may be
+        used to establish these.
+
+        .. versionchanged:: 1.3.13 Added support for prefixes.
+           In particular - MATERIALIZED and NOT MATERIALIZED.
 
         :param name: name given to the common table expression.  Like
          :meth:`._FromClause.alias`, the name can be left as ``None``

diff --git a/test/sql/test_cte.py b/test/sql/test_cte.py
index 26f367e9f448d3b389709579ecd0d412a4284dd4..ff5a124817aecd4027837c6903514d6e55bdff01 100644 (file)
--- a/test/sql/test_cte.py
+++ b/test/sql/test_cte.py
@@ -903,6 +903,28 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             'ON anon_1."order" = "order"."order"',
         )
 
+    def test_prefixes(self):
+        orders = table("order", column("order"))
+        s = select([orders.c.order]).cte("regional_sales")
+        s = s.prefix_with("NOT MATERIALIZED", dialect="postgresql")
+        stmt = select([orders]).where(orders.c.order > s.c.order)
+
+        self.assert_compile(
+            stmt,
+            'WITH regional_sales AS (SELECT "order"."order" AS "order" '
+            'FROM "order") SELECT "order"."order" FROM "order", '
+            'regional_sales WHERE "order"."order" > regional_sales."order"',
+        )
+
+        self.assert_compile(
+            stmt,
+            'WITH regional_sales AS NOT MATERIALIZED '
+            '(SELECT "order"."order" AS "order" '
+            'FROM "order") SELECT "order"."order" FROM "order", '
+            'regional_sales WHERE "order"."order" > regional_sales."order"',
+            dialect="postgresql"
+        )
+
     def test_suffixes(self):
         orders = table("order", column("order"))
         s = select([orders.c.order]).cte("regional_sales")