MySQL: support variable assignments by assignment operator :=

src/sqlfluff/dialects/dialect_mysql.py

@@ -125,6 +125,7 @@
         insert=[
             Ref("SessionVariableNameSegment"),
             Ref("LocalVariableNameSegment"),
+            Ref("VariableAssignmentSegment"),
         ]
     ),
     DateTimeLiteralGrammar=Sequence(
@@ -188,6 +189,14 @@
         StringParser("NOT", KeywordSegment, type="keyword"),
         StringParser("!", CodeSegment, type="not_operator"),
     ),
+    Expression_C_Grammar=Sequence(
+        Sequence(
+            Ref("SessionVariableNameSegment"),
+            Ref("WalrusOperatorSegment"),
+            optional=True,
+        ),
+        ansi_dialect.get_grammar("Expression_C_Grammar"),
+    ),
 )
 
 mysql_dialect.add(
@@ -760,6 +769,12 @@ class IntervalExpressionSegment(BaseSegment):
         CodeSegment,
         type="variable",
     ),
+    WalrusOperatorSegment=StringParser(":=", SymbolSegment, type="assignment_operator"),
+    VariableAssignmentSegment=Sequence(
+        Ref("SessionVariableNameSegment"),
+        Ref("WalrusOperatorSegment"),
+        Ref("BaseExpressionElementGrammar"),
+    ),
     BooleanDynamicSystemVariablesGrammar=OneOf(
         # Boolean dynamic system varaiables can be set to ON/OFF, TRUE/FALSE, or 0/1:
         # https://dev.mysql.com/doc/refman/8.0/en/dynamic-system-variables.html
@@ -800,6 +815,14 @@ class IntervalExpressionSegment(BaseSegment):
 )
 
 
+mysql_dialect.insert_lexer_matchers(
+    [
+        StringLexer("walrus_operator", ":=", CodeSegment),
+    ],
+    before="equals",
+)
+
+
 class RoleReferenceSegment(ansi.RoleReferenceSegment):
     """A reference to an account, role, or user.
 
@@ -1151,7 +1174,10 @@ class SetAssignmentStatementSegment(BaseSegment):
                 OneOf(
                     Ref("SessionVariableNameSegment"), Ref("LocalVariableNameSegment")
                 ),
-                Ref("EqualsSegment"),
+                OneOf(
+                    Ref("EqualsSegment"),
+                    Ref("WalrusOperatorSegment"),
+                ),
                 AnyNumberOf(
                     Ref("QuotedLiteralSegment"),
                     Ref("DoubleQuotedLiteralSegment"),

src/sqlfluff/rules/L006.py

@@ -44,7 +44,9 @@ class Rule_L006(BaseRule):
     """
 
     groups = ("all", "core")
-    crawl_behaviour = ParentOfSegmentCrawler({"binary_operator", "comparison_operator"})
+    crawl_behaviour = ParentOfSegmentCrawler(
+        {"binary_operator", "comparison_operator", "assignment_operator"}
+    )
     # L006 works on operators so requires three operators.
     # However some rules that inherit from here (e.g. L048) do not.
     # So allow this to be configurable.
@@ -53,6 +55,7 @@ class Rule_L006(BaseRule):
     _target_elems: List[Tuple[str, str]] = [
         ("type", "binary_operator"),
         ("type", "comparison_operator"),
+        ("type", "assignment_operator"),
     ]
 
     @staticmethod

test/fixtures/dialects/mysql/variable_assignment.sql

@@ -0,0 +1,7 @@
+SELECT @var1:=COUNT(*) FROM t1;
+
+SET @var1:=0;
+
+SET @var1:=@var2:=0;
+
+UPDATE t1 SET c1 = 2 WHERE c1 = @var1:= 1;

test/fixtures/dialects/mysql/variable_assignment.yml

@@ -0,0 +1,74 @@
+# 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: afaf7d13bb288b44720e378bce9b97f290a3d003ca4bc17e02440b6b1438ba95
+file:
+- statement:
+    select_statement:
+      select_clause:
+        keyword: SELECT
+        select_clause_element:
+          expression:
+            variable: '@var1'
+            assignment_operator: :=
+            function:
+              function_name:
+                function_name_identifier: COUNT
+              bracketed:
+                start_bracket: (
+                star: '*'
+                end_bracket: )
+      from_clause:
+        keyword: FROM
+        from_expression:
+          from_expression_element:
+            table_expression:
+              table_reference:
+                naked_identifier: t1
+- statement_terminator: ;
+- statement:
+    set_statement:
+    - keyword: SET
+    - variable: '@var1'
+    - assignment_operator: :=
+    - variable: '0'
+- statement_terminator: ;
+- statement:
+    set_statement:
+      keyword: SET
+      variable: '@var1'
+      assignment_operator: :=
+      expression:
+        variable: '@var2'
+        assignment_operator: :=
+        numeric_literal: '0'
+- statement_terminator: ;
+- statement:
+    update_statement:
+      keyword: UPDATE
+      from_expression:
+        from_expression_element:
+          table_expression:
+            table_reference:
+              naked_identifier: t1
+      set_clause_list:
+        keyword: SET
+        set_clause:
+          column_reference:
+            naked_identifier: c1
+          comparison_operator:
+            raw_comparison_operator: '='
+          numeric_literal: '2'
+      where_clause:
+        keyword: WHERE
+        expression:
+          column_reference:
+            naked_identifier: c1
+          comparison_operator:
+            raw_comparison_operator: '='
+          variable: '@var1'
+          assignment_operator: :=
+          numeric_literal: '1'
+- statement_terminator: ;