Fix status line display of the executed SQL command

`sqlparse.sql.Statement.token_first()` provides the `skip_cm` option
already:

  if *skip_cm* is ``True`` (default: ``False``), comments are
  ignored too.

Using `statement.get_type()` would be even more DWIM:

  The returned value is a string holding an upper-cased reprint of
  the first DML or DDL keyword. If the first token in this group
  isn't a DML or DDL keyword "UNKNOWN" is returned.

  Whitespaces and comments at the beginning of the statement
  are ignored.

However, that would introduce a regression that the outcome of
non-DML/DDL keywords like `BEGIN` would yield `UNKNOWN`.

CHANGES.txt

@@ -5,6 +5,9 @@ Changes for crash
 Unreleased
 ==========
 
+- Fixed status line display of the executed SQL command
+
+
 2024/02/08 0.31.2
 =================
 

crate/crash/command.py

@@ -26,12 +26,12 @@
 
 import logging
 import os
-import re
 import sys
 from argparse import ArgumentParser, ArgumentTypeError
 from collections import namedtuple
 from getpass import getpass
 from operator import itemgetter
+from typing import Union
 
 import sqlparse
 import urllib3
@@ -274,7 +274,7 @@ def _process_lines(self, lines):
 
     def _process_sql(self, text):
         sql = sqlparse.format(text, strip_comments=False)
-        for statement in sqlparse.split(sql):
+        for statement in sqlparse.parse(sql):
             self._exec_and_print(statement)
 
     def exit(self):
@@ -442,7 +442,7 @@ def _try_exec_cmd(self, line):
         return False
 
     def _exec(self, statement: str) -> bool:
-        """Execute the statement, prints errors if any occurr but no results."""
+        """Execute the statement, prints errors if any, but no results."""
         try:
             self.cursor.execute(statement)
             return True
@@ -457,9 +457,10 @@ def _exec(self, statement: str) -> bool:
                 self.logger.critical('\n' + e.error_trace)
         return False
 
-    def _exec_and_print(self, statement: str) -> bool:
+    def _exec_and_print(self, expression: Union[str, sqlparse.sql.Statement]) -> bool:
         """Execute the statement and print the output."""
-        success = self._exec(statement)
+        statement = to_statement(expression)
+        success = self._exec(str(statement).strip())
         self.exit_code = self.exit_code or int(not success)
         if not success:
             return False
@@ -482,9 +483,26 @@ def _exec_and_print(self, statement: str) -> bool:
         return True
 
 
-def stmt_type(statement):
+def stmt_type(expression: Union[str, sqlparse.sql.Statement]):
     """Extract type of statement, e.g. SELECT, INSERT, UPDATE, DELETE, ..."""
-    return re.findall(r'[\w]+', statement)[0].upper()
+    statement = to_statement(expression)
+    return str(statement.token_first(skip_ws=True, skip_cm=True)).upper()
+
+
+def to_statement(expression: Union[str, sqlparse.sql.Statement]) -> sqlparse.sql.Statement:
+    """
+    Convert SQL expression to sqlparse Statement.
+
+    This is mostly for backward-compatibility reasons, because the test cases
+    also submit *string types* to both `_exec_and_print` and `stmt_type`.
+    """
+    if isinstance(expression, sqlparse.sql.Statement):
+        statement = expression
+    elif isinstance(expression, str):
+        statement = sqlparse.parse(expression)[0]
+    else:
+        raise TypeError(f"Unknown type for expression: {type(expression)}")
+    return statement
 
 
 def get_lines_from_stdin():

crate/crash/output.txt

@@ -109,3 +109,26 @@ Status messages show the first word only::
     | 1 | 2 |
     +---+---+
     SELECT 1 row in set (... sec)
+
+
+    cr> /* foo */ select
+    ... 1,2
+    ... from sys.cluster limit 1;
+    +---+---+
+    | 1 | 2 |
+    +---+---+
+    | 1 | 2 |
+    +---+---+
+    SELECT 1 row in set (... sec)
+
+
+    cr> -- foo
+    ... select
+    ... 1,2
+    ... from sys.cluster limit 1;
+    +---+---+
+    | 1 | 2 |
+    +---+---+
+    | 1 | 2 |
+    +---+---+
+    SELECT 1 row in set (... sec)

