Introduce isolation mode option (#35)

Isolation mode option

Author: EnchoMishinevQC

src/pytsql/tsql.py

@@ -83,16 +83,19 @@ def __init__(self):
         # batches.
         self.dynamics: List[str] = []
 
-    def visit(self, tree: tsqlParser.Sql_clauseContext) -> str:
+    def visit(
+        self, tree: tsqlParser.Sql_clauseContext, prepend_dynamics: bool = True
+    ) -> str:
         dynamics = self.dynamics[:]
 
         chunks = tree.accept(self)
 
         # CREATE SCHEMA/VIEW must be the only statement in a batch
-        if tree.ddl_clause() is not None and (
+        is_create_schema_or_view = tree.ddl_clause() is not None and (
             tree.ddl_clause().create_schema() is not None
             or tree.ddl_clause().create_view() is not None
-        ):
+        )
+        if not prepend_dynamics or is_create_schema_or_view:
             return " ".join(chunks)
 
         return " ".join(dynamics + chunks)
@@ -143,7 +146,7 @@ def syntaxError(
         raise ValueError(f"Error parsing SQL script: {error_message}")
 
 
-def _split(code: str) -> List[str]:
+def _split(code: str, isolate_top_level_statements: bool = True) -> List[str]:
     if not USE_CPP_IMPLEMENTATION:
         warnings.warn(
             "Can not find C++ version of the parser, Python version will be used instead."
@@ -160,15 +163,24 @@ def _split(code: str) -> List[str]:
     tree = parse(InputStream(data=code), "tsql_file", error_listener)
     visitor = _TSQLVisitor()
 
-    # Our current definition of a 'batch' is a single top-level SQL clause.
+    # Our current definition of a 'batch' in isolation mode is a single top-level SQL clause.
     # Note that this differs from the grammar definition of a batch, which is
-    # a group of clauses between GO statements
+    # a group of clauses between GO statements. The latter matches the definition of batches
+    # in non-isolation mode.
     batches = []
     for batch in tree.batch():
+        clauses = []
+        first_clause_in_batch = True
         for sql_clause in batch.sql_clauses().sql_clause():
-            batch_query = visitor.visit(sql_clause)
-            if batch_query != "":
-                batches.append(batch_query)
+            prepend_dynamics = first_clause_in_batch or isolate_top_level_statements
+            clause = visitor.visit(sql_clause, prepend_dynamics=prepend_dynamics)
+            if clause != "":
+                clauses.append(clause)
+                first_clause_in_batch = False
+        if isolate_top_level_statements:
+            batches.extend(clauses)
+        else:
+            batches.append("\n".join(clauses))
 
     logging.info("SQL script parsed successfully.")
 
@@ -186,6 +198,7 @@ def executes(
     code: str,
     engine: sqlalchemy.engine.Engine,
     parameters: Optional[Dict[str, Any]] = None,
+    isolate_top_level_statements=True,
 ) -> None:
     """Execute a given sql string through a sqlalchemy.engine.Engine connection.
 
@@ -197,6 +210,7 @@ def executes(
     code T-SQL string to be executed
     engine (sqlalchemy.engine.Engine): established mssql connection
     parameters An optional dictionary of parameters to substituted in the sql script
+    isolate_top_level_statements: whether to execute statements one by one or in whole batches
 
     Returns
     -------
@@ -209,7 +223,7 @@ def executes(
         # connection is closed. Caveat: sqlalchemy engines can pool connections, so we still have to drop it preemtively.
         conn.execute(f"DROP TABLE IF EXISTS {_PRINTS_TABLE}")
         conn.execute(f"CREATE TABLE {_PRINTS_TABLE} (p NVARCHAR(4000))")
-        for batch in _split(parametrized_code):
+        for batch in _split(parametrized_code, isolate_top_level_statements):
             conn.execute(batch)
             _fetch_and_clear_prints(conn)
 
@@ -218,6 +232,7 @@ def execute(
     path: Union[str, Path],
     engine: sqlalchemy.engine.Engine,
     parameters: Optional[Dict[str, Any]] = None,
+    isolate_top_level_statements=True,
     encoding: str = "utf-8",
 ) -> None:
     """Execute a given sql script through a sqlalchemy.engine.Engine connection.
@@ -227,10 +242,11 @@ def execute(
     path (Path or str): Path to the sql file to be executed
     engine (sqlalchemy.engine.Engine): established mssql connection
     encoding: file encoding of the sql script (default: utf-8)
+    isolate_top_level_statements: whether to execute statements one by one or in whole batches
 
     Returns
     -------
     None
 
     """
-    executes(_code(path, encoding), engine, parameters)
+    executes(_code(path, encoding), engine, parameters, isolate_top_level_statements)

tests/integration/test_non_isolation_mode.py

@@ -0,0 +1,45 @@
+import logging
+
+from pytsql import executes
+
+
+def test_rowcount(engine, caplog):
+    caplog.set_level(logging.INFO)
+
+    seed = """
+    USE [tempdb]
+    GO
+    DROP TABLE IF EXISTS [test_table]
+    CREATE TABLE [test_table] (
+        col VARCHAR(3)
+    )
+    GO
+    INSERT INTO [test_table] (col)
+    VALUES ('A'), ('AB'), ('ABC')
+    PRINT('Affected ' + CAST(@@ROWCOUNT AS VARCHAR) + ' rows')
+    """
+
+    executes(seed, engine, isolate_top_level_statements=False)
+
+    assert "Affected 3 rows" in caplog.text
+
+
+def test_semi_persistent_set(engine, caplog):
+    caplog.set_level(logging.INFO)
+
+    seed = """
+    DECLARE @A INT = 12
+    DECLARE @B INT = 34
+    SET @A = 56
+    SET @B = 78
+    PRINT(@A)
+    GO
+    PRINT(@B)
+    """
+
+    executes(seed, engine, isolate_top_level_statements=False)
+
+    assert "56" in caplog.text
+    assert "34" in caplog.text
+    assert "12" not in caplog.text
+    assert "78" not in caplog.text

tests/unit/test_dynamics.py

@@ -44,3 +44,28 @@ def test_dont_append_dynamics_on_create_schema():
         """,
     )
     assert splits[3] == "CREATE VIEW y AS SELECT 1"
+
+
+def test_dynamics_no_isolation():
+    seed = """
+    DECLARE @A INT = 5
+    SELECT @A
+    GO
+    SELECT @A
+    """
+    splits = _split(seed, isolate_top_level_statements=False)
+    assert len(splits) == 2
+    assert_strings_equal_disregarding_whitespace(
+        splits[0],
+        """
+        DECLARE @A INT = 5
+        SELECT @A
+        """,
+    )
+    assert_strings_equal_disregarding_whitespace(
+        splits[1],
+        """
+        DECLARE @A INT = 5
+        SELECT @A
+        """,
+    )

tests/unit/test_split.py

@@ -313,3 +313,15 @@ def test_linked_server():
     FROM linked_server.db.schm.table2
     """
     assert len(_split(seed)) == 1
+
+
+def test_rowcount_no_isolation():
+    seed = """
+    USE new_db
+    PRINT('Using new_db')
+    GO
+    SELECT * FROM my_table
+    PRINT(@@ROWCOUNT)
+    """
+    assert len(_split(seed)) == 4
+    assert len(_split(seed, isolate_top_level_statements=False)) == 2