feat(optimizer)!: annotate type for Snowflake REGEXP_SUBSTR_ALL function

Author: fivetran-BradfordPaskewitz

sqlglot/dialects/snowflake.py

@@ -333,7 +333,9 @@ def _builder(args: t.List) -> E | exp.Escape:
     return _builder
 
 
-def _regexpextract_sql(self, expression: exp.RegexpExtract | exp.RegexpExtractAll) -> str:
+def _regexpextract_sql(
+    self, expression: exp.RegexpExtract | exp.RegexpExtractAll | exp.RegexpSubstrAll
+) -> str:
     # Other dialects don't support all of the following parameters, so we need to
     # generate default values as necessary to ensure the transpilation is correct
     group = expression.args.get("group")
@@ -347,8 +349,15 @@ def _regexpextract_sql(self, expression: exp.RegexpExtract | exp.RegexpExtractAl
     occurrence = expression.args.get("occurrence") or (parameters and exp.Literal.number(1))
     position = expression.args.get("position") or (occurrence and exp.Literal.number(1))
 
+    if isinstance(expression, exp.RegexpExtract):
+        func_name = "REGEXP_SUBSTR"
+    elif isinstance(expression, exp.RegexpSubstrAll):
+        func_name = "REGEXP_SUBSTR_ALL"
+    else:  # exp.RegexpExtractAll
+        func_name = "REGEXP_EXTRACT_ALL"
+
     return self.func(
-        "REGEXP_SUBSTR" if isinstance(expression, exp.RegexpExtract) else "REGEXP_EXTRACT_ALL",
+        func_name,
         expression.this,
         expression.expression,
         position,
@@ -585,6 +594,7 @@ class Snowflake(Dialect):
         },
         exp.DataType.Type.ARRAY: {
             exp.Split,
+            exp.RegexpSubstrAll,
         },
         exp.DataType.Type.OBJECT: {
             exp.ParseUrl,
@@ -755,7 +765,7 @@ class Parser(parser.Parser):
             "REGEXP_EXTRACT_ALL": _build_regexp_extract(exp.RegexpExtractAll),
             "REGEXP_REPLACE": _build_regexp_replace,
             "REGEXP_SUBSTR": _build_regexp_extract(exp.RegexpExtract),
-            "REGEXP_SUBSTR_ALL": _build_regexp_extract(exp.RegexpExtractAll),
+            "REGEXP_SUBSTR_ALL": _build_regexp_extract(exp.RegexpSubstrAll),
             "REPLACE": build_replace_with_optional_replacement,
             "RLIKE": exp.RegexpLike.from_arg_list,
             "SHA1_BINARY": exp.SHA1Digest.from_arg_list,
@@ -1424,6 +1434,7 @@ class Generator(generator.Generator):
             exp.Pivot: transforms.preprocess([_unqualify_pivot_columns]),
             exp.RegexpExtract: _regexpextract_sql,
             exp.RegexpExtractAll: _regexpextract_sql,
+            exp.RegexpSubstrAll: _regexpextract_sql,
             exp.RegexpILike: _regexpilike_sql,
             exp.Rand: rename_func("RANDOM"),
             exp.Select: transforms.preprocess(

sqlglot/expressions.py

@@ -7256,6 +7256,17 @@ class RegexpExtractAll(Func):
     }
 
 
+class RegexpSubstrAll(Func):
+    arg_types = {
+        "this": True,
+        "expression": True,
+        "position": False,
+        "occurrence": False,
+        "parameters": False,
+        "group": False,
+    }
+
+
 class RegexpReplace(Func):
     arg_types = {
         "this": True,

tests/dialects/test_snowflake.py

@@ -2416,10 +2416,15 @@ def test_regexp_substr(self, logger):
             },
         )
 
+        self.validate_identity("REGEXP_SUBSTR_ALL(subject, pattern)")
+        self.validate_identity("REGEXP_SUBSTR_ALL(subject, pattern, 1)")
+        self.validate_identity("REGEXP_SUBSTR_ALL(subject, pattern, 1, 1)")
+        self.validate_identity("REGEXP_SUBSTR_ALL(subject, pattern, 1, 1, 'i')")
         self.validate_identity(
-            "REGEXP_SUBSTR_ALL(subject, pattern)",
-            "REGEXP_EXTRACT_ALL(subject, pattern)",
+            "REGEXP_SUBSTR_ALL(subject, pattern, 1, 1, 'i', 0)",
+            "REGEXP_SUBSTR_ALL(subject, pattern, 1, 1, 'i')",
         )
+        self.validate_identity("REGEXP_EXTRACT_ALL(subject, pattern)")
 
         self.validate_identity("SELECT REGEXP_COUNT('hello world', 'l')")
         self.validate_identity("SELECT REGEXP_COUNT('hello world', 'l', 1)")

tests/fixtures/optimizer/annotate_functions.sql

@@ -1927,6 +1927,26 @@ VARCHAR;
 REGEXP_SUBSTR('hello world', 'world', 1, 1, 'e', NULL);
 VARCHAR;
 
+# dialect: snowflake
+REGEXP_SUBSTR_ALL('hello world', 'world');
+ARRAY;
+
+# dialect: snowflake
+REGEXP_SUBSTR_ALL('hello world', 'world', 1);
+ARRAY;
+
+# dialect: snowflake
+REGEXP_SUBSTR_ALL('hello world', 'world', 1, 1);
+ARRAY;
+
+# dialect: snowflake
+REGEXP_SUBSTR_ALL('hello world', 'world', 1, 1, 'i');
+ARRAY;
+
+# dialect: snowflake
+REGEXP_SUBSTR_ALL('hello world', 'world', 1, 1, 'i', 0);
+ARRAY;
+
 # dialect: snowflake
 REPEAT('hello', 3);
 VARCHAR;