BigQuery: Support SAFE functions (#3048)

src/sqlfluff/dialects/dialect_bigquery.py

@@ -452,75 +452,82 @@ class FunctionSegment(ansi.FunctionSegment):
     for our purposes.
     """
 
-    match_grammar = OneOf(
-        Sequence(
-            # BigQuery EXTRACT allows optional TimeZone
-            Ref("ExtractFunctionNameSegment"),
-            Bracketed(
-                Ref("DatetimeUnitSegment"),
-                "FROM",
-                Ref("ExpressionSegment"),
-                Ref("TimeZoneGrammar", optional=True),
-            ),
-        ),
-        Sequence(
-            # Treat functions which take date parts separately
-            # So those functions parse date parts as DatetimeUnitSegment
-            # rather than identifiers.
-            Ref(
-                "DatePartFunctionNameSegment",
-                exclude=Ref("ExtractFunctionNameSegment"),
-            ),
-            Bracketed(
-                Delimited(
+    match_grammar = Sequence(
+        # BigQuery Function names can be prefixed by the keyword SAFE to
+        # return NULL instead of error.
+        # https://cloud.google.com/bigquery/docs/reference/standard-sql/functions-reference#safe_prefix
+        Sequence("SAFE", Ref("DotSegment"), optional=True),
+        OneOf(
+            Sequence(
+                # BigQuery EXTRACT allows optional TimeZone
+                Ref("ExtractFunctionNameSegment"),
+                Bracketed(
                     Ref("DatetimeUnitSegment"),
-                    Ref(
-                        "FunctionContentsGrammar",
-                        ephemeral_name="FunctionContentsGrammar",
-                    ),
+                    "FROM",
+                    Ref("ExpressionSegment"),
+                    Ref("TimeZoneGrammar", optional=True),
                 ),
             ),
-        ),
-        Sequence(
             Sequence(
+                # Treat functions which take date parts separately
+                # So those functions parse date parts as DatetimeUnitSegment
+                # rather than identifiers.
                 Ref(
-                    "FunctionNameSegment",
-                    exclude=OneOf(
-                        Ref("DatePartFunctionNameSegment"),
-                        Ref("ValuesClauseSegment"),
-                    ),
+                    "DatePartFunctionNameSegment",
+                    exclude=Ref("ExtractFunctionNameSegment"),
                 ),
                 Bracketed(
-                    Ref(
-                        "FunctionContentsGrammar",
-                        # The brackets might be empty for some functions...
-                        optional=True,
-                        ephemeral_name="FunctionContentsGrammar",
-                    )
+                    Delimited(
+                        Ref("DatetimeUnitSegment"),
+                        Ref(
+                            "FunctionContentsGrammar",
+                            ephemeral_name="FunctionContentsGrammar",
+                        ),
+                    ),
                 ),
             ),
-            # Functions returning ARRAYS in BigQuery can have optional
-            # Array Accessor clauses
-            Ref("ArrayAccessorSegment", optional=True),
-            # Functions returning STRUCTs in BigQuery can have the fields
-            # elements referenced (e.g. ".a"), including wildcards (e.g. ".*")
-            # or multiple nested fields (e.g. ".a.b", or ".a.b.c")
             Sequence(
-                Ref("DotSegment"),
-                AnyNumberOf(
-                    Sequence(
-                        Ref("ParameterNameSegment"),
-                        Ref("DotSegment"),
+                Sequence(
+                    Ref(
+                        "FunctionNameSegment",
+                        exclude=OneOf(
+                            Ref("DatePartFunctionNameSegment"),
+                            Ref("ValuesClauseSegment"),
+                        ),
+                    ),
+                    Bracketed(
+                        Ref(
+                            "FunctionContentsGrammar",
+                            # The brackets might be empty for some functions...
+                            optional=True,
+                            ephemeral_name="FunctionContentsGrammar",
+                        )
                     ),
                 ),
-                OneOf(
-                    Ref("ParameterNameSegment"),
-                    Ref("StarSegment"),
+                # Functions returning ARRAYS in BigQuery can have optional
+                # Array Accessor clauses
+                Ref("ArrayAccessorSegment", optional=True),
+                # Functions returning STRUCTs in BigQuery can have the fields
+                # elements referenced (e.g. ".a"), including wildcards (e.g. ".*")
+                # or multiple nested fields (e.g. ".a.b", or ".a.b.c")
+                Sequence(
+                    Ref("DotSegment"),
+                    AnyNumberOf(
+                        Sequence(
+                            Ref("ParameterNameSegment"),
+                            Ref("DotSegment"),
+                        ),
+                    ),
+                    OneOf(
+                        Ref("ParameterNameSegment"),
+                        Ref("StarSegment"),
+                    ),
+                    optional=True,
                 ),
-                optional=True,
+                Ref("PostFunctionGrammar", optional=True),
             ),
-            Ref("PostFunctionGrammar", optional=True),
         ),
+        allow_gaps=False,
     )
 
 

src/sqlfluff/dialects/dialect_bigquery_keywords.py

@@ -214,6 +214,7 @@
 ROLLBACK
 ROW
 ROUTINE
+SAFE
 SCHEMA
 SCHEMAS
 SECOND

test/fixtures/dialects/bigquery/select_safe_function.sql

@@ -0,0 +1,7 @@
+SELECT
+    TRUE AS col1,
+    SAFE.SUBSTR('foo', 0, -2) AS col2,
+    SAFE.DATEADD(DAY, -2, CURRENT_DATE),
+    SAFE.MY_FUNCTION(column1)
+FROM
+    table1;

test/fixtures/dialects/bigquery/select_safe_function.yml

@@ -0,0 +1,79 @@
+# YML test files are auto-generated from SQL files and should not be edited by
+# hand. To help enforce this, the "hash" field in the file must match a hash
+# computed by SQLFluff when running the tests. Please run
+# `python test/generate_parse_fixture_yml.py`  to generate them after adding or
+# altering SQL files.
+_hash: f7907313f6ed80ed70a7803d4289bf65b1776f4753f359835361ed6e6353ce30
+file:
+  statement:
+    select_statement:
+      select_clause:
+      - keyword: SELECT
+      - select_clause_element:
+          literal: 'TRUE'
+          alias_expression:
+            keyword: AS
+            identifier: col1
+      - comma: ','
+      - select_clause_element:
+          function:
+            keyword: SAFE
+            dot: .
+            function_name:
+              function_name_identifier: SUBSTR
+            bracketed:
+            - start_bracket: (
+            - expression:
+                literal: "'foo'"
+            - comma: ','
+            - expression:
+                literal: '0'
+            - comma: ','
+            - expression:
+                numeric_literal:
+                  binary_operator: '-'
+                  literal: '2'
+            - end_bracket: )
+          alias_expression:
+            keyword: AS
+            identifier: col2
+      - comma: ','
+      - select_clause_element:
+          function:
+            keyword: SAFE
+            dot: .
+            function_name:
+              function_name_identifier: DATEADD
+            bracketed:
+            - start_bracket: (
+            - date_part: DAY
+            - comma: ','
+            - expression:
+                numeric_literal:
+                  binary_operator: '-'
+                  literal: '2'
+            - comma: ','
+            - expression:
+                bare_function: CURRENT_DATE
+            - end_bracket: )
+      - comma: ','
+      - select_clause_element:
+          function:
+            keyword: SAFE
+            dot: .
+            function_name:
+              function_name_identifier: MY_FUNCTION
+            bracketed:
+              start_bracket: (
+              expression:
+                column_reference:
+                  identifier: column1
+              end_bracket: )
+      from_clause:
+        keyword: FROM
+        from_expression:
+          from_expression_element:
+            table_expression:
+              table_reference:
+                identifier: table1
+  statement_terminator: ;

test/fixtures/rules/std_rule_cases/L014.yml

@@ -238,3 +238,9 @@ test_fail_inconsistent_capitalisation_properties_naked_identifier_2:
   configs:
     core:
       dialect: sparksql
+
+test_pass_bigquer_safe_does_not_trigger:
+  pass_str: SELECT SAFE.myFunction(1) AS col1
+  configs:
+    core:
+      dialect: bigquery