fix count statement when displaylimit=None (#815)

CHANGELOG.md

@@ -15,6 +15,7 @@
 * [Fix] `--persist/--persist-replace` perform `ROLLBACK` automatically when needed
 * [Fix] `ResultSet` footer (when `displaylimit` truncates results and when showing how to convert to a data frame) now appears in the `ResultSet` plain text representation (#682)
 * [Fix] Improve error when calling `%sqlcmd` (#761)
+* [Fix] Fix count statement's result not displayed when `displaylimit=None` (#801)
 * [API Change] When loading connections from a `.ini` file via `%sql --section section_name`, the section name is set as the connection alias
 * [API Change] Starting connections from a `.ini` file via `%sql [section_name]` has been deprecated
 * [Doc] Fixes documentation inaccuracy that said `:variable` was deprecated (we brought it back in `0.9.0`)

src/sql/run/resultset.py

@@ -97,7 +97,9 @@ def _extend_results(self, elements):
         """Store the DB fetched results into the internal list of results"""
         to_add = self._config.displaylimit - len(self._results)
         self._results.extend(elements)
-        self._pretty_table.add_rows(elements[:to_add])
+        self._pretty_table.add_rows(
+            elements if self._config.displaylimit == 0 else elements[:to_add]
+        )
 
     def mark_fetching_as_done(self):
         self._mark_fetching_as_done = True

src/tests/test_magic.py

@@ -677,6 +677,28 @@ def test_displaylimit_with_conditional_clause(
         assert f"Truncated to {DISPLAYLIMIT_LINK} of 10" in out._repr_html_()
 
 
+@pytest.mark.parametrize(
+    "config_value",
+    [
+        (1),
+        (0),
+        (None),
+    ],
+)
+def test_displaylimit_with_count_statement(ip, load_penguin, config_value):
+    ip.run_cell(f"%config SqlMagic.displaylimit = {config_value}")
+    result = ip.run_line_magic("sql", "select count(*) from penguins.csv")
+
+    assert isinstance(result, ResultSet)
+    assert str(result) == (
+        "+--------------+\n"
+        "| count_star() |\n"
+        "+--------------+\n"
+        "|     344      |\n"
+        "+--------------+"
+    )
+
+
 def test_column_local_vars(ip):
     ip.run_line_magic("config", "SqlMagic.column_local_vars = True")
     result = runsql(ip, "SELECT * FROM author;")