tests/test_command.py

@@ -113,6 +113,11 @@ def test_stmt_type(self):
         # statements with trailing or leading spaces/tabs/linebreaks
         self.assertEqual(stmt_type(' SELECT 1 ;'), 'SELECT')
         self.assertEqual(stmt_type('\nSELECT\n1\n;\n'), 'SELECT')
+        # statements with trailing or leading comments
+        self.assertEqual(stmt_type('/* foo */ SELECT 1;'), 'SELECT')
+        self.assertEqual(stmt_type('SELECT 1; /* foo */'), 'SELECT')
+        self.assertEqual(stmt_type('-- foo \n SELECT 1;'), 'SELECT')
+        self.assertEqual(stmt_type('SELECT 1; -- foo'), 'SELECT')
 
     def test_decode_timeout_success(self):
         self.assertEqual(_decode_timeout(None), None)

tests/test_commands.py

@@ -24,7 +24,7 @@
 import sys
 import tempfile
 from unittest import SkipTest, TestCase
-from unittest.mock import MagicMock, call, patch
+from unittest.mock import MagicMock, Mock, call, patch
 
 from verlib2 import Version
 
@@ -38,6 +38,7 @@
     ToggleAutocompleteCommand,
     ToggleVerboseCommand,
 )
+from tests.util import fake_cursor
 
 
 class ReadFileCommandTest(TestCase):
@@ -221,6 +222,7 @@ def test_check_command_with_node_check(self, cmd):
         cmd.logger.info.assert_called_with('NODE CHECK OK')
 
 
+@patch('crate.client.connection.Cursor', fake_cursor())
 class CommentsTest(TestCase):
 
     def test_sql_comments(self):
@@ -236,9 +238,9 @@ def test_sql_comments(self):
 comment */ 4;
 """
         cmd = CrateShell()
-        cmd._exec_and_print = MagicMock()
+        cmd._exec = Mock()
         cmd.process_iterable(sql.splitlines())
-        self.assertListEqual(cmd._exec_and_print.mock_calls, [
+        self.assertListEqual(cmd._exec.mock_calls, [
             call("-- Just a dummy SELECT statement.\nSELECT 1;"),
             call("-- Another SELECT statement.\nSELECT 2;"),
             call('\n'.join([
@@ -267,29 +269,68 @@ def test_js_comments(self):
         """
 
         cmd = CrateShell()
-        cmd._exec_and_print = MagicMock()
+        cmd._exec = Mock()
         cmd.process(sql)
-        self.assertEqual(1, cmd._exec_and_print.call_count)
-        self.assertIn("CREATE FUNCTION", cmd._exec_and_print.mock_calls[0].args[0])
+        self.assertEqual(1, cmd._exec.call_count)
+        self.assertIn("CREATE FUNCTION", cmd._exec.mock_calls[0].args[0])
 
+    def test_exec_type_with_block_comment(self):
+        """
+        Make sure block SQL statements are executed as-is, without removing any comments.
 
+        However, also verify the status line displays the correct canonical
+        statement type, even when the SQL statement is prefixed using an
+        SQL comment.
+        """
+        sql = "/* foo */ select 1;"
+        cmd = CrateShell()
+        cmd._exec = Mock()
+        cmd.logger = Mock()
+        cmd.process(sql)
+        self.assertEqual(1, cmd._exec.call_count)
+        self.assertListEqual(cmd._exec.mock_calls, [
+            call("/* foo */ select 1;"),
+        ])
+        self.assertIn("SELECT 1 row in set", cmd.logger.mock_calls[0].args[0])
+
+    def test_exec_type_with_line_comment(self):
+        """
+        Make sure per-line SQL statements are executed as-is, without removing any comments.
+
+        However, also verify the status line displays the correct canonical
+        statement type, even when the SQL statement is prefixed using an
+        SQL comment.
+        """
+        sql = "-- foo \n select 1;"
+        cmd = CrateShell()
+        cmd._exec = Mock()
+        cmd.logger = Mock()
+        cmd.process(sql)
+        self.assertEqual(1, cmd._exec.call_count)
+        self.assertListEqual(cmd._exec.mock_calls, [
+            call("-- foo\nselect 1;"),
+        ])
+        self.assertIn("SELECT 1 row in set", cmd.logger.mock_calls[0].args[0])
+
+
+@patch('crate.client.connection.Cursor', fake_cursor())
 class MultipleStatementsTest(TestCase):
 
     def test_single_line_multiple_sql_statements(self):
         cmd = CrateShell()
