feat(oracle): parse WITH FUNCTION PL/SQL declarations [CLAUDE]

sqlglot/generators/oracle.py

@@ -145,6 +145,30 @@ def version_sql(self, expression: exp.Version) -> str:
     def offset_sql(self, expression: exp.Offset) -> str:
         return f"{super().offset_sql(expression)} ROWS"
 
+    def with_sql(self, expression: exp.With) -> str:
+        if not expression.args.get("udfs"):
+            return super().with_sql(expression)
+
+        functions = self.expressions(expression, key="udfs", sep=" ", flat=True)
+        ctes = self.expressions(expression, flat=True)
+        search = self.sql(expression, "search")
+        search = f" {search}" if search else ""
+        cycle = self.sql(expression, "cycle")
+        cycle = f" {cycle}" if cycle else ""
+
+        return f"WITH {functions} {ctes}{search}{cycle}" if ctes else f"WITH {functions}"
+
+    def functionspecification_sql(self, expression: exp.FunctionSpecification) -> str:
+        characteristics = expression.args.get("characteristics")
+        returns = characteristics and characteristics.find(exp.ReturnsProperty)
+        returns_sql = f" RETURN {self.sql(returns, 'this')}" if returns else ""
+
+        body = expression.expression
+        if not isinstance(body, exp.Var):
+            self.unsupported("Inline function bodies must be PL/SQL blocks in Oracle")
+
+        return f"FUNCTION {self.sql(expression, 'this')}{returns_sql} IS {self.sql(body)};"
+
     def add_column_sql(self, expression: exp.Expr) -> str:
         return f"ADD {self.sql(expression)}"
 

sqlglot/parsers/oracle.py

@@ -6,7 +6,7 @@
 from sqlglot.dialects.dialect import build_formatted_time, build_timetostr_or_tochar, build_trunc
 from sqlglot.helper import ensure_list, seq_get
 from sqlglot.parser import OPTIONS_TYPE, build_coalesce
-from sqlglot.tokens import TokenType
+from sqlglot.tokens import Token, TokenType
 
 if t.TYPE_CHECKING:
     from sqlglot._typing import E
@@ -136,6 +136,137 @@ class OracleParser(parser.Parser):
         ),
     }
 
