Handle escaping when converting delimiters to regex expr

Author: treysp

sqlglot/dialects/bigquery.py

@@ -360,7 +360,7 @@ class BigQuery(Dialect):
 
     # https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#initcap
     INITCAP_SUPPORTS_CUSTOM_DELIMITERS = True
-    INITCAP_DEFAULT_DELIMITER_CHARS = r' \t\n\r\f\v\[\](){}/|\<>!?@"^#$&~_,.:;*%+\-'
+    INITCAP_DEFAULT_DELIMITER_CHARS = ' \t\n\r\f\v\\[\\](){}/|\<>!?@"^#$&~_,.:;*%+\\-'
 
     # https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#case_sensitivity
     NORMALIZATION_STRATEGY = NormalizationStrategy.CASE_INSENSITIVE

sqlglot/dialects/dialect.py

@@ -540,7 +540,7 @@ class Dialect(metaclass=_Dialect):
     # Whether the INITCAP function supports custom delimiter characters as the second argument
     # Default delimiter characters for INITCAP function: whitespace and non-alphanumeric characters
     INITCAP_SUPPORTS_CUSTOM_DELIMITERS = False
-    INITCAP_DEFAULT_DELIMITER_CHARS = r" \t\n\r\f\v!\"#$%&'()*+,\-./:;<=>?@\[\\]^_`{|}~"
+    INITCAP_DEFAULT_DELIMITER_CHARS = " \t\n\r\f\v!\"#$%&'()*+,\\-./:;<=>?@\\[\\]^_`{|}~"
 
     BYTE_STRING_IS_BYTES_TYPE: bool = False
     """

sqlglot/dialects/duckdb.py

@@ -268,54 +268,110 @@ def _json_extract_value_array_sql(
     return self.sql(exp.cast(json_extract, to=exp.DataType.build(data_type)))
 
 
-def _initcap_sql(self: DuckDB.Generator, expression: exp.Initcap) -> str:
-    def build_capitalize_sql(
-        value_to_split: str, delimiters_sql: str, convert_delim_to_regex: bool = True
-    ) -> str:
-        # empty string delimiter --> treat value as one word, no need to split
-        if delimiters_sql == "''":
-            return f"UPPER(LEFT({value_to_split}, 1)) || LOWER(SUBSTR({value_to_split}, 2))"
-
-        delim_regex_sql = delimiters_sql
-        split_regex_sql = delimiters_sql
-        if convert_delim_to_regex:
-            delim_regex_sql = f"CONCAT('[', {delimiters_sql}, ']')"
-            split_regex_sql = f"CONCAT('([', {delimiters_sql}, ']+|[^', {delimiters_sql}, ']+)')"
-
-        # REGEXP_EXTRACT_ALL produces a list of string segments, alternating between delimiter and non-delimiter segments.
-        # We do not know whether the first segment is a delimiter or not, so we check the first character of the string
-        # with REGEXP_MATCHES. If the first char is a delimiter, we capitalize even list indexes, otherwise capitalize odd.
-        return self.func(
-            "ARRAY_TO_STRING",
-            exp.case()
-            .when(
-                f"REGEXP_MATCHES(LEFT({value_to_split}, 1), {delim_regex_sql})",
-                self.func(
-                    "LIST_TRANSFORM",
-                    self.func("REGEXP_EXTRACT_ALL", value_to_split, split_regex_sql),
-                    "(seg, idx) -> CASE WHEN idx % 2 = 0 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTR(seg, 2)) ELSE seg END",
-                ),
-            )
-            .else_(
-                self.func(
-                    "LIST_TRANSFORM",
-                    self.func("REGEXP_EXTRACT_ALL", value_to_split, split_regex_sql),
-                    "(seg, idx) -> CASE WHEN idx % 2 = 1 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTR(seg, 2)) ELSE seg END",
-                ),
+def _escape_regex_metachars(
+    self: DuckDB.Generator, delimiters: t.Optional[exp.Expression], delimiters_sql: str
+) -> str:
+    if not delimiters:
+        return delimiters_sql
+
+    REGEX_LITERAL_ESCAPES = {
+        "\\": "\\\\",  # literals need two slashes inside []
+        "-": "\\-",
+        "^": "\\^",
+        "[": "\\[",
+        "]": "\\]",
+    }
+
+    if isinstance(delimiters, exp.Literal) and delimiters.is_string:
+        literal_value = delimiters.this
+        escaped_literal = "".join(REGEX_LITERAL_ESCAPES.get(ch, ch) for ch in literal_value)
+        return self.sql(exp.Literal.string(escaped_literal))
+
+    REGEX_ESCAPE_REPLACEMENTS = (
+        ("\\", "\\\\"),
+        ("-", r"\-"),
+        ("^", r"\^"),
+        ("[", r"\["),
+        ("]", r"\]"),
+    )
+
+    escaped_sql = delimiters_sql
+    for raw, escaped in REGEX_ESCAPE_REPLACEMENTS:
+        escaped_sql = self.func(
+            "REPLACE",
+            escaped_sql,
+            self.sql(exp.Literal.string(raw)),
+            self.sql(exp.Literal.string(escaped)),
+        )
+
+    return escaped_sql
+
+
+def _build_capitalization_sql(
+    self: DuckDB.Generator,
+    value_to_split: str,
+    raw_delimiters_sql: str,
+    escaped_delimiters_sql: t.Optional[str] = None,
+    convert_delim_to_regex: bool = True,
+) -> str:
+    # empty string delimiter --> treat value as one word, no need to split
+    if raw_delimiters_sql == "''":
+        return f"UPPER(LEFT({value_to_split}, 1)) || LOWER(SUBSTR({value_to_split}, 2))"
+
+    regex_ready_sql = escaped_delimiters_sql or raw_delimiters_sql
+    delim_regex_sql = regex_ready_sql
+    split_regex_sql = regex_ready_sql
+    if convert_delim_to_regex:
+        delim_regex_sql = f"CONCAT('[', {regex_ready_sql}, ']')"
+        split_regex_sql = f"CONCAT('([', {regex_ready_sql}, ']+|[^', {regex_ready_sql}, ']+)')"
+
+    # REGEXP_EXTRACT_ALL produces a list of string segments, alternating between delimiter and non-delimiter segments.
+    # We do not know whether the first segment is a delimiter or not, so we check the first character of the string
+    # with REGEXP_MATCHES. If the first char is a delimiter, we capitalize even list indexes, otherwise capitalize odd.
+    return self.func(
+        "ARRAY_TO_STRING",
+        exp.case()
+        .when(
+            f"REGEXP_MATCHES(LEFT({value_to_split}, 1), {delim_regex_sql})",
+            self.func(
+                "LIST_TRANSFORM",
+                self.func("REGEXP_EXTRACT_ALL", value_to_split, split_regex_sql),
+                "(seg, idx) -> CASE WHEN idx % 2 = 0 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTR(seg, 2)) ELSE seg END",
             ),
-            "''",
         )
+        .else_(
+            self.func(
+                "LIST_TRANSFORM",
+                self.func("REGEXP_EXTRACT_ALL", value_to_split, split_regex_sql),
+                "(seg, idx) -> CASE WHEN idx % 2 = 1 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTR(seg, 2)) ELSE seg END",
+            ),
+        ),
+        "''",
+    )
 
+
+def _initcap_sql(self: DuckDB.Generator, expression: exp.Initcap) -> str:
     this_sql = self.sql(expression, "this")
     delimiters = expression.args.get("expression")
     delimiters_sql = self.sql(delimiters)
+    escaped_delimiters_sql = (
+        _escape_regex_metachars(self, delimiters, delimiters_sql)
+        if not isinstance(delimiters, exp.Null)
+        else delimiters_sql
+    )
 
     if delimiters and (isinstance(delimiters, exp.Literal) and delimiters.is_string):
-        return f"CASE WHEN {this_sql} IS NULL THEN NULL ELSE {build_capitalize_sql(this_sql, delimiters_sql)} END"
+        return (
+            f"CASE WHEN {this_sql} IS NULL THEN NULL ELSE "
+            f"{_build_capitalization_sql(self, this_sql, delimiters_sql, escaped_delimiters_sql)} END"
+        )
 
-    # delimiters arg is SQL expression or NULL
-    capitalize_sql = build_capitalize_sql(
-        this_sql, delimiters_sql, convert_delim_to_regex=not isinstance(delimiters, exp.Null)
+    capitalize_sql = _build_capitalization_sql(
+        self,
+        this_sql,
+        delimiters_sql,
+        escaped_delimiters_sql,
+        convert_delim_to_regex=not isinstance(delimiters, exp.Null),
     )
     return f"CASE WHEN {this_sql} IS NULL OR {delimiters_sql} IS NULL THEN NULL ELSE {capitalize_sql} END"
 

sqlglot/dialects/snowflake.py

@@ -558,7 +558,7 @@ class Snowflake(Dialect):
 
     # https://docs.snowflake.com/en/en/sql-reference/functions/initcap
     INITCAP_SUPPORTS_CUSTOM_DELIMITERS = True
-    INITCAP_DEFAULT_DELIMITER_CHARS = r' \t\n\r\f\v!?@"^#$&~_,.:;+\-*%/|\[\](){}<>'
+    INITCAP_DEFAULT_DELIMITER_CHARS = ' \t\n\r\f\v!?@"^#$&~_,.:;+\\-*%/|\\[\\](){}<>'
 
     TIME_MAPPING = {
         "YYYY": "%Y",

sqlglot/dialects/spark2.py

@@ -121,7 +121,7 @@ class Spark2(Hive):
 
     # https://spark.apache.org/docs/latest/api/sql/index.html#initcap
     # https://docs.databricks.com/aws/en/sql/language-manual/functions/initcap
-    INITCAP_DEFAULT_DELIMITER_CHARS = r" \t\n\r\f\v"
+    INITCAP_DEFAULT_DELIMITER_CHARS = " \t\n\r\f\v"
 
     class Tokenizer(Hive.Tokenizer):
         HEX_STRINGS = [("X'", "'"), ("x'", "'")]

tests/dialects/test_dialect.py

@@ -4313,6 +4313,18 @@ def test_initcap(self):
             "spark": Spark2.INITCAP_DEFAULT_DELIMITER_CHARS,
         }
 
+        REGEX_LITERAL_ESCAPES = {
+            "\\": "\\\\",
+            "-": "\\-",
+            "^": "\\^",
+            "[": "\\[",
+            "]": "\\]",
+        }
+
+        def duckdb_regex_literal_sql(delimiters: str) -> str:
+            escaped_literal = "".join(REGEX_LITERAL_ESCAPES.get(ch, ch) for ch in delimiters)
+            return exp.Literal.string(escaped_literal).sql("duckdb")
+
         # default delimiters not present in roundtrip
         for dialect in delimiter_chars.keys():
             with self.subTest(
@@ -4338,13 +4350,13 @@ def test_initcap(self):
 
         for dialect, default_delimiters in delimiter_chars.items():
             with self.subTest(f"DuckDB rewrite for {dialect or 'default'} default delimiters"):
-                literal = exp.Literal.string(default_delimiters).sql()
+                escaped_literal = duckdb_regex_literal_sql(default_delimiters)
                 expected = (
                     "CASE WHEN col IS NULL THEN NULL ELSE ARRAY_TO_STRING("
-                    f"CASE WHEN REGEXP_MATCHES(LEFT(col, 1), '[' || {literal} || ']') "
-                    f"THEN LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || {literal} || ']+|[^' || {literal} || ']+)'), "
+                    f"CASE WHEN REGEXP_MATCHES(LEFT(col, 1), '[' || {escaped_literal} || ']') "
+                    f"THEN LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || {escaped_literal} || ']+|[^' || {escaped_literal} || ']+)'), "
                     f"(seg, idx) -> CASE WHEN idx % 2 = 0 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) "
-                    f"ELSE LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || {literal} || ']+|[^' || {literal} || ']+)'), "
+                    f"ELSE LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || {escaped_literal} || ']+|[^' || {escaped_literal} || ']+)'), "
                     f"(seg, idx) -> CASE WHEN idx % 2 = 1 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) "
                     "END, '') END"
                 )
@@ -4372,35 +4384,61 @@ def test_initcap(self):
                     "END, '') END",
                 )
 
-            for custom_delimiter in (" ", "@", " _@"):
+            for custom_delimiter in (" ", "@", " _@", r"\\"):
                 with self.subTest(
                     f"DuckDB generation for INITCAP(col, {custom_delimiter}) from {dialect}"
                 ):
+                    literal_sql = exp.Literal.string(custom_delimiter).sql(dialect)
+                    expression = parse_one(f"INITCAP(col, {literal_sql})", read=dialect)
+                    duckdb_sql = expression.sql("duckdb")
+                    escaped_custom_delimiter = duckdb_regex_literal_sql(custom_delimiter)
                     self.assertEqual(
-                        parse_one(f"INITCAP(col, '{custom_delimiter}')", read=dialect).sql(
-                            "duckdb"
-                        ),
+                        duckdb_sql,
                         "CASE WHEN col IS NULL THEN NULL ELSE ARRAY_TO_STRING("
-                        f"CASE WHEN REGEXP_MATCHES(LEFT(col, 1), '[' || '{custom_delimiter}' || ']') "
-                        f"THEN LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || '{custom_delimiter}' || ']+|[^' || '{custom_delimiter}' || ']+)'), "
+                        f"CASE WHEN REGEXP_MATCHES(LEFT(col, 1), '[' || {escaped_custom_delimiter} || ']') "
+                        f"THEN LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || {escaped_custom_delimiter} || ']+|[^' || {escaped_custom_delimiter} || ']+)'), "
                         f"(seg, idx) -> CASE WHEN idx % 2 = 0 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) "
-                        f"ELSE LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || '{custom_delimiter}' || ']+|[^' || '{custom_delimiter}' || ']+)'), "
+                        f"ELSE LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || {escaped_custom_delimiter} || ']+|[^' || {escaped_custom_delimiter} || ']+)'), "
                         f"(seg, idx) -> CASE WHEN idx % 2 = 1 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) "
                         "END, '') END",
                     )
 
+            def escape_expression_sql(sql: str) -> str:
+                escaped_sql = sql
+                for raw, escaped in (
+                    ("\\", "\\\\"),
+                    ("-", r"\-"),
+                    ("^", r"\^"),
+                    ("[", r"\["),
+                    ("]", r"\]"),
+                ):
+                    raw_sql = exp.Literal.string(raw).sql()
+                    escaped_literal_sql = exp.Literal.string(escaped).sql()
+                    escaped_sql = f"REPLACE({escaped_sql}, {raw_sql}, {escaped_literal_sql})"
+
+                return escaped_sql
+
             with self.subTest(
                 f"DuckDB generation for INITCAP subquery as custom delimiter arg from {dialect}"
             ):
+                escaped_subquery = escape_expression_sql("(SELECT delimiter FROM settings LIMIT 1)")
                 self.assertEqual(
                     parse_one(
                         "INITCAP(col, (SELECT delimiter FROM settings LIMIT 1))", read=dialect
                     ).sql("duckdb"),
                     "CASE WHEN col IS NULL OR (SELECT delimiter FROM settings LIMIT 1) IS NULL THEN NULL ELSE ARRAY_TO_STRING("
-                    + "CASE WHEN REGEXP_MATCHES(LEFT(col, 1), '[' || (SELECT delimiter FROM settings LIMIT 1) || ']') "
-                    "THEN LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || (SELECT delimiter FROM settings LIMIT 1) || ']+|[^' || (SELECT delimiter FROM settings LIMIT 1) || ']+)'), "
+                    + f"CASE WHEN REGEXP_MATCHES(LEFT(col, 1), '[' || {escaped_subquery} || ']') "
+                    "THEN LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || "
+                    + escaped_subquery
+                    + " || ']+|[^' || "
+                    + escaped_subquery
+                    + " || ']+)'), "
                     "(seg, idx) -> CASE WHEN idx % 2 = 0 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) "
-                    "ELSE LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || (SELECT delimiter FROM settings LIMIT 1) || ']+|[^' || (SELECT delimiter FROM settings LIMIT 1) || ']+)'), "
+                    "ELSE LIST_TRANSFORM(REGEXP_EXTRACT_ALL(col, '([' || "
+                    + escaped_subquery
+                    + " || ']+|[^' || "
+                    + escaped_subquery
+                    + " || ']+)'), "
                     "(seg, idx) -> CASE WHEN idx % 2 = 1 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) "
                     "END, '') END",
                 )

tests/dialects/test_hive.py

@@ -1,6 +1,6 @@
 from tests.dialects.test_dialect import Validator
-
 from sqlglot import exp
+from sqlglot.dialects import Hive
 
 
 class TestHive(Validator):
@@ -685,11 +685,32 @@ def test_hive(self):
                 "spark": "LOCATE('a', x, 3)",
             },
         )
+
+        REGEX_LITERAL_ESCAPES = {
+            "\\": "\\\\",
+            "-": "\\-",
+            "^": "\\^",
+            "[": "\\[",
+            "]": "\\]",
+        }
+
+        def duckdb_regex_literal_sql(delimiters: str) -> str:
+            escaped_literal = "".join(REGEX_LITERAL_ESCAPES.get(ch, ch) for ch in delimiters)
+            return exp.Literal.string(escaped_literal).sql("duckdb")
+
+        hive_escaped_delimiters = duckdb_regex_literal_sql(Hive.INITCAP_DEFAULT_DELIMITER_CHARS)
         self.validate_all(
             "INITCAP('new york')",
             write={
-                "duckdb": r"CASE WHEN 'new york' IS NULL THEN NULL ELSE ARRAY_TO_STRING(CASE WHEN REGEXP_MATCHES(LEFT('new york', 1), '[' || ' \t\n\r\f\v!\"#$%&''()*+,\-./:;<=>?@\[\\]^_`{|}~' || ']') THEN LIST_TRANSFORM(REGEXP_EXTRACT_ALL('new york', '([' || ' \t\n\r\f\v!\"#$%&''()*+,\-./:;<=>?@\[\\]^_`{|}~' || ']+|[^' || ' \t\n\r\f\v!\"#$%&''()*+,\-./:;<=>?@\[\\]^_`{|}~' || ']+)'), (seg, idx) -> CASE WHEN idx % 2 = 0 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) ELSE LIST_TRANSFORM(REGEXP_EXTRACT_ALL('new york', '([' || ' \t\n\r\f\v!\"#$%&''()*+,\-./:;<=>?@\[\\]^_`{|}~' || ']+|[^' || ' \t\n\r\f\v!\"#$%&''()*+,\-./:;<=>?@\[\\]^_`{|}~' || ']+)'), (seg, idx) -> CASE WHEN idx % 2 = 1 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) END, '') END",
-                "presto": r"REGEXP_REPLACE('new york', '(\w)(\w*)', x -> UPPER(x[1]) || LOWER(x[2]))",
+                "duckdb": (
+                    "CASE WHEN 'new york' IS NULL THEN NULL ELSE ARRAY_TO_STRING("
+                    f"CASE WHEN REGEXP_MATCHES(LEFT('new york', 1), '[' || {hive_escaped_delimiters} || ']') "
+                    f"THEN LIST_TRANSFORM(REGEXP_EXTRACT_ALL('new york', '([' || {hive_escaped_delimiters} || ']+|[^' || {hive_escaped_delimiters} || ']+)'), "
+                    "(seg, idx) -> CASE WHEN idx % 2 = 0 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) "
+                    f"ELSE LIST_TRANSFORM(REGEXP_EXTRACT_ALL('new york', '([' || {hive_escaped_delimiters} || ']+|[^' || {hive_escaped_delimiters} || ']+)'), "
+                    "(seg, idx) -> CASE WHEN idx % 2 = 1 THEN UPPER(LEFT(seg, 1)) || LOWER(SUBSTRING(seg, 2)) ELSE seg END) "
+                    "END, '') END"
+                ),
                 "hive": "INITCAP('new york')",
                 "spark": "INITCAP('new york')",
             },