Feat(tsql, oracle): add support for NEXT VALUE FOR clause (#1521)

tobymao · May 2, 2023 · 20cacba · 20cacba
1 parent c3db2b8
commit 20cacba
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 1 deletion.
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -4075,6 +4075,12 @@ class When(Func):
     arg_types = {"matched": True, "source": False, "condition": False, "then": True}
 
 
+# https://docs.oracle.com/javadb/10.8.3.0/ref/rrefsqljnextvaluefor.html
+# https://learn.microsoft.com/en-us/sql/t-sql/functions/next-value-for-transact-sql?view=sql-server-ver16
+class NextValueFor(Func):
+    arg_types = {"this": True, "order": False}
+
+
 def _norm_arg(arg):
     return arg.lower() if type(arg) is str else arg
 

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -1629,6 +1629,11 @@ def constraint_sql(self, expression: exp.Constraint) -> str:
         expressions = self.expressions(expression, flat=True)
         return f"CONSTRAINT {this} {expressions}"
 
+    def nextvaluefor_sql(self, expression: exp.NextValueFor) -> str:
+        order = expression.args.get("order")
+        order = f" OVER ({self.order_sql(order, flat=True)})" if order else ""
+        return f"NEXT VALUE FOR {self.sql(expression, 'this')}{order}"
+
     def extract_sql(self, expression: exp.Extract) -> str:
         this = self.sql(expression, "this")
         expression_sql = self.sql(expression, "expression")

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -685,9 +685,14 @@ class Parser(metaclass=_Parser):
     SCHEMA_UNNAMED_CONSTRAINTS = {"CHECK", "FOREIGN KEY", "LIKE", "PRIMARY KEY", "UNIQUE"}
 
     NO_PAREN_FUNCTION_PARSERS = {
+        TokenType.ANY: lambda self: self.expression(exp.Any, this=self._parse_bitwise()),
         TokenType.CASE: lambda self: self._parse_case(),
         TokenType.IF: lambda self: self._parse_if(),
-        TokenType.ANY: lambda self: self.expression(exp.Any, this=self._parse_bitwise()),
+        TokenType.NEXT_VALUE_FOR: lambda self: self.expression(
+            exp.NextValueFor,
+            this=self._parse_column(),
+            order=self._match(TokenType.OVER) and self._parse_wrapped(self._parse_order),
+        ),
     }
 
     FUNCTION_PARSERS: t.Dict[str, t.Callable] = {

diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
@@ -232,6 +232,7 @@ class TokenType(AutoName):
     MOD = auto()
     NATURAL = auto()
     NEXT = auto()
+    NEXT_VALUE_FOR = auto()
     NO_ACTION = auto()
     NOTNULL = auto()
     NULL = auto()
@@ -573,6 +574,7 @@ class Tokenizer(metaclass=_Tokenizer):
         "MERGE": TokenType.MERGE,
         "NATURAL": TokenType.NATURAL,
         "NEXT": TokenType.NEXT,
+        "NEXT VALUE FOR": TokenType.NEXT_VALUE_FOR,
         "NO ACTION": TokenType.NO_ACTION,
         "NOT": TokenType.NOT,
         "NOTNULL": TokenType.NOTNULL,

diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
@@ -818,3 +818,5 @@ JSON_OBJECT('x': NULL, 'y': 1 ABSENT ON NULL WITH UNIQUE KEYS)
 JSON_OBJECT('x': 1 RETURNING VARCHAR(100))
 JSON_OBJECT('x': 1 RETURNING VARBINARY FORMAT JSON ENCODING UTF8)
 SELECT if.x
+SELECT NEXT VALUE FOR db.schema.sequence_name
+SELECT NEXT VALUE FOR db.schema.sequence_name OVER (ORDER BY foo), col