Fix formatting on INSERT (fixes andialbrecht#329)

fredyw · andialbrecht · commit fa7def01d00f · 2019-03-10T08:03:05.000+01:00
This patch fixes the formatting on INSERT by creating a new
instance of sql.Values to group all the values.

SQL: insert into foo values (1, 'foo'), (2, 'bar'), (3, 'baz')

Before:
insert into foo
values (1,
        'foo'), (2,
                 'bar'), (3,
                          'baz')

After:
insert into foo
values (1, 'foo'),
       (2, 'bar'),
       (3, 'baz')
diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
@@ -327,6 +327,18 @@ def align_comments(tlist):
         tidx, token = tlist.token_next_by(i=sql.Comment, idx=tidx)
 
 
+def group_values(tlist):
+    tidx, token = tlist.token_next_by(m=(T.Keyword, 'VALUES'))
+    start_idx = tidx
+    end_idx = -1
+    while token:
+        if isinstance(token, sql.Parenthesis):
+            end_idx = tidx
+        tidx, token = tlist.token_next(tidx)
+    if end_idx != -1:
+        tlist.group_tokens(sql.Values, start_idx, end_idx, extend=True)
+
+
 def group(stmt):
     for func in [
         group_comments,
@@ -354,6 +366,7 @@ def group(stmt):
 
         align_comments,
         group_identifier_list,
+        group_values,
     ]:
         func(stmt)
     return stmt
diff --git a/sqlparse/filters/reindent.py b/sqlparse/filters/reindent.py
@@ -133,7 +133,7 @@ def _process_identifierlist(self, tlist):
             first = next(identifiers.pop(0).flatten())
             num_offset = 1 if self.char == '\t' else self._get_offset(first)
 
-        if not tlist.within(sql.Function):
+        if not tlist.within(sql.Function) and not tlist.within(sql.Values):
             with offset(self, num_offset):
                 position = 0
                 for token in identifiers:
@@ -206,6 +206,23 @@ def _process_case(self, tlist):
             if end_idx is not None:
                 tlist.insert_before(end_idx, self.nl())
 
+    def _process_values(self, tlist):
+        tlist.insert_before(0, self.nl())
+        tidx, token = tlist.token_next_by(i=sql.Parenthesis)
+        first_token = token
+        while token:
+            ptidx, ptoken = tlist.token_next_by(m=(T.Punctuation, ','),
+                                                idx=tidx)
+            if ptoken:
+                if self.comma_first:
+                    adjust = -2
+                    offset = self._get_offset(first_token) + adjust
+                    tlist.insert_before(ptoken, self.nl(offset))
+                else:
+                    tlist.insert_after(ptoken,
+                                       self.nl(self._get_offset(token)))
+            tidx, token = tlist.token_next_by(i=sql.Parenthesis, idx=tidx)
+
     def _process_default(self, tlist, stmts=True):
         self._split_statements(tlist) if stmts else None
         self._split_kwds(tlist)
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
@@ -621,3 +621,7 @@ class Begin(TokenList):
 
 class Operation(TokenList):
     """Grouping of operations"""
+
+
+class Values(TokenList):
+    """Grouping of values"""
diff --git a/tests/test_format.py b/tests/test_format.py
@@ -554,6 +554,49 @@ def test_identifier_and_functions(self):
             '       nvl(1)',
             'from dual'])
 
+    def test_insert_values(self):
+        # issue 329
+        f = lambda sql: sqlparse.format(sql, reindent=True)
+        s = 'insert into foo values (1, 2)'
+        assert f(s) == '\n'.join([
+            'insert into foo',
+            'values (1, 2)'])
+
+        s = 'insert into foo values (1, 2), (3, 4), (5, 6)'
+        assert f(s) == '\n'.join([
+            'insert into foo',
+            'values (1, 2),',
+            '       (3, 4),',
+            '       (5, 6)'])
+
+        s = 'insert into foo(a, b) values (1, 2), (3, 4), (5, 6)'
+        assert f(s) == '\n'.join([
+            'insert into foo(a, b)',
+            'values (1, 2),',
+            '       (3, 4),',
+            '       (5, 6)'])
+
+        f = lambda sql: sqlparse.format(sql, reindent=True,
+                                        comma_first=True)
+        s = 'insert into foo values (1, 2)'
+        assert f(s) == '\n'.join([
+            'insert into foo',
+            'values (1, 2)'])
+
+        s = 'insert into foo values (1, 2), (3, 4), (5, 6)'
+        assert f(s) == '\n'.join([
+            'insert into foo',
+            'values (1, 2)',
+            '     , (3, 4)',
+            '     , (5, 6)'])
+
+        s = 'insert into foo(a, b) values (1, 2), (3, 4), (5, 6)'
+        assert f(s) == '\n'.join([
+            'insert into foo(a, b)',
+            'values (1, 2)',
+            '     , (3, 4)',
+            '     , (5, 6)'])
+
 
 class TestOutputFormat(object):
     def test_python(self):
diff --git a/tests/test_grouping.py b/tests/test_grouping.py
@@ -49,10 +49,13 @@ def test_grouping_identifiers():
     assert str(parsed) == s
     assert isinstance(parsed.tokens[-1].tokens[3], sql.Identifier)
 
-    s = "INSERT INTO `test` VALUES('foo', 'bar');"
-    parsed = sqlparse.parse(s)[0]
-    types = [l.ttype for l in parsed.tokens if not l.is_whitespace]
-    assert types == [T.DML, T.Keyword, None, T.Keyword, None, T.Punctuation]
+    for s in ["INSERT INTO `test` VALUES('foo', 'bar');",
+              "INSERT INTO `test` VALUES(1, 2), (3, 4), (5, 6);",
+              "INSERT INTO `test(a, b)` VALUES(1, 2), (3, 4), (5, 6);"]:
+        parsed = sqlparse.parse(s)[0]
+        types = [l.ttype for l in parsed.tokens if not l.is_whitespace]
+        assert types == [T.DML, T.Keyword, None, None, T.Punctuation]
+        assert isinstance(parsed.tokens[6], sql.Values)
 
     s = "select 1.0*(a+b) as col, sum(c)/sum(d) from myschema.mytable"
     parsed = sqlparse.parse(s)[0]
