From 421b07e9fedb64a0eedb8d5cb3be044bd0ea104b Mon Sep 17 00:00:00 2001
From: James LaChance <fatboyxpc@gmail.com>
Date: Mon, 23 Sep 2024 02:03:03 -0400
Subject: [PATCH] Add collapsed output special command

It's sometimes annoying to have auto_vertical_output enabled in
myclirc and end up wanting to have the output not be expanded. The
current workaround for this is to modify the setting in the rc file
and restart mycli (or start a new session).
---
 mycli/AUTHORS                        | 1 +
 mycli/main.py                        | 4 ++++
 mycli/packages/special/iocommands.py | 9 +++++++++
 mycli/sqlexecute.py                  | 6 ++++++
 test/test_sqlexecute.py              | 5 +++++
 5 files changed, 25 insertions(+)

diff --git a/mycli/AUTHORS b/mycli/AUTHORS
index b8344520..7149be51 100644
--- a/mycli/AUTHORS
+++ b/mycli/AUTHORS
@@ -32,6 +32,7 @@ Contributors:
   * Daniel West
   * Daniël van Eeden
   * Fabrizio Gennari
+  * FatBoyXPC
   * François Pietka
   * Frederic Aoustin
   * Georgy Frolov
diff --git a/mycli/main.py b/mycli/main.py
index be15e343..ed608096 100755
--- a/mycli/main.py
+++ b/mycli/main.py
@@ -674,6 +674,7 @@ def one_iteration(text=None):
                     return
 
                 special.set_expanded_output(False)
+                special.set_forced_horizontal_output(False)
 
                 try:
                     text = self.handle_editor_command(text)
@@ -743,6 +744,9 @@ def one_iteration(text=None):
                     else:
                         max_width = None
 
+                    if special.forced_horizontal():
+                        max_width = None
+
                     formatted = self.format_output(title, cur, headers, special.is_expanded_output(), max_width)
 
                     t = time() - start
diff --git a/mycli/packages/special/iocommands.py b/mycli/packages/special/iocommands.py
index 87b53667..e3950c34 100644
--- a/mycli/packages/special/iocommands.py
+++ b/mycli/packages/special/iocommands.py
@@ -20,6 +20,7 @@
 
 TIMING_ENABLED = False
 use_expanded_output = False
+force_horizontal_output = False
 PAGER_ENABLED = True
 tee_file = None
 once_file = None
@@ -97,6 +98,14 @@ def set_expanded_output(val):
 def is_expanded_output():
     return use_expanded_output
 
+@export
+def set_forced_horizontal_output(val):
+    global force_horizontal_output
+    force_horizontal_output = val
+
+@export
+def forced_horizontal():
+    return force_horizontal_output
 
 _logger = logging.getLogger(__name__)
 
diff --git a/mycli/sqlexecute.py b/mycli/sqlexecute.py
index d5b6db6f..89d4ba6b 100644
--- a/mycli/sqlexecute.py
+++ b/mycli/sqlexecute.py
@@ -298,6 +298,12 @@ def run(self, statement):
             if sql.endswith("\\G"):
                 special.set_expanded_output(True)
                 sql = sql[:-2].strip()
+            # \g is treated specially since we might want collapsed output when
+            # auto vertical output is enabled
+            elif sql.endswith('\\g'):
+                special.set_expanded_output(False)
+                special.set_forced_horizontal_output(True)
+                sql = sql[:-2].strip()
 
             cur = self.conn.cursor()
             try:  # Special command
diff --git a/test/test_sqlexecute.py b/test/test_sqlexecute.py
index 17e082b5..37587cbb 100644
--- a/test/test_sqlexecute.py
+++ b/test/test_sqlexecute.py
@@ -172,6 +172,11 @@ def test_favorite_query_expanded_output(executor):
     results = run(executor, "\\fd test-ae")
     assert_result_equal(results, status="test-ae: Deleted")
 
+@dbtest
+def test_collapsed_output_special_command(executor):
+    set_expanded_output(True)
+    results = run(executor, 'select 1\\g')
+    assert is_expanded_output() is False
 
 @dbtest
 def test_special_command(executor):