Feat(duckdb): support USING KEY (...) in recursive DuckDB CTEs

sqlglot/expressions.py

@@ -1727,6 +1727,7 @@ class CTE(DerivedTable):
         "alias": True,
         "scalar": False,
         "materialized": False,
+        "key_expressions": False,
     }
 
 

sqlglot/generator.py

@@ -1332,7 +1332,10 @@ def cte_sql(self, expression: exp.CTE) -> str:
         elif materialized:
             materialized = "MATERIALIZED "
 
-        return f"{alias_sql} AS {materialized or ''}{self.wrap(expression)}"
+        key_expressions = self.expressions(expression, key="key_expressions", flat=True)
+        key_expressions = f" USING KEY ({key_expressions})" if key_expressions else ""
+
+        return f"{alias_sql}{key_expressions} AS {materialized or ''}{self.wrap(expression)}"
 
     def tablealias_sql(self, expression: exp.TableAlias) -> str:
         alias = self.sql(expression, "this")

sqlglot/parser.py

@@ -3434,6 +3434,10 @@ def _parse_cte(self) -> t.Optional[exp.CTE]:
         if not alias or not alias.this:
             self.raise_error("Expected CTE to have alias")
 
+        key_expressions = (
+            self._parse_wrapped_id_vars() if self._match_text_seq("USING", "KEY") else None
+        )
+
         if not self._match(TokenType.ALIAS) and not self.OPTIONAL_ALIAS_TOKEN_CTE:
             self._retreat(index)
             return None
@@ -3452,6 +3456,7 @@ def _parse_cte(self) -> t.Optional[exp.CTE]:
             this=self._parse_wrapped(self._parse_statement),
             alias=alias,
             materialized=materialized,
+            key_expressions=key_expressions,
             comments=comments,
         )
 

tests/dialects/test_duckdb.py

@@ -1850,3 +1850,11 @@ def test_install(self):
         self.validate_identity("FORCE INSTALL httpfs FROM community")
         self.validate_identity("FORCE INSTALL httpfs FROM 'https://extensions.duckdb.org'")
         self.validate_identity("FORCE CHECKPOINT db", check_command_warning=True)
+
+    def test_cte_using_key(self):
+        self.validate_identity(
+            "WITH RECURSIVE tbl(a, b) USING KEY (a) AS (SELECT a, b FROM (VALUES (1, 3), (2, 4)) AS t(a, b) UNION SELECT a + 1, b FROM tbl WHERE a < 3) SELECT * FROM tbl"
+        )
+        self.validate_identity(
+            "WITH RECURSIVE tbl(a, b) USING KEY (a, b) AS (SELECT a, b FROM (VALUES (1, 3), (2, 4)) AS t(a, b) UNION SELECT a + 1, b FROM tbl WHERE a < 3) SELECT * FROM tbl"
+        )