Merge pull request #248 from MarketSquare/split_script_feature

New keyword "Split SQL Script" and external script parser

Author: amochin

pyproject.toml

@@ -2,7 +2,8 @@
 requires = [
     "setuptools>=61.0",
     "robotframework>=5.0.1",
-    "robotframework-assertion-engine"
+    "robotframework-assertion-engine",
+    "sqlparse"
     ]
 build-backend = "setuptools.build_meta"
 
@@ -15,7 +16,8 @@ readme = "README.md"
 requires-python = ">=3.8.1"
 dependencies = [
     "robotframework>=5.0.1",
-    "robotframework-assertion-engine"
+    "robotframework-assertion-engine",
+    "sqlparse"
 ]
 classifiers = [
     "Programming Language :: Python :: 3",

src/DatabaseLibrary/query.py

@@ -18,6 +18,7 @@
 import sys
 from typing import List, Optional, Tuple
 
+import sqlparse
 from robot.api import logger
 from robot.utils.dotdict import DotDict
 
@@ -263,6 +264,7 @@ def execute_sql_script(
         no_transaction: bool = False,
         alias: Optional[str] = None,
         split: bool = True,
+        external_parser=False,
         *,
         sqlScriptFileName: Optional[str] = None,
         sansTran: Optional[bool] = None,
@@ -274,6 +276,8 @@ def execute_sql_script(
         Set ``split`` to _False_ to disable this behavior  - in this case the entire script content
         will be passed to the database module for execution as a single command.
 
+        Set `external_parser` to _True_ to use the external `sqlparse` library for splitting the script.
+
         Set ``no_transaction`` to _True_ to run command without explicit transaction commit
         or rollback in case of error.
         See `Commit behavior` for details.
@@ -293,84 +297,104 @@ def execute_sql_script(
         | Execute SQL Script | insert_data_in_person_table.sql | split=False |
         """
         db_connection = self.connection_store.get_connection(alias)
-        with open(script_path, encoding="UTF-8") as sql_file:
-            cur = None
-            try:
-                cur = db_connection.client.cursor()
-                if not split:
+        cur = None
+        try:
+            cur = db_connection.client.cursor()
+            if not split:
+                with open(script_path, encoding="UTF-8") as sql_file:
                     logger.info("Statements splitting disabled - pass entire script content to the database module")
                     self._execute_sql(
                         cur,
                         sql_file.read(),
                         omit_trailing_semicolon=db_connection.omit_trailing_semicolon,
                     )
-                else:
-                    logger.info("Splitting script file into statements...")
-                    statements_to_execute = []
-                    current_statement = ""
-                    inside_statements_group = False
-                    proc_start_pattern = re.compile("create( or replace)? (procedure|function){1}( )?")
+            else:
+                statements_to_execute = self.split_sql_script(script_path, external_parser=external_parser)
+                for statement in statements_to_execute:
                     proc_end_pattern = re.compile("end(?!( if;| loop;| case;| while;| repeat;)).*;()?")
-                    for line in sql_file:
-                        line = line.strip()
-                        if line.startswith("#") or line.startswith("--") or line == "/":
+                    line_ends_with_proc_end = re.compile(r"(\s|;)" + proc_end_pattern.pattern + "$")
+                    omit_semicolon = not line_ends_with_proc_end.search(statement.lower())
+                    self._execute_sql(cur, statement, omit_semicolon)
+            self._commit_if_needed(db_connection, no_transaction)
+        except Exception as e:
+            self._rollback_and_raise(db_connection, no_transaction, e)
+
+    def split_sql_script(
+        self,
+        script_path: str,
+        external_parser=False,
+    ):
+        """
+        Splits the content of the SQL script file loaded from `script_path` into individual SQL commands
+        and returns them as a list of strings.
+        SQL commands are expected to be delimited by a semicolon (';').
+
+        Set `external_parser` to _True_ to use the external `sqlparse` library.
+        """
+        with open(script_path, encoding="UTF-8") as sql_file:
+            logger.info("Splitting script file into statements...")
+            statements_to_execute = []
+            if external_parser:
+                statements_to_execute = sqlparse.split(sql_file.read())
+            else:
+                current_statement = ""
+                inside_statements_group = False
+                proc_start_pattern = re.compile("create( or replace)? (procedure|function){1}( )?")
+                proc_end_pattern = re.compile("end(?!( if;| loop;| case;| while;| repeat;)).*;()?")
+                for line in sql_file:
+                    line = line.strip()
+                    if line.startswith("#") or line.startswith("--") or line == "/":
+                        continue
+
+                    # check if the line matches the creating procedure regexp pattern
+                    if proc_start_pattern.match(line.lower()):
+                        inside_statements_group = True
+                    elif line.lower().startswith("begin"):
+                        inside_statements_group = True
+
+                    # semicolons inside the line? use them to separate statements
+                    # ... but not if they are inside a begin/end block (aka. statements group)
+                    sqlFragments = line.split(";")
+                    # no semicolons
+                    if len(sqlFragments) == 1:
+                        current_statement += line + " "
+                        continue
+                    quotes = 0
+                    # "select * from person;" -> ["select..", ""]
+                    for sqlFragment in sqlFragments:
+                        if len(sqlFragment.strip()) == 0:
                             continue
 
-                        # check if the line matches the creating procedure regexp pattern
-                        if proc_start_pattern.match(line.lower()):
+                        if inside_statements_group:
+                            # if statements inside a begin/end block have semicolns,
+                            # they must persist - even with oracle
+                            sqlFragment += "; "
+
+                        if proc_end_pattern.match(sqlFragment.lower()):
+                            inside_statements_group = False
+                        elif proc_start_pattern.match(sqlFragment.lower()):
                             inside_statements_group = True
-                        elif line.lower().startswith("begin"):
+                        elif sqlFragment.lower().startswith("begin"):
                             inside_statements_group = True
 
-                        # semicolons inside the line? use them to separate statements
-                        # ... but not if they are inside a begin/end block (aka. statements group)
-                        sqlFragments = line.split(";")
-                        # no semicolons
-                        if len(sqlFragments) == 1:
-                            current_statement += line + " "
-                            continue
-                        quotes = 0
-                        # "select * from person;" -> ["select..", ""]
-                        for sqlFragment in sqlFragments:
-                            if len(sqlFragment.strip()) == 0:
-                                continue
-
-                            if inside_statements_group:
-                                # if statements inside a begin/end block have semicolns,
-                                # they must persist - even with oracle
-                                sqlFragment += "; "
-
-                            if proc_end_pattern.match(sqlFragment.lower()):
-                                inside_statements_group = False
-                            elif proc_start_pattern.match(sqlFragment.lower()):
-                                inside_statements_group = True
-                            elif sqlFragment.lower().startswith("begin"):
-                                inside_statements_group = True
-
-                            # check if the semicolon is a part of the value (quoted string)
-                            quotes += sqlFragment.count("'")
-                            quotes -= sqlFragment.count("\\'")
-                            inside_quoted_string = quotes % 2 != 0
-                            if inside_quoted_string:
-                                sqlFragment += ";"  # restore the semicolon
-
-                            current_statement += sqlFragment
-                            if not inside_statements_group and not inside_quoted_string:
-                                statements_to_execute.append(current_statement.strip())
-                                current_statement = ""
-                                quotes = 0
-
-                    current_statement = current_statement.strip()
-                    if len(current_statement) != 0:
-                        statements_to_execute.append(current_statement)
-
-                    for statement in statements_to_execute:
-                        line_ends_with_proc_end = re.compile(r"(\s|;)" + proc_end_pattern.pattern + "$")
-                        omit_semicolon = not line_ends_with_proc_end.search(statement.lower())
-                        self._execute_sql(cur, statement, omit_semicolon)
-                self._commit_if_needed(db_connection, no_transaction)
-            except Exception as e:
-                self._rollback_and_raise(db_connection, no_transaction, e)
+                        # check if the semicolon is a part of the value (quoted string)
+                        quotes += sqlFragment.count("'")
+                        quotes -= sqlFragment.count("\\'")
+                        inside_quoted_string = quotes % 2 != 0
+                        if inside_quoted_string:
+                            sqlFragment += ";"  # restore the semicolon
+
+                        current_statement += sqlFragment
+                        if not inside_statements_group and not inside_quoted_string:
+                            statements_to_execute.append(current_statement.strip())
+                            current_statement = ""
+                            quotes = 0
+
+                current_statement = current_statement.strip()
+                if len(current_statement) != 0:
+                    statements_to_execute.append(current_statement)
+
+            return statements_to_execute
 
     @renamed_args(
         mapping={

test/resources/script_file_tests/split_commands.sql

@@ -0,0 +1,2 @@
+SELECT * FROM person;
+SELECT * FROM person WHERE id=1;

test/tests/common_tests/script_files.robot

@@ -6,6 +6,9 @@ Suite Teardown      Disconnect From Database
 Test Setup          Create Person Table
 Test Teardown       Drop Tables Person And Foobar
 
+*** Variables ***
+${Script files dir}    ${CURDIR}/../../resources/script_file_tests
+
 
 *** Test Cases ***
 Semicolons As Statement Separators In One Line
@@ -35,9 +38,30 @@ Semicolons And Quotes In Values
     Should Be Equal As Strings    ${results}[0]    (5, 'Miles', "O'Brian")
     Should Be Equal As Strings    ${results}[1]    (6, 'Keiko', "O'Brian")
 
+Split Script Into Statements - Internal Parser
+    Insert Data In Person Table Using SQL Script
+    @{Expected commands}=    Create List
+    ...    SELECT * FROM person
+    ...    SELECT * FROM person WHERE id=1
+    ${extracted commands}=    Split Sql Script    ${Script files dir}/split_commands.sql
+    Lists Should Be Equal    ${Expected commands}    ${extracted commands}
+    FOR    ${command}    IN    @{extracted commands}
+        ${results}=    Query    ${command}
+    END
+
+Split Script Into Statements - External Parser
+    Insert Data In Person Table Using SQL Script
+    @{Expected commands}=    Create List
+    ...    SELECT * FROM person;
+    ...    SELECT * FROM person WHERE id=1;
+    ${extracted commands}=    Split Sql Script    ${Script files dir}/split_commands.sql    external_parser=True
+    Lists Should Be Equal    ${Expected commands}    ${extracted commands}
+    FOR    ${command}    IN    @{extracted commands}
+        ${results}=    Query    ${command}
+    END
+
 
 *** Keywords ***
 Run SQL Script File
     [Arguments]    ${File Name}
-    ${Script files dir}=    Set Variable    ${CURDIR}/../../resources/script_file_tests
     Execute Sql Script    ${Script files dir}/${File Name}.sql