From c97048cfc8c0fed01d7b32d3173be508135b9769 Mon Sep 17 00:00:00 2001
From: Joel Natividad <1980690+jqnatividad@users.noreply.github.com>
Date: Wed, 3 Jan 2024 07:41:19 -0500
Subject: [PATCH] `sqlp`: automatically use tab delimiters when output format
 is CSV and file extension is TSV or TAB (like the rest of qsv)

---
 src/cmd/sqlp.rs    | 17 +++++++++++++++--
 tests/test_sqlp.rs | 31 +++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/cmd/sqlp.rs b/src/cmd/sqlp.rs
index d3be4e94e..45a9748b8 100644
--- a/src/cmd/sqlp.rs
+++ b/src/cmd/sqlp.rs
@@ -240,7 +240,7 @@ struct Args {
     flag_quiet:            bool,
 }
 
-#[derive(Default, Clone)]
+#[derive(Default, Clone, PartialEq)]
 enum OutputMode {
     #[default]
     Csv,
@@ -258,7 +258,7 @@ impl OutputMode {
         &self,
         query: &str,
         ctx: &mut SQLContext,
-        delim: u8,
+        mut delim: u8,
         args: Args,
     ) -> CliResult<(usize, usize)> {
         let mut df = DataFrame::default();
@@ -275,6 +275,19 @@ impl OutputMode {
             let w = match args.flag_output {
                 Some(x) => {
                     let path = Path::new(&x);
+                    // if the output file ends with ".tsv" or ".tab", we use tab as the delimiter
+                    delim = if *self == OutputMode::Csv
+                        && path
+                            .extension()
+                            .map_or(false, |ext| ext.eq_ignore_ascii_case("tsv"))
+                        || path
+                            .extension()
+                            .map_or(false, |ext| ext.eq_ignore_ascii_case("tab"))
+                    {
+                        b'\t'
+                    } else {
+                        delim
+                    };
                     Box::new(File::create(path)?) as Box<dyn Write>
                 },
                 None => Box::new(io::stdout()) as Box<dyn Write>,
diff --git a/tests/test_sqlp.rs b/tests/test_sqlp.rs
index 2c22a0e54..8abb0bb5c 100644
--- a/tests/test_sqlp.rs
+++ b/tests/test_sqlp.rs
@@ -1149,3 +1149,34 @@ fn sqlp_sql_from_subquery() {
 
     assert_eq!(got, expected);
 }
+
+#[test]
+fn sqlp_sql_tsv() {
+    let wrk = Workdir::new("sqlp_sql_tsv");
+    wrk.create(
+        "test.csv",
+        vec![
+            svec!["idx", "val"],
+            svec!["0", "ABC"],
+            svec!["1", "abc"],
+            svec!["2", "000"],
+            svec!["3", "A0C"],
+            svec!["4", "a0c"],
+        ],
+    );
+
+    let output_file = wrk.path("output.tsv").to_string_lossy().to_string();
+
+    let mut cmd = wrk.command("sqlp");
+    cmd.arg("test.csv")
+        .arg("SELECT * FROM test")
+        .args(["--output", &output_file]);
+
+    wrk.assert_success(&mut cmd);
+
+    let got = wrk.read_to_string(&output_file);
+
+    let expected = "idx\tval\n0\tABC\n1\tabc\n2\t000\n3\tA0C\n4\ta0c\n";
+
+    assert_eq!(got, expected);
+}