Copy only top-level DECLARE statements (#44)

In case of isolate_top_level_statements=True, copy DECLARE statements to subsequent statements only if they are top level.

Author: windiana42

src/pytsql/grammar/tsqlParser.py

@@ -24890,6 +24890,13 @@ def getRuleIndex(self):
             return tsqlParser.RULE_data_type
 
 
+    @staticmethod
+    def is_top_level_statement(node: ParserRuleContext):
+        """Check wether node is a top level SQL statement."""
+        cur = node.parentCtx
+        while isinstance(cur, tsqlParser.Sql_clauseContext) or isinstance(cur, tsqlParser.Sql_clausesContext):
+            cur = cur.parentCtx
+        return isinstance(cur, tsqlParser.BatchContext)
 
 
     def data_type(self):

src/pytsql/tsql.py

@@ -112,7 +112,7 @@ def visitChildren(self, node: antlr4.ParserRuleContext) -> List[str]:
         else:
             result = super().visitChildren(node)
 
-        if isinstance(node, tsqlParser.Declare_statementContext):
+        if isinstance(node, tsqlParser.Declare_statementContext) and tsqlParser.is_top_level_statement(node):
             self.dynamics.extend(result)
 
         return result

tests/integration/test_multiple_statements.py

@@ -182,6 +182,23 @@ def test_stored_procedure_declaration(engine):
     executes(statement, engine, None)
 
 
+def test_top_level_declaration(engine):
+    statement = """
+    DROP DATABASE IF EXISTS top_level_declaration
+    CREATE DATABASE top_level_declaration
+    USE top_level_declaration
+GO
+
+DECLARE @Current AS DATE = '2022-01-01'
+GO
+SELECT @Current as a INTO dummy01
+GO
+SELECT @Current as b INTO dummy02
+GO
+"""
+    executes(statement, engine, None)
+
+
 def get_table(
     engine: Engine, table_name: str, schema: Optional[str] = None
 ) -> sa.Table:

tests/unit/test_dynamics.py

@@ -7,20 +7,36 @@ def test_declaration_in_control_flow():
     seed = """
     IF 1 = 1
         DECLARE @A INT = 5
-    SELECT * FROM x
+    SELECT @A
     """
     splits = _split(seed)
     assert len(splits) == 2
 
     assert_strings_equal_disregarding_whitespace(
         splits[0], "IF 1 = 1 DECLARE @A INT = 5"
     )
+    # unfortunately we can't be right here because otherwise we would need to get
+    # the output of the declaration
     assert_strings_equal_disregarding_whitespace(
         splits[1],
-        """
+        """SELECT @A""",
+    )
+
+
+def test_select_in_control_flow():
+    seed = """
+    IF 1 = 0
+    BEGIN
         DECLARE @A INT = 5
         SELECT * FROM x
-        """,
+    END
+    """
+    splits = _split(seed)
+    assert len(splits) == 1
+
+    # this is beyond the complexity we want to manage with isolate_top_level_statements=True
+    assert_strings_equal_disregarding_whitespace(
+        splits[0], "IF 1 = 0 BEGIN DECLARE @A INT = 5 SELECT * FROM x END"
     )