Repository: David-Wobrock/sqlvalidator

Issue:
Implement automatic 'SELECT' detection
Currently to format a SQL string, you need to mark it with `# sqlformat` comment (so that it's easy to not format a SQL string when the sqlvalidator is buggy).

We should also be able to automatically let the `SELECT` statements be formatted (and have a comment to not format some)

But should we have both ways or just stick to one?

Discussion:

Patch:
diff --git a/README.md b/README.md
index fbfa604..2461e3f 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,7 @@ SQL queries formatting, syntactic and semantic validation
 _sql.py_
 ```
 def fun():
-    return "select col1, column2 from table"  # sqlformat
+    return "select col1, column2 from table"
 ```
 
 Command line:
@@ -33,12 +33,11 @@ SELECT
  col1,
  column2
 FROM table
-"""  # sqlformat
+"""
 
 ```
 
-The `sqlformat` comment is required to indicated to `sqlvalidator` that this string is a SQL statement
-and should be formatted.
+A `nosqlformat` comment can be appended to indicate to `sqlvalidator` that this string should not be formatted.
 
 
 ### Check SQL format
diff --git a/sqlvalidator/file_formatter.py b/sqlvalidator/file_formatter.py
index 04fa268..d93f429 100644
--- a/sqlvalidator/file_formatter.py
+++ b/sqlvalidator/file_formatter.py
@@ -5,7 +5,7 @@ from typing import IO, List, Optional, Set, Tuple
 
 from . import sql_formatter
 
-SQLFORMAT_COMMENT = "sqlformat"
+NO_SQLFORMAT_COMMENT = "nosqlformat"
 
 POSSIBLE_QUOTES = ('"""', "'''", '"', "'")
 STRING_PREFIXES = (
@@ -90,6 +90,10 @@ def format_sql_string(sql_string):
     return quoted_sql
 
 
+def is_select_string(token_value: str) -> bool:
+    return token_value.lower().lstrip(" \n'\"`ufbr").startswith("select")
+
+
 def get_formatted_file_content(file: IO) -> Tuple[int, str]:
     count_changed_sql = 0
     tokens = []
@@ -100,19 +104,31 @@ def get_formatted_file_content(file: IO) -> Tuple[int, str]:
         nonlocal tokens
         nonlocal count_changed_sql
 
-        following_tokens = []
+        if not is_select_string(token_value):
+            tokens.append((token_type, token_value, starting, ending, line))
+            return
 
-        # Find if the next comment without a string in-between
+        following_tokens = []
+        # Find if the next comment without a string in-between if there is one
         next_token, next_token_value, next_starting, next_ending, next_line = next(
             token_generator, (None, None, None, None, None)
         )
         if next_token is None:
-            tokens.append((token_type, token_value, starting, ending, line))
+            formatted_sql = format_sql_string(token_value)
+            tokens.append((token_type, formatted_sql, starting, ending, line))
+            tokens += following_tokens
+            if formatted_sql != token_value:
+                count_changed_sql += 1
             return
+
         following_tokens.append(
             (next_token, next_token_value, next_starting, next_ending, next_line)
         )
-        while next_token != tokenize.COMMENT and next_token != tokenize.STRING:
+        while (
+            next_token
+            and next_token != tokenize.COMMENT
+            and next_token != tokenize.STRING
+        ):
             next_token, next_token_value, next_starting, next_ending, next_line = next(
                 token_generator, (None, None, None, None, None)
             )
@@ -124,15 +140,17 @@ def get_formatted_file_content(file: IO) -> Tuple[int, str]:
                 (next_token, next_token_value, next_starting, next_ending, next_line)
             )
 
-        if next_token == tokenize.COMMENT and SQLFORMAT_COMMENT in next_token_value:
+        if not (
+            next_token == tokenize.COMMENT and NO_SQLFORMAT_COMMENT in next_token_value
+        ):
             formatted_sql = format_sql_string(token_value)
             tokens.append((token_type, formatted_sql, starting, ending, line))
-            tokens += following_tokens
-
             if formatted_sql != token_value:
                 count_changed_sql += 1
-        elif next_token == tokenize.STRING:
+        else:
             tokens.append((token_type, token_value, starting, ending, line))
+
+        if next_token == tokenize.STRING:
             tokens += following_tokens[:-1]
             handle_string_token(
                 token_generator,
@@ -143,7 +161,6 @@ def get_formatted_file_content(file: IO) -> Tuple[int, str]:
                 next_line,
             )
         else:
-            tokens.append((token_type, token_value, starting, ending, line))
             tokens += following_tokens
 
     token_generator = tokenize.generate_tokens(file.readline)