+    def _parse_with(self, skip_with_token: bool = False) -> exp.With | None:
+        # https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/SELECT.html
+        # PL/SQL function declarations precede the CTEs of the same WITH clause, and each one is
+        # terminated by its own `END;` rather than separated by a comma
+        if not skip_with_token and not self._match(TokenType.WITH):
+            return None
+
+        comments = self._prev_comments
+        udfs = []
+
+        while (
+            self._match(TokenType.FUNCTION, advance=False)
+            and self._next
+            and self._next.token_type in self.ID_VAR_TOKENS
+        ):
+            self._advance()
+            udfs.append(self._parse_plsql_function())
+
+        if not udfs:
+            return super()._parse_with(skip_with_token=True)
+
+        if self._match_set(self.ID_VAR_TOKENS, advance=False):
+            with_ = super()._parse_with(skip_with_token=True)
+        else:
+            with_ = self.expression(exp.With())
+
+        if with_:
+            with_.set("udfs", udfs)
+            with_.add_comments(comments)
+
+        return with_
+
+    def _parse_plsql_function(self) -> exp.FunctionSpecification:
+        this = self._parse_user_defined_function(kind=TokenType.FUNCTION)
+
+        if not self._match_text_seq("RETURN"):
+            self.raise_error("Expected RETURN")
+
+        returns = self.expression(exp.ReturnsProperty(this=self._parse_types()))
+
+        if not self._match_set((TokenType.IS, TokenType.ALIAS)):
+            self.raise_error("Expected IS or AS")
+
+        # The body is kept verbatim, since there is no PL/SQL statement grammar to parse it with
+        start = self._prev
+        end = self._advance_plsql_body()
+        body = self.sql[start.end + 1 : end.end + 1].strip() if end else ""
+
+        return self.expression(
+            exp.FunctionSpecification(
+                this=this,
+                characteristics=self.expression(exp.Properties(expressions=[returns])),
+                expression=exp.var(body),
+            )
+        )
+
+    def _advance_plsql_token(self) -> bool:
+        # The body's own `;` separators split it across statement chunks, so stepping through
+        # it has to carry on into the next chunk whenever the current one runs out
+        self._advance()
+
+        while not self._curr or self._curr.token_type == TokenType.SEMICOLON:
+            if self._curr:
+                self._advance()
+            elif self._chunk_index < len(self._chunks):
+                self._advance_chunk()
+            else:
+                return False
+
+        return True
+
+    def _advance_plsql_body(self) -> Token | None:
+        """
+        Advances past a subprogram body (its declarations, then BEGIN ... END [name];) and
+        returns the body's last token, leaving the parser on the token after the closing `;`.
+        """
+        while self._curr and self._curr.token_type != TokenType.BEGIN:
+            if self._curr.token_type in (TokenType.FUNCTION, TokenType.PROCEDURE):
+                # A nested subprogram is either a forward declaration, which ends at the next
+                # `;`, or a definition whose IS/AS is followed by a body of its own
+                chunk_index = self._chunk_index
+                while (
+                    self._advance_plsql_token()
+                    and self._chunk_index == chunk_index
+                    and not self._match_set((TokenType.IS, TokenType.ALIAS), advance=False)
+                ):
+                    pass
+
+                if self._curr and self._chunk_index == chunk_index:
+                    self._advance()
+                    self._advance_plsql_body()
+
+                continue
+
+            self._advance_plsql_token()
+
+        # END IF and END LOOP close blocks that are not counted here, whereas END CASE and a
+        # bare END close a BEGIN or a CASE
+        depth = 0
+        while self._curr:
+            token_type = self._curr.token_type
+
+            if token_type in (TokenType.BEGIN, TokenType.CASE):
+                depth += 1
+            elif token_type == TokenType.END:
+                closes = self._next.text.upper() if self._next else ""
+
+                if closes in ("IF", "LOOP"):
+                    self._advance()
+                else:
+                    if closes == "CASE":
+                        self._advance()
+
+                    depth -= 1
+
+                    if not depth:
+                        if self._next:
+                            self._advance()
+
+                        end = self._curr
+                        if self._next:
+                            self.raise_error("Expected ; after END")
+
+                        self._advance_plsql_token()
+                        return end
+
+            self._advance_plsql_token()
+
+        self.raise_error("Expected END to close the PL/SQL body")
+        return None
+
     def _parse_physical_property(self) -> exp.PhysicalProperty:
         name = self._prev.text.upper()
         value = self._parse_number()

tests/dialects/test_oracle.py

@@ -897,6 +897,93 @@ def test_recursive_cte(self):
         self.assertEqual(with_.args["cycle"].args["to"], exp.Literal.string("Y"))
         self.assertEqual(with_.args["cycle"].args["default"], exp.Literal.string("N"))
 
