feat(oracle): Support for WITH READ ONLY / CHECK OPTION (#3639)

VaggelisD · web-flow · commit 6abd2c943896 · 2024-06-12T20:43:37.000+03:00
* feat(oracle): Support for WITH READ ONLY / CHECK OPTION

* PR Feedback 1

* PR Feedback 2

* PR Feedback 3
diff --git a/sqlglot/dialects/oracle.py b/sqlglot/dialects/oracle.py
@@ -13,6 +13,7 @@
     trim_sql,
 )
 from sqlglot.helper import seq_get
+from sqlglot.parser import OPTIONS_TYPE
 from sqlglot.tokens import TokenType
 
 if t.TYPE_CHECKING:
@@ -36,8 +37,8 @@ def eliminate_join_marks(ast: exp.Expression) -> exp.Expression:
     from sqlglot.optimizer.scope import traverse_scope
 
     """Remove join marks from an expression
-    
-    SELECT * FROM a, b WHERE a.id = b.id(+)   
+
+    SELECT * FROM a, b WHERE a.id = b.id(+)
     becomes:
     SELECT * FROM a LEFT JOIN b ON a.id = b.id
 
@@ -299,6 +300,7 @@ class Parser(parser.Parser):
         QUERY_MODIFIER_PARSERS = {
             **parser.Parser.QUERY_MODIFIER_PARSERS,
             TokenType.ORDER_SIBLINGS_BY: lambda self: ("order", self._parse_order()),
+            TokenType.WITH: lambda self: ("options", self._parse_query_restrictions()),
         }
 
         TYPE_LITERAL_PARSERS = {
@@ -311,6 +313,13 @@ class Parser(parser.Parser):
         # Reference: https://stackoverflow.com/a/336455
         DISTINCT_TOKENS = {TokenType.DISTINCT, TokenType.UNIQUE}
 
+        QUERY_RESTRICTIONS: OPTIONS_TYPE = {
+            "WITH": (
+                ("READ", "ONLY"),
+                ("CHECK", "OPTION"),
+            ),
+        }
+
         def _parse_xml_table(self) -> exp.XMLTable:
             this = self._parse_string()
 
@@ -354,6 +363,18 @@ def _parse_hint(self) -> t.Optional[exp.Hint]:
 
             return None
 
+        def _parse_query_restrictions(self) -> t.Optional[exp.Expression]:
+            kind = self._parse_var_from_options(self.QUERY_RESTRICTIONS, raise_unmatched=False)
+
+            if not kind:
+                return None
+
+            return self.expression(
+                exp.QueryOption,
+                this=kind,
+                expression=self._match(TokenType.CONSTRAINT) and self._parse_field(),
+            )
+
     class Generator(generator.Generator):
         LOCKING_READS_SUPPORTED = True
         JOIN_HINTS = False
@@ -442,3 +463,10 @@ def add_column_sql(self, expression: exp.AlterTable) -> str:
             if len(expression.args.get("actions", [])) > 1:
                 return f"ADD ({actions})"
             return f"ADD {actions}"
+
+        def queryoption_sql(self, expression: exp.QueryOption) -> str:
+            option = self.sql(expression, "this")
+            value = self.sql(expression, "expression")
+            value = f" CONSTRAINT {value}" if value else ""
+
+            return f"{option}{value}"
diff --git a/sqlglot/dialects/tsql.py b/sqlglot/dialects/tsql.py
@@ -1118,3 +1118,7 @@ def declareitem_sql(self, expression: exp.DeclareItem) -> str:
                 kind = f"TABLE {kind}"
 
             return f"{variable} AS {kind}{default}"
+
+        def options_modifier(self, expression: exp.Expression) -> str:
+            options = self.expressions(expression, key="options")
+            return f" OPTION{self.wrap(options)}" if options else ""
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -2244,10 +2244,6 @@ def query_modifiers(self, expression: exp.Expression, *sqls: str) -> str:
         elif self.LIMIT_FETCH == "FETCH" and isinstance(limit, exp.Limit):
             limit = exp.Fetch(direction="FIRST", count=exp.maybe_copy(limit.expression))
 
-        options = self.expressions(expression, key="options")
-        if options:
-            options = f" OPTION{self.wrap(options)}"
-
         return csv(
             *sqls,
             *[self.sql(join) for join in expression.args.get("joins") or []],
@@ -2262,10 +2258,14 @@ def query_modifiers(self, expression: exp.Expression, *sqls: str) -> str:
             self.sql(expression, "order"),
             *self.offset_limit_modifiers(expression, isinstance(limit, exp.Fetch), limit),
             *self.after_limit_modifiers(expression),
-            options,
+            self.options_modifier(expression),
             sep="",
         )
 
+    def options_modifier(self, expression: exp.Expression) -> str:
+        options = self.sql(expression, "options")
+        return f" {options}" if options else ""
+
     def queryoption_sql(self, expression: exp.QueryOption) -> str:
         return ""
 
diff --git a/tests/dialects/test_oracle.py b/tests/dialects/test_oracle.py
@@ -468,3 +468,12 @@ def test_eliminate_join_marks(self):
                     eliminate_join_marks(self.parse_one(original)).sql(dialect=self.dialect),
                     expected,
                 )
+
+    def test_query_restrictions(self):
+        for restriction in ("READ ONLY", "CHECK OPTION"):
+            for constraint_name in (" CONSTRAINT name", ""):
+                with self.subTest(f"Restriction: {restriction}"):
+                    self.validate_identity(f"SELECT * FROM tbl WITH {restriction}{constraint_name}")
+                    self.validate_identity(
+                        f"CREATE VIEW view AS SELECT * FROM tbl WITH {restriction}{constraint_name}"
+                    )