-        cmd._exec_and_print = MagicMock()
+        cmd._exec = Mock()
         cmd.process("SELECT 1; SELECT 2; SELECT 3;")
-        self.assertListEqual(cmd._exec_and_print.mock_calls, [
+        self.assertListEqual(cmd._exec.mock_calls, [
             call("SELECT 1;"),
             call("SELECT 2;"),
             call("SELECT 3;"),
         ])
 
     def test_multiple_lines_multiple_sql_statements(self):
         cmd = CrateShell()
-        cmd._exec_and_print = MagicMock()
+        cmd._exec = Mock()
         cmd.process("SELECT 1;\nSELECT 2; SELECT 3;\nSELECT\n4;")
-        self.assertListEqual(cmd._exec_and_print.mock_calls, [
+        self.assertListEqual(cmd._exec.mock_calls, [
             call("SELECT 1;"),
             call("SELECT 2;"),
             call("SELECT 3;"),
@@ -300,9 +341,9 @@ def test_single_sql_statement_multiple_lines(self):
         """When processing single SQL statements, new lines are preserved."""
 
         cmd = CrateShell()
-        cmd._exec_and_print = MagicMock()
+        cmd._exec = Mock()
         cmd.process("\nSELECT\n1\nWHERE\n2\n=\n3\n;\n")
-        self.assertListEqual(cmd._exec_and_print.mock_calls, [
+        self.assertListEqual(cmd._exec.mock_calls, [
             call("SELECT\n1\nWHERE\n2\n=\n3\n;"),
         ])
 
@@ -321,10 +362,10 @@ def test_commands_and_multiple_sql_statements_interleaved(self):
         """Combine all test cases above to be sure everything integrates well."""
 
         cmd = CrateShell()
-        mock_manager = MagicMock()
+        mock_manager = Mock()
 
         cmd._try_exec_cmd = mock_manager.cmd
-        cmd._exec_and_print = mock_manager.sql
+        cmd._exec = mock_manager.sql
 
         cmd.process("""
     \\?
@@ -349,7 +390,7 @@ def test_comments_along_multiple_statements(self):
         """Test multiple types of comments along multi-statement input."""
 
         cmd = CrateShell()
-        cmd._exec_and_print = MagicMock()
+        cmd._exec = Mock()
 
         cmd.process("""
 -- Multiple statements and comments on same line
@@ -370,7 +411,7 @@ def test_comments_along_multiple_statements(self):
 comment */ 6;
         """)
 
-        self.assertListEqual(cmd._exec_and_print.mock_calls, [
+        self.assertListEqual(cmd._exec.mock_calls, [
             call('-- Multiple statements and comments on same line\n\nSELECT /* inner comment */ 1;'),
             call('/* this is a single-line comment */ SELECT /* inner comment */ 2;'),
 
@@ -379,4 +420,4 @@ def test_comments_along_multiple_statements(self):
 
             call('-- Multiple statements on multiple lines with multi-line comments between and inside them\n\nSELECT /* inner multi-line\ncomment */ 5 /* this is a multi-line\ncomment before statement end */;'),
             call('/* this is another multi-line\ncomment */ SELECT /* inner multi-line\ncomment */ 6;')
-        ])
+        ])

tests/util.py

@@ -0,0 +1,23 @@
+from unittest.mock import Mock
+
+from crate.client.cursor import Cursor
+
+
+def mocked_cursor(description, records, duration=0.1):
+    """
+    Provide a mocked `crate.client.cursor.Cursor` instance.
+    """
+    rowcount = len(records)
+    fake_cursor = Mock(name='fake_cursor', description=description, rowcount=rowcount, duration=duration)
+    fake_cursor.fetchall.return_value = records
+    FakeCursor = Mock(name='FakeCursor', spec=Cursor)
+    FakeCursor.return_value = fake_cursor
+    return FakeCursor
+
+
+def fake_cursor():
+    """
+    Provide an empty/minimal mocked cursor object,
+    that just works if you do not care about results.
+    """
+    return mocked_cursor(description=[('undef',)], records=[('undef', None)])