+    def test_with_function(self):
+        select = self.validate_identity(
+            """WITH FUNCTION f_demo(v_in_ch IN CHAR) RETURN CHAR
+IS
+BEGIN
+    RETURN v_in_ch;
+END;
+
+SELECT
+    'C' || f_demo('d')
+FROM dual""",
+            "WITH FUNCTION f_demo(v_in_ch IN CHAR) RETURN CHAR IS BEGIN\n    RETURN v_in_ch;\nEND; SELECT 'C' || F_DEMO('d') FROM dual",
+        )
+        udf = select.args["with_"].args["udfs"][0]
+        self.assertIsInstance(udf, exp.FunctionSpecification)
+        self.assertEqual(udf.find(exp.ReturnsProperty).this.sql("oracle"), "CHAR")
+        self.assertEqual(udf.expression.name, "BEGIN\n    RETURN v_in_ch;\nEND")
+
+        self.validate_identity(
+            "WITH FUNCTION a(x NUMBER) RETURN NUMBER IS BEGIN RETURN x + 1; END; "
+            "FUNCTION b RETURN NUMBER IS BEGIN RETURN 2; END b; "
+            "t AS (SELECT 1 AS v FROM dual) SELECT A(v) + b FROM t"
+        )
+        self.validate_identity(
+            "WITH FUNCTION f RETURN NUMBER AS BEGIN RETURN 1; END; SELECT f FROM dual",
+            "WITH FUNCTION f RETURN NUMBER IS BEGIN RETURN 1; END; SELECT f FROM dual",
+        )
+        self.validate_identity(
+            "WITH FUNCTION f RETURN NUMBER IS BEGIN RETURN 1; END; SELECT f FROM dual WITH READ ONLY"
+        )
+
+        # The END that closes the function is found past nested blocks, END IF / END LOOP /
+        # END CASE, a CASE expression, and statement separators inside strings and comments
+        self.validate_identity(
+            "WITH FUNCTION f(n NUMBER) RETURN VARCHAR2 IS "
+            "s VARCHAR2(20) := 'END; BEGIN'; "
+            "BEGIN "
+            "IF n > 0 THEN s := 'pos'; ELSIF n < 0 THEN s := 'neg'; ELSE s := 'zero'; END IF; "
+            "<<outer>> FOR i IN 1..3 LOOP s := s || i; END LOOP outer; "
+            "CASE n WHEN 1 THEN s := s || 'one'; ELSE NULL; END CASE; "
+            "s := s || CASE WHEN n = 2 THEN 'two' ELSE '' END; "
+            "BEGIN NULL; EXCEPTION WHEN OTHERS THEN NULL; END; "
+            "RETURN s /* END; */; "
+            "END f; "
+            "SELECT F(1) FROM dual"
+        )
+
+        # A nested subprogram, both forward declared and defined, has a BEGIN ... END of its own
+        self.validate_identity(
+            "WITH FUNCTION outer_f(x NUMBER) RETURN NUMBER IS "
+            "FUNCTION helper(y NUMBER) RETURN NUMBER; "
+            "FUNCTION helper(y NUMBER) RETURN NUMBER IS BEGIN RETURN y * 2; END helper; "
+            "BEGIN RETURN helper(x); END outer_f; "
+            "SELECT OUTER_F(3) FROM dual"
+        )
+
+        expressions = parse(
+            "WITH FUNCTION f RETURN NUMBER IS BEGIN RETURN 1; END; SELECT f FROM dual; SELECT 2 FROM dual",
+            read="oracle",
+        )
+        self.assertEqual(
+            [e.sql("oracle") for e in expressions],
+            [
+                "WITH FUNCTION f RETURN NUMBER IS BEGIN RETURN 1; END; SELECT f FROM dual",
+                "SELECT 2 FROM dual",
+            ],
+        )
+
+        self.validate_identity(
+            "WITH FUNCTION f RETURN NUMBER IS BEGIN RETURN 1; /* done */ END; SELECT f FROM dual"
+        )
+        self.validate_identity("WITH function AS (SELECT 1 AS x FROM dual) SELECT x FROM function")
+
+        for sql in (
+            "WITH FUNCTION f NUMBER IS BEGIN RETURN 1; END; SELECT f FROM dual",
+            "WITH FUNCTION f RETURN NUMBER BEGIN RETURN 1; END; SELECT f FROM dual",
+            "WITH FUNCTION f RETURN NUMBER IS BEGIN RETURN 1; SELECT f FROM dual",
+            "WITH FUNCTION f RETURN NUMBER IS BEGIN RETURN 1; END f g; SELECT f FROM dual",
+        ):
+            with self.subTest(sql), self.assertRaises(ParseError):
+                parse_one(sql, read="oracle")
+
+        with self.assertRaises(UnsupportedError):
+            parse_one("WITH FUNCTION f(x INT) RETURNS INT RETURN x SELECT f(1)", read="trino").sql(
+                "oracle", unsupported_level=ErrorLevel.RAISE
+            )
+
     def test_query_restrictions(self):
         for restriction in ("READ ONLY", "CHECK OPTION"):
             for constraint_name in (" CONSTRAINT name", ""):