Fix(tsql): do not bubble up CTEs in TVF DDLs with inline RETURN queries closes #7721

Author: georgesittas

sqlglot/generators/tsql.py

@@ -417,6 +417,15 @@ def create_sql(self, expression: exp.Create) -> str:
                 # but CREATE VIEW actually requires the WITH clause to come after it so we need
                 # to amend the AST by moving the CTEs to the CREATE VIEW statement's query.
                 ctas_expression.set("with_", with_.pop())
+        elif (
+            kind == "FUNCTION"
+            and isinstance(ctas_expression, exp.Return)
+            and isinstance(body := ctas_expression.this.unnest(), exp.Query)
+            and (with_ := expression.args.get("with_"))
+        ):
+            # Similar to the VIEW branch, the table-valued functions require the WITH clause
+            # to stay inside the RETURN body, so we move back any CTEs that were bubbled up.
+            body.set("with_", with_.pop())
 
         table = expression.find(exp.Table)
 

tests/dialects/test_tsql.py

@@ -1374,6 +1374,9 @@ def test_udf(self):
 
         self.validate_identity("CREATE FUNCTION foo(@bar INTEGER) RETURNS TABLE AS RETURN SELECT 1")
         self.validate_identity("CREATE FUNCTION dbo.ISOweek(@DATE DATETIME2) RETURNS INTEGER")
+        self.validate_identity(
+            "CREATE FUNCTION dbo.f() RETURNS TABLE AS RETURN (WITH subquery AS (SELECT id AS id FROM subtable) SELECT other_id FROM main_table AS mt INNER JOIN subquery ON subquery.id = mt.other_id)"
+        )
 
         # The following two cases don't necessarily correspond to valid TSQL, but they are used to verify
         # that the syntax RETURNS @return_variable TABLE <table_type_definition> ... is parsed correctly.