Feat(oracle): support unicode strings u'...' (#3641)

georgesittas · web-flow · commit 227e0544ede5 · 2024-06-12T18:39:09.000+03:00
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
@@ -169,6 +169,7 @@ def get_start_end(token_type: TokenType) -> t.Tuple[t.Optional[str], t.Optional[
 
         if enum not in ("", "athena", "presto", "trino"):
             klass.generator_class.TRY_SUPPORTED = False
+            klass.generator_class.SUPPORTS_UESCAPE = False
 
         if enum not in ("", "databricks", "hive", "spark", "spark2"):
             modifier_transforms = klass.generator_class.AFTER_HAVING_MODIFIER_TRANSFORMS.copy()
diff --git a/sqlglot/dialects/oracle.py b/sqlglot/dialects/oracle.py
@@ -231,6 +231,12 @@ class Oracle(Dialect):
     class Tokenizer(tokens.Tokenizer):
         VAR_SINGLE_TOKENS = {"@", "$", "#"}
 
+        UNICODE_STRINGS = [
+            (prefix + q, q)
+            for q in t.cast(t.List[str], tokens.Tokenizer.QUOTES)
+            for prefix in ("U", "u")
+        ]
+
         KEYWORDS = {
             **tokens.Tokenizer.KEYWORDS,
             "(+)": TokenType.JOIN_MARKER,
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -356,6 +356,9 @@ class Generator(metaclass=_Generator):
     # Whether the conditional TRY(expression) function is supported
     TRY_SUPPORTED = True
 
+    # Whether the UESCAPE syntax in unicode strings is supported
+    SUPPORTS_UESCAPE = True
+
     # The keyword to use when generating a star projection with excluded columns
     STAR_EXCEPT = "EXCEPT"
 
@@ -1143,16 +1146,23 @@ def unicodestring_sql(self, expression: exp.UnicodeString) -> str:
         escape = expression.args.get("escape")
 
         if self.dialect.UNICODE_START:
-            escape = f" UESCAPE {self.sql(escape)}" if escape else ""
-            return f"{self.dialect.UNICODE_START}{this}{self.dialect.UNICODE_END}{escape}"
+            escape_substitute = r"\\\1"
+            left_quote, right_quote = self.dialect.UNICODE_START, self.dialect.UNICODE_END
+        else:
+            escape_substitute = r"\\u\1"
+            left_quote, right_quote = self.dialect.QUOTE_START, self.dialect.QUOTE_END
 
         if escape:
-            pattern = re.compile(rf"{escape.name}(\d+)")
+            escape_pattern = re.compile(rf"{escape.name}(\d+)")
+            escape_sql = f" UESCAPE {self.sql(escape)}" if self.SUPPORTS_UESCAPE else ""
         else:
-            pattern = ESCAPED_UNICODE_RE
+            escape_pattern = ESCAPED_UNICODE_RE
+            escape_sql = ""
+
+        if not self.dialect.UNICODE_START or (escape and not self.SUPPORTS_UESCAPE):
+            this = escape_pattern.sub(escape_substitute, this)
 
-        this = pattern.sub(r"\\u\1", this)
-        return f"{self.dialect.QUOTE_START}{this}{self.dialect.QUOTE_END}"
+        return f"{left_quote}{this}{right_quote}{escape_sql}"
 
     def rawstring_sql(self, expression: exp.RawString) -> str:
         string = self.escape_str(expression.this.replace("\\", "\\\\"), escape_backslash=False)
diff --git a/tests/dialects/test_presto.py b/tests/dialects/test_presto.py
@@ -564,6 +564,7 @@ def test_unicode_string(self):
             self.validate_all(
                 f"{prefix}'Hello winter \\2603 !'",
                 write={
+                    "oracle": "U'Hello winter \\2603 !'",
                     "presto": "U&'Hello winter \\2603 !'",
                     "snowflake": "'Hello winter \\u2603 !'",
                     "spark": "'Hello winter \\u2603 !'",
@@ -572,6 +573,7 @@ def test_unicode_string(self):
             self.validate_all(
                 f"{prefix}'Hello winter #2603 !' UESCAPE '#'",
                 write={
+                    "oracle": "U'Hello winter \\2603 !'",
                     "presto": "U&'Hello winter #2603 !' UESCAPE '#'",
                     "snowflake": "'Hello winter \\u2603 !'",
                     "spark": "'Hello winter \\u2603 !'",
