From 57274fa40b448953df48ee27ab593e04fa458efb Mon Sep 17 00:00:00 2001
From: CK-7vn <keighan.robichaudgriffin@maine.edu>
Date: Fri, 17 Jan 2025 13:59:34 -0500
Subject: [PATCH] Correct CLI comment handling to mimic sqlite behavior

---
 cli/app.rs            | 39 ++++++++++++++++++++--
 testing/shelltests.py | 78 ++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 106 insertions(+), 11 deletions(-)

diff --git a/cli/app.rs b/cli/app.rs
index 25eb43592..0d9645c3e 100644
--- a/cli/app.rs
+++ b/cli/app.rs
@@ -434,9 +434,6 @@ impl Limbo {
         line: &str,
         rl: &mut rustyline::DefaultEditor,
     ) -> anyhow::Result<()> {
-        if line.trim_start().starts_with("--") {
-            return Ok(());
-        }
         if self.input_buff.is_empty() {
             if line.is_empty() {
                 return Ok(());
@@ -448,6 +445,42 @@ impl Limbo {
                 return Ok(());
             }
         }
+        if line.trim_start().starts_with("--") {
+            if let Some(remaining) = line.split_once('\n') {
+                let after_comment = remaining.1.trim();
+                if !after_comment.is_empty() {
+                    rl.add_history_entry(after_comment.to_owned())?;
+                    self.buffer_input(after_comment);
+
+                    if after_comment.ends_with(';') {
+                        if self.opts.echo {
+                            let _ = self.writeln(after_comment);
+                        }
+                        let conn = self.conn.clone();
+                        let runner = conn.query_runner(after_comment.as_bytes());
+                        for output in runner {
+                            if let Err(e) = self.print_query_result(after_comment, output) {
+                                let _ = self.writeln(e.to_string());
+                            }
+                        }
+                        self.reset_input();
+                    } else {
+                        self.set_multiline_prompt();
+                    }
+                    self.interrupt_count.store(0, Ordering::SeqCst);
+                    return Ok(());
+                }
+            }
+            return Ok(());
+        }
+
+        if let Some(comment_pos) = line.find("--") {
+            let before_comment = line[..comment_pos].trim();
+            if !before_comment.is_empty() {
+                return self.handle_input_line(before_comment, rl);
+            }
+        }
+
         if line.ends_with(';') {
             self.buffer_input(line);
             let buff = self.input_buff.clone();
diff --git a/testing/shelltests.py b/testing/shelltests.py
index bfe0b50f7..81b0cd78f 100755
--- a/testing/shelltests.py
+++ b/testing/shelltests.py
@@ -52,7 +52,9 @@ def execute_sql(pipe, sql):
 
     output = ""
     while True:
-        ready_to_read, _, error_in_pipe = select.select([stdout, stderr], [], [stdout, stderr])
+        ready_to_read, _, error_in_pipe = select.select(
+            [stdout, stderr], [], [stdout, stderr]
+        )
         ready_to_read_or_err = set(ready_to_read + error_in_pipe)
         if stderr in ready_to_read_or_err:
             exit_on_error(stderr)
@@ -67,6 +69,7 @@ def execute_sql(pipe, sql):
     output = strip_each_line(output)
     return output
 
+
 def strip_each_line(lines: str) -> str:
     lines = lines.split("\n")
     lines = [line.strip() for line in lines if line != ""]
@@ -237,8 +240,10 @@ csv_file = "./test_files/test.csv"
 write_to_pipe(".open :memory:")
 
 
-def test_import_csv(test_name: str, options: str, import_output: str, table_output: str):
-    csv_table_name = f'csv_table_{test_name}'
+def test_import_csv(
+    test_name: str, options: str, import_output: str, table_output: str
+):
+    csv_table_name = f"csv_table_{test_name}"
     write_to_pipe(f"CREATE TABLE {csv_table_name} (c1 INT, c2 REAL, c3 String);")
     do_execshell_test(
         pipe,
@@ -253,11 +258,15 @@ def test_import_csv(test_name: str, options: str, import_output: str, table_outp
         table_output,
     )
 
-test_import_csv('no_options', '--csv', '', '1|2.0|String\'1\n3|4.0|String2')
-test_import_csv('verbose', '--csv -v', 
-                'Added 2 rows with 0 errors using 2 lines of input' 
-                ,'1|2.0|String\'1\n3|4.0|String2')
-test_import_csv('skip', '--csv --skip 1', '' ,'3|4.0|String2')
+
+test_import_csv("no_options", "--csv", "", "1|2.0|String'1\n3|4.0|String2")
+test_import_csv(
+    "verbose",
+    "--csv -v",
+    "Added 2 rows with 0 errors using 2 lines of input",
+    "1|2.0|String'1\n3|4.0|String2",
+)
+test_import_csv("skip", "--csv --skip 1", "", "3|4.0|String2")
 
 
 # Verify the output file exists and contains expected content
@@ -297,6 +306,59 @@ else:
     print(f"File contents:\n{file_contents}")
     exit(1)
 
+do_execshell_test(
+    pipe,
+    "test-single-line-comment",
+    "-- this is a comment\nSELECT 1;",
+    "1",
+)
+
+do_execshell_test(
+    pipe,
+    "test-multi-line-single-line-comments-in-succession",
+    """-- First of the comments
+    -- Second line of the comments
+    SELECT 2;""",
+    "2",
+)
+do_execshell_test(
+    pipe,
+    "test-multi-line-comments",
+    """/*
+    This is a multi-line comment
+    */
+    SELECT 3;""",
+    "3",
+)
+
+# readd some data to test inline comments
+write_to_pipe("""
+CREATE TABLE users (id INTEGER PRIMARY KEY, first_name TEXT, last_name TEXT, age INTEGER);
+INSERT INTO users (id, first_name, last_name, age) VALUES
+(1, 'Alice', 'Smith', 30), (2, 'Bob', 'Johnson', 25), (3, 'Charlie', 'Brown', 66), (4, 'David', 'Nichols', 70);
+""")
+
+do_execshell_test(
+    pipe,
+    "test-inline-comments",
+    """SELECT id, -- this is a comment until newline 
+    first_name
+    FROM users
+    LIMIT 1; """,
+    "1|Alice",
+)
+
+do_execshell_test(
+    pipe,
+    "test-multiple-inline-comments",
+    """SELECT id, --first inline 
+    --second inline 
+    first_name
+    --third inline 
+    FROM users
+    LIMIT 1; """,
+    "1|Alice",
+)
 # Cleanup
 os.remove(filepath)
 pipe.terminate()