Add materialized argument to _with (tobymao#4036)

hex-inc · Sep 2, 2024 · aa83ba6 · aa83ba6
1 parent 85cc7ad
commit aa83ba6
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 3 deletions.
diff --git a/Makefile b/Makefile
@@ -7,6 +7,7 @@ install-dev-rs-release:
 	cd sqlglotrs/ && python -m maturin develop -r
 
 install-dev-rs:
+	@unset CONDA_PREFIX && \
 	cd sqlglotrs/ && python -m maturin develop
 
 install-dev-core:

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -1195,6 +1195,7 @@ def with_(
         alias: ExpOrStr,
         as_: ExpOrStr,
         recursive: t.Optional[bool] = None,
+        materialized: t.Optional[bool] = None,
         append: bool = True,
         dialect: DialectType = None,
         copy: bool = True,
@@ -1213,6 +1214,7 @@ def with_(
             as_: the SQL code string to parse as the table expression.
                 If an `Expression` instance is passed, it will be used as-is.
             recursive: set the RECURSIVE part of the expression. Defaults to `False`.
+            materialized: set the MATERIALIZED part of the expression.
             append: if `True`, add to any existing expressions.
                 Otherwise, this resets the expressions.
             dialect: the dialect used to parse the input expression.
@@ -1223,7 +1225,15 @@ def with_(
             The modified expression.
         """
         return _apply_cte_builder(
-            self, alias, as_, recursive=recursive, append=append, dialect=dialect, copy=copy, **opts
+            self,
+            alias,
+            as_,
+            recursive=recursive,
+            materialized=materialized,
+            append=append,
+            dialect=dialect,
+            copy=copy,
+            **opts,
         )
 
     def union(
@@ -2180,6 +2190,7 @@ def with_(
         alias: ExpOrStr,
         as_: ExpOrStr,
         recursive: t.Optional[bool] = None,
+        materialized: t.Optional[bool] = None,
         append: bool = True,
         dialect: DialectType = None,
         copy: bool = True,
@@ -2198,6 +2209,7 @@ def with_(
             as_: the SQL code string to parse as the table expression.
                 If an `Expression` instance is passed, it will be used as-is.
             recursive: set the RECURSIVE part of the expression. Defaults to `False`.
+            materialized: set the MATERIALIZED part of the expression.
             append: if `True`, add to any existing expressions.
                 Otherwise, this resets the expressions.
             dialect: the dialect used to parse the input expression.
@@ -2208,7 +2220,15 @@ def with_(
             The modified expression.
         """
         return _apply_cte_builder(
-            self, alias, as_, recursive=recursive, append=append, dialect=dialect, copy=copy, **opts
+            self,
+            alias,
+            as_,
+            recursive=recursive,
+            materialized=materialized,
+            append=append,
+            dialect=dialect,
+            copy=copy,
+            **opts,
         )
 
 
@@ -6514,14 +6534,15 @@ def _apply_cte_builder(
     alias: ExpOrStr,
     as_: ExpOrStr,
     recursive: t.Optional[bool] = None,
+    materialized: t.Optional[bool] = None,
     append: bool = True,
     dialect: DialectType = None,
     copy: bool = True,
     **opts,
 ) -> E:
     alias_expression = maybe_parse(alias, dialect=dialect, into=TableAlias, **opts)
     as_expression = maybe_parse(as_, dialect=dialect, **opts)
-    cte = CTE(this=as_expression, alias=alias_expression)
+    cte = CTE(this=as_expression, alias=alias_expression, materialized=materialized)
     return _apply_child_list_builder(
         cte,
         instance=instance,

diff --git a/tests/test_build.py b/tests/test_build.py
@@ -358,12 +358,36 @@ def test_build(self):
                 lambda: select("x").from_("tbl").with_("tbl", as_="SELECT x FROM tbl2"),
                 "WITH tbl AS (SELECT x FROM tbl2) SELECT x FROM tbl",
             ),
+            (
+                lambda: select("x")
+                .from_("tbl")
+                .with_("tbl", as_="SELECT x FROM tbl2", materialized=True),
+                "WITH tbl AS MATERIALIZED (SELECT x FROM tbl2) SELECT x FROM tbl",
+            ),
+            (
+                lambda: select("x")
+                .from_("tbl")
+                .with_("tbl", as_="SELECT x FROM tbl2", materialized=False),
+                "WITH tbl AS NOT MATERIALIZED (SELECT x FROM tbl2) SELECT x FROM tbl",
+            ),
             (
                 lambda: select("x")
                 .from_("tbl")
                 .with_("tbl", as_="SELECT x FROM tbl2", recursive=True),
                 "WITH RECURSIVE tbl AS (SELECT x FROM tbl2) SELECT x FROM tbl",
             ),
+            (
+                lambda: select("x")
+                .from_("tbl")
+                .with_("tbl", as_=select("x").from_("tbl2"), recursive=True, materialized=True),
+                "WITH RECURSIVE tbl AS MATERIALIZED (SELECT x FROM tbl2) SELECT x FROM tbl",
+            ),
+            (
+                lambda: select("x")
+                .from_("tbl")
+                .with_("tbl", as_=select("x").from_("tbl2"), recursive=True, materialized=False),
+                "WITH RECURSIVE tbl AS NOT MATERIALIZED (SELECT x FROM tbl2) SELECT x FROM tbl",
+            ),
             (
                 lambda: select("x").from_("tbl").with_("tbl", as_=select("x").from_("tbl2")),
                 "WITH tbl AS (SELECT x FROM tbl2) SELECT x FROM tbl",
@@ -676,6 +700,18 @@ def test_build(self):
                 lambda: exp.insert("SELECT * FROM cte", "t").with_("cte", as_="SELECT x FROM tbl"),
                 "WITH cte AS (SELECT x FROM tbl) INSERT INTO t SELECT * FROM cte",
             ),
+            (
+                lambda: exp.insert("SELECT * FROM cte", "t").with_(
+                    "cte", as_="SELECT x FROM tbl", materialized=True
+                ),
+                "WITH cte AS MATERIALIZED (SELECT x FROM tbl) INSERT INTO t SELECT * FROM cte",
+            ),
+            (
+                lambda: exp.insert("SELECT * FROM cte", "t").with_(
+                    "cte", as_="SELECT x FROM tbl", materialized=False
+                ),
+                "WITH cte AS NOT MATERIALIZED (SELECT x FROM tbl) INSERT INTO t SELECT * FROM cte",
+            ),
             (
                 lambda: exp.convert((exp.column("x"), exp.column("y"))).isin((1, 2), (3, 4)),
                 "(x, y) IN ((1, 2), (3, 4))",