Feat(oracle): parse hints (#4249)

mkmoisen · mmoisen · web-flow · commit 8f5efc7bc01b · 2024-10-17T13:36:13.000+03:00
* feat(oracle)!: support oracle hints

* feat(oracle): support oracle hints

* feat(oracle)!: support oracle hints

---------

Co-authored-by: mmoisen &lt;mmoisen@cisco.com&gt;
diff --git a/sqlglot/dialects/oracle.py b/sqlglot/dialects/oracle.py
@@ -15,6 +15,7 @@
 from sqlglot.helper import seq_get
 from sqlglot.parser import OPTIONS_TYPE, build_coalesce
 from sqlglot.tokens import TokenType
+from sqlglot.errors import ParseError
 
 if t.TYPE_CHECKING:
     from sqlglot._typing import E
@@ -205,6 +206,57 @@ def _parse_json_array(self, expr_type: t.Type[E], **kwargs) -> E:
             )
 
         def _parse_hint(self) -> t.Optional[exp.Hint]:
+            start_index = self._index
+            should_fallback_to_string = False
+
+            if not self._match(TokenType.HINT):
+                return None
+
+            hints = []
+
+            try:
+                for hint in iter(
+                    lambda: self._parse_csv(
+                        lambda: self._parse_hint_function_call() or self._parse_var(upper=True),
+                    ),
+                    [],
+                ):
+                    hints.extend(hint)
+            except ParseError:
+                should_fallback_to_string = True
+
+            if not self._match_pair(TokenType.STAR, TokenType.SLASH):
+                should_fallback_to_string = True
+
+            if should_fallback_to_string:
+                self._retreat(start_index)
+                return self._parse_hint_fallback_to_string()
+
+            return self.expression(exp.Hint, expressions=hints)
+
+        def _parse_hint_function_call(self) -> t.Optional[exp.Expression]:
+            if not self._curr or not self._next or self._next.token_type != TokenType.L_PAREN:
+                return None
+
+            this = self._curr.text
+
+            self._advance(2)
+            args = self._parse_hint_args()
+            this = self.expression(exp.Anonymous, this=this, expressions=args)
+            self._match_r_paren(this)
+            return this
+
+        def _parse_hint_args(self):
+            args = []
+            result = self._parse_var()
+
+            while result:
+                args.append(result)
+                result = self._parse_var()
+
+            return args
+
+        def _parse_hint_fallback_to_string(self) -> t.Optional[exp.Hint]:
             if self._match(TokenType.HINT):
                 start = self._curr
                 while self._curr and not self._match_pair(TokenType.STAR, TokenType.SLASH):
@@ -271,6 +323,7 @@ class Generator(generator.Generator):
         LAST_DAY_SUPPORTS_DATE_PART = False
         SUPPORTS_SELECT_INTO = True
         TZ_TO_WITH_TIME_ZONE = True
+        QUERY_HINT_SEP = " "
 
         TYPE_MAPPING = {
             **generator.Generator.TYPE_MAPPING,
@@ -370,3 +423,23 @@ def into_sql(self, expression: exp.Into) -> str:
                 return f"{self.seg(into)} {self.sql(expression, 'this')}"
 
             return f"{self.seg(into)} {self.expressions(expression)}"
+
+        def hint_sql(self, expression: exp.Hint) -> str:
+            expressions = []
+
+            for expression in expression.expressions:
+                if isinstance(expression, exp.Anonymous):
+                    formatted_args = self._format_hint_function_args(*expression.expressions)
+                    expressions.append(f"{self.sql(expression, 'this')}({formatted_args})")
+                else:
+                    expressions.append(self.sql(expression))
+
+            return f" /*+ {self.expressions(sqls=expressions, sep=self.QUERY_HINT_SEP).strip()} */"
+
+        def _format_hint_function_args(self, *args: t.Optional[str | exp.Expression]) -> str:
+            arg_sqls = tuple(self.sql(arg) for arg in args)
+            if self.pretty and self.too_wide(arg_sqls):
+                return self.indent(
+                    "\n" + "\n".join(arg_sqls) + "\n", skip_first=True, skip_last=True
+                )
+            return " ".join(arg_sqls)
diff --git a/tests/dialects/test_oracle.py b/tests/dialects/test_oracle.py
@@ -329,6 +329,57 @@ def test_hints(self):
         )
         self.validate_identity("INSERT /*+ APPEND */ INTO IAP_TBL (id, col1) VALUES (2, 'test2')")
         self.validate_identity("INSERT /*+ APPEND_VALUES */ INTO dest_table VALUES (i, 'Value')")
