TSQL: Allow DECLARE/SET statements to parse using ExpressionStatement (…

…#1649) * Fix TSQL DECLARE and SET statements to fully parse expression assignments * Fix TSQL DECLARE and SET statements to fully parse expression assignments Co-authored-by: jpersons <jpersons@iuhealth.org>
sqlfluff · Oct 14, 2021 · 6e4f2fc · 6e4f2fc
1 parent 52b85c5
commit 6e4f2fc
Show file tree

Hide file tree

Showing 6 changed files with 83 additions and 73 deletions.
diff --git a/src/sqlfluff/dialects/dialect_tsql.py b/src/sqlfluff/dialects/dialect_tsql.py
@@ -17,7 +17,6 @@
     Delimited,
     Matchable,
     NamedParser,
-    StartsWith,
     OptionallyBracketed,
     Dedent,
     BaseFileSegment,
@@ -435,12 +434,7 @@ class DeclareStatementSegment(BaseSegment):
         Ref("DatatypeSegment"),
         Sequence(
             Ref("EqualsSegment"),
-            OneOf(
-                Ref("LiteralGrammar"),
-                Bracketed(Ref("SelectStatementSegment")),
-                Ref("BareFunctionSegment"),
-                Ref("FunctionSegment"),
-            ),
+            Ref("ExpressionSegment"),
             optional=True,
         ),
         Ref("DelimiterSegment", optional=True),
@@ -628,11 +622,12 @@ class SetStatementSegment(BaseSegment):
 
     Setting an already declared variable or global variable.
     https://docs.microsoft.com/en-us/sql/t-sql/statements/set-statements-transact-sql?view=sql-server-ver15
+
+    https://docs.microsoft.com/en-us/sql/t-sql/language-elements/set-local-variable-transact-sql?view=sql-server-ver15
     """
 
     type = "set_segment"
-    match_grammar = StartsWith("SET")
-    parse_grammar = Sequence(
+    match_grammar = Sequence(
         "SET",
         OneOf(
             Ref("ParameterNameSegment"),
@@ -682,25 +677,7 @@ class SetStatementSegment(BaseSegment):
             "OFF",
             Sequence(
                 Ref("EqualsSegment"),
-                OneOf(
-                    Delimited(
-                        OneOf(
-                            Ref("LiteralGrammar"),
-                            Bracketed(Ref("SelectStatementSegment")),
-                            Ref("FunctionSegment"),
-                            Bracketed(
-                                Delimited(
-                                    OneOf(
-                                        Ref("LiteralGrammar"),
-                                        Bracketed(Ref("SelectStatementSegment")),
-                                        Ref("BareFunctionSegment"),
-                                        Ref("FunctionSegment"),
-                                    )
-                                )
-                            ),
-                        )
-                    )
-                ),
+                Ref("ExpressionSegment"),
             ),
         ),
     )

diff --git a/test/fixtures/dialects/tsql/cast_variable.yml b/test/fixtures/dialects/tsql/cast_variable.yml
@@ -3,7 +3,7 @@
 # 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: 5df341c4d4c382af3d6656685aa7c53f42f4477d61d7d54eba911dd63eec7c14
+_hash: a70b2187383021670cb3d18d029294bda1d9652b38d34fc2a2fefb47072e0fe3
 file:
   batch:
   - statement:
@@ -13,36 +13,37 @@ file:
         data_type:
           identifier: date
         comparison_operator: '='
-        function:
-          function_name:
-            function_name_identifier: ISNULL
-          bracketed:
-          - start_bracket: (
-          - expression:
-              function:
-                function_name:
-                  identifier: Shared
-                  dot: .
-                  function_name_identifier: GetESTDateTime
-                bracketed:
-                  start_bracket: (
-                  expression:
-                    function:
-                      function_name:
-                        function_name_identifier: GETDATE
-                      bracketed:
-                        start_bracket: (
-                        end_bracket: )
-                  end_bracket: )
-          - comma: ','
-          - expression:
-              function:
-                function_name:
-                  function_name_identifier: GETDATE
-                bracketed:
-                  start_bracket: (
-                  end_bracket: )
-          - end_bracket: )
+        expression:
+          function:
+            function_name:
+              function_name_identifier: ISNULL
+            bracketed:
+            - start_bracket: (
+            - expression:
+                function:
+                  function_name:
+                    identifier: Shared
+                    dot: .
+                    function_name_identifier: GetESTDateTime
+                  bracketed:
+                    start_bracket: (
+                    expression:
+                      function:
+                        function_name:
+                          function_name_identifier: GETDATE
+                        bracketed:
+                          start_bracket: (
+                          end_bracket: )
+                    end_bracket: )
+            - comma: ','
+            - expression:
+                function:
+                  function_name:
+                    function_name_identifier: GETDATE
+                  bracketed:
+                    start_bracket: (
+                    end_bracket: )
+            - end_bracket: )
   - statement:
       select_statement:
         select_clause:

diff --git a/test/fixtures/dialects/tsql/declare_with_following_statements.sql b/test/fixtures/dialects/tsql/declare_with_following_statements.sql
@@ -3,6 +3,11 @@ BEGIN
 
 	DECLARE @DateNow DATE = GETDATE();
 
+	DECLARE @EOMONTH DATE = ('1900-01-01')
+
+	SET @EOMONTH = ('2000-01-01')
+
 	IF OBJECT_ID('tempdb..#UP') IS NOT NULL DROP TABLE #UP;
 
+
 END
diff --git a/test/fixtures/dialects/tsql/declare_with_following_statements.yml b/test/fixtures/dialects/tsql/declare_with_following_statements.yml
@@ -3,7 +3,7 @@
 # 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: 15bacd81d8b826a7ae2fef591b17ea988be9763d23b09a1f00b58be87c65310f
+_hash: 9f92368365708cb2c3720db344c9a1e4f43acb64808c7d59336769be0ff7932a
 file:
   batch:
     create_procedure_statement:
@@ -24,13 +24,38 @@ file:
               data_type:
                 identifier: DATE
               comparison_operator: '='
-              function:
-                function_name:
-                  function_name_identifier: GETDATE
+              expression:
+                function:
+                  function_name:
+                    function_name_identifier: GETDATE
+                  bracketed:
+                    start_bracket: (
+                    end_bracket: )
+              statement_terminator: ;
+        - statement:
+            declare_segment:
+              keyword: DECLARE
+              parameter: '@EOMONTH'
+              data_type:
+                identifier: DATE
+              comparison_operator: '='
+              expression:
                 bracketed:
                   start_bracket: (
+                  expression:
+                    literal: "'1900-01-01'"
+                  end_bracket: )
+        - statement:
+            set_segment:
+              keyword: SET
+              parameter: '@EOMONTH'
+              comparison_operator: '='
+              expression:
+                bracketed:
+                  start_bracket: (
+                  expression:
+                    literal: "'2000-01-01'"
                   end_bracket: )
-              statement_terminator: ;
         - statement:
             if_then_statement:
               keyword: IF

diff --git a/test/fixtures/dialects/tsql/print.yml b/test/fixtures/dialects/tsql/print.yml
@@ -3,7 +3,7 @@
 # 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: ae5fdbfdc047812dbf5359045d2a417e4e7c09cbe83683737b5e96d1900da119
+_hash: ccc3078339f97b77c695b4327a6d0b12d1d6007c5abfaa5460ca49f812164c4e
 file:
   batch:
   - statement:
@@ -18,7 +18,8 @@ file:
               literal: '20'
             end_bracket: )
         comparison_operator: '='
-        literal: "'Test Print'"
+        expression:
+          literal: "'Test Print'"
   - statement:
       print_statement:
         keyword: PRINT

diff --git a/test/fixtures/dialects/tsql/stored_procedured_mixed_statements.yml b/test/fixtures/dialects/tsql/stored_procedured_mixed_statements.yml
@@ -3,7 +3,7 @@
 # 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: 5a9199e249ff2e03ef14c6833eb96ea0924cd942867ba7b8fa5f128d3a0f3986
+_hash: 1d71a43a3c51b8860b18d9da2256b0c0d635a475a00841b5b9fc8db8ba541e38
 file:
   batch:
     create_procedure_statement:
@@ -22,12 +22,13 @@ file:
             data_type:
               identifier: DATE
             comparison_operator: '='
-            function:
-              function_name:
-                function_name_identifier: GETDATE
-              bracketed:
-                start_bracket: (
-                end_bracket: )
+            expression:
+              function:
+                function_name:
+                  function_name_identifier: GETDATE
+                bracketed:
+                  start_bracket: (
+                  end_bracket: )
       - statement:
           drop_statement:
           - keyword: DROP