Test patch:
diff --git a/tests/integration/test_file_formatting.py b/tests/integration/test_file_formatting.py
index a1fa99b..bafcf17 100644
--- a/tests/integration/test_file_formatting.py
+++ b/tests/integration/test_file_formatting.py
@@ -4,51 +4,90 @@ from sqlvalidator import file_formatter
 
 
 def test_format_file():
-    file_content = "'''select id from table_stmt;'''#sqlformat"
+    file_content = "'''select id from table_stmt;'''"
     input_file = StringIO(file_content)
     num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
-    expected_file_content = "'''\nSELECT id\nFROM table_stmt;\n'''#sqlformat"
+    expected_file_content = "'''\nSELECT id\nFROM table_stmt;\n'''"
     assert num_changed_sql == 1
     assert new_content == expected_file_content
 
 
 def test_format_file_multilines():
-    file_content = "'select id from table_stmt'  #sqlformat"
+    file_content = "'select id from table_stmt'"
     input_file = StringIO(file_content)
     _, new_content = file_formatter.get_formatted_file_content(input_file)
-    expected_file_content = "'''\nSELECT id\nFROM table_stmt\n'''  #sqlformat"
+    expected_file_content = "'''\nSELECT id\nFROM table_stmt\n'''"
     assert new_content == expected_file_content
 
 
 def test_format_file_multiline_sql():
-    file_content = "'''\nselect\n   id  \n   from table_stmt\n;\n'''#sqlformat"
+    file_content = "'''\nselect\n   id  \n   from table_stmt\n;\n'''"
     input_file = StringIO(file_content)
     num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
-    expected_file_content = "'''\nSELECT id\nFROM table_stmt;\n'''#sqlformat"
+    expected_file_content = "'''\nSELECT id\nFROM table_stmt;\n'''"
     assert num_changed_sql == 1
     assert new_content == expected_file_content
 
 
 def test_format_file_multiline_sql_with_prefix():
-    file_content = "r'''\nselect\n   id  \n   from table_stmt\n;\n'''#sqlformat"
+    file_content = "r'''\nselect\n   id  \n   from table_stmt\n;\n'''"
     input_file = StringIO(file_content)
     num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
-    expected_file_content = "r'''\nSELECT id\nFROM table_stmt;\n'''#sqlformat"
+    expected_file_content = "r'''\nSELECT id\nFROM table_stmt;\n'''"
     assert num_changed_sql == 1
     assert new_content == expected_file_content
 
 
-def test_missing_formatting_comment():
-    file_content = "'select id from table_stmt'"
+def test_format_sql_string_with_strformat_call():
+    file_content = "'select id from {}'.format(__file__)"
     input_file = StringIO(file_content)
-    _, new_content = file_formatter.get_formatted_file_content(input_file)
+    num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
+    expected_file_content = "'''\nSELECT id\nFROM {}\n'''.format(__file__)"
+    assert num_changed_sql == 1
+    assert new_content == expected_file_content
+
+
+def test_no_format_file():
+    file_content = "'select id from t'  # nosqlformat"
+    input_file = StringIO(file_content)
+    num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
+    assert num_changed_sql == 0
     assert new_content == file_content
 
 
-def test_format_sql_string_with_strformat_call():
-    file_content = "'select id from {}'.format(__file__) # sqlformat"
+def test_no_format_sql_string_with_strformat_call():
+    file_content = "'select id from {}'.format(__file__)  # nosqlformat"
     input_file = StringIO(file_content)
     num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
-    expected_file_content = "'''\nSELECT id\nFROM {}\n'''.format(__file__) # sqlformat"
+    assert num_changed_sql == 0
+    assert new_content == file_content
+
+
+def test_two_following_statements():
+    file_content = "sql = ('select * from t;', 'select count(*) from t;')"
+    input_file = StringIO(file_content)
+    num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
+    expected_file_content = (
+        "sql = ('''\nSELECT *\nFROM t;\n''', '''\nSELECT COUNT(*)\nFROM t;\n''')"
+    )
+    assert num_changed_sql == 2
+    assert new_content == expected_file_content
+
+
+def test_two_following_statements_second_noformat():
+    file_content = (
+        "sql = (\n'select * from t;', 'select count(*) from t;'  # nosqlformat\n)"
+    )
+    input_file = StringIO(file_content)
+    num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
+    expected_file_content = "sql = (\n'''\nSELECT *\nFROM t;\n''', 'select count(*) from t;'  # nosqlformat\n)"  # noqa
     assert num_changed_sql == 1
     assert new_content == expected_file_content
+
+
+def test_two_following_statements_both_no_format():
+    file_content = "sql = (\n'select * from t;'  # nosqlformat\n, 'select count(*) from t;'  # nosqlformat\n)"  # noqa
+    input_file = StringIO(file_content)
+    num_changed_sql, new_content = file_formatter.get_formatted_file_content(input_file)
+    assert num_changed_sql == 0
+    assert new_content == file_content