+        self.validate_identity(
+            "SELECT /*+ LEADING(departments employees) USE_NL(employees) */ * FROM employees JOIN departments ON employees.department_id = departments.department_id",
+            """SELECT /*+ LEADING(departments employees)
+  USE_NL(employees) */
+  *
+FROM employees
+JOIN departments
+  ON employees.department_id = departments.department_id""",
+            pretty=True,
+        )
+        self.validate_identity(
+            "SELECT /*+ USE_NL(bbbbbbbbbbbbbbbbbbbbbbbb) LEADING(aaaaaaaaaaaaaaaaaaaaaaaa bbbbbbbbbbbbbbbbbbbbbbbb cccccccccccccccccccccccc dddddddddddddddddddddddd) INDEX(cccccccccccccccccccccccc) */ * FROM aaaaaaaaaaaaaaaaaaaaaaaa JOIN bbbbbbbbbbbbbbbbbbbbbbbb ON aaaaaaaaaaaaaaaaaaaaaaaa.id = bbbbbbbbbbbbbbbbbbbbbbbb.a_id JOIN cccccccccccccccccccccccc ON bbbbbbbbbbbbbbbbbbbbbbbb.id = cccccccccccccccccccccccc.b_id JOIN dddddddddddddddddddddddd ON cccccccccccccccccccccccc.id = dddddddddddddddddddddddd.c_id",
+        )
+        self.validate_identity(
+            "SELECT /*+ USE_NL(bbbbbbbbbbbbbbbbbbbbbbbb) LEADING(aaaaaaaaaaaaaaaaaaaaaaaa bbbbbbbbbbbbbbbbbbbbbbbb cccccccccccccccccccccccc dddddddddddddddddddddddd) INDEX(cccccccccccccccccccccccc) */ * FROM aaaaaaaaaaaaaaaaaaaaaaaa JOIN bbbbbbbbbbbbbbbbbbbbbbbb ON aaaaaaaaaaaaaaaaaaaaaaaa.id = bbbbbbbbbbbbbbbbbbbbbbbb.a_id JOIN cccccccccccccccccccccccc ON bbbbbbbbbbbbbbbbbbbbbbbb.id = cccccccccccccccccccccccc.b_id JOIN dddddddddddddddddddddddd ON cccccccccccccccccccccccc.id = dddddddddddddddddddddddd.c_id",
+            """SELECT /*+ USE_NL(bbbbbbbbbbbbbbbbbbbbbbbb)
+  LEADING(
+    aaaaaaaaaaaaaaaaaaaaaaaa
+    bbbbbbbbbbbbbbbbbbbbbbbb
+    cccccccccccccccccccccccc
+    dddddddddddddddddddddddd
+  )
+  INDEX(cccccccccccccccccccccccc) */
+  *
+FROM aaaaaaaaaaaaaaaaaaaaaaaa
+JOIN bbbbbbbbbbbbbbbbbbbbbbbb
+  ON aaaaaaaaaaaaaaaaaaaaaaaa.id = bbbbbbbbbbbbbbbbbbbbbbbb.a_id
+JOIN cccccccccccccccccccccccc
+  ON bbbbbbbbbbbbbbbbbbbbbbbb.id = cccccccccccccccccccccccc.b_id
+JOIN dddddddddddddddddddddddd
+  ON cccccccccccccccccccccccc.id = dddddddddddddddddddddddd.c_id""",
+            pretty=True,
+        )
+        # Test that parsing error with keywords like select where etc falls back
+        self.validate_identity(
+            "SELECT /*+ LEADING(departments employees) USE_NL(employees) select where group by is order by */ * FROM employees JOIN departments ON employees.department_id = departments.department_id",
+            """SELECT /*+ LEADING(departments employees) USE_NL(employees) select where group by is order by */
+  *
+FROM employees
+JOIN departments
+  ON employees.department_id = departments.department_id""",
+            pretty=True,
+        )
+        # Test that parsing error with , inside hint function falls back
+        self.validate_identity(
+            "SELECT /*+ LEADING(departments, employees) */ * FROM employees JOIN departments ON employees.department_id = departments.department_id"
+        )
+        # Test that parsing error with keyword inside hint function falls back
+        self.validate_identity(
+            "SELECT /*+ LEADING(departments select) */ * FROM employees JOIN departments ON employees.department_id = departments.department_id"
+        )
 
     def test_xml_table(self):
         self.validate_identity("XMLTABLE('x')")
