From d48f7c4bce918ac728ec514d1f56fb84d8ccb18e Mon Sep 17 00:00:00 2001
From: xxchan <xxchan22f@gmail.com>
Date: Sat, 16 Nov 2024 16:08:26 +0800
Subject: [PATCH] fix count

Signed-off-by: xxchan <xxchan22f@gmail.com>
---
 CHANGELOG.md               |  1 +
 sqllogictest/src/runner.rs | 24 +++++++++++++++++++++---
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 55d36e1..91f3d3f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - runner: Previously, `query` returning 0 rows will become `statement ok`. Now it returns `statement count 0`.
   - bin: Now `--override` will not change the type chars of `query <types>`, since in practice
     it becomes `?`s which might cause confusion.
+* runner: `statement count <n>` is incorrectly handled when the result is a `query`.
 
 ## [0.22.1] - 2024-11-11
 
diff --git a/sqllogictest/src/runner.rs b/sqllogictest/src/runner.rs
index 5635ee0..71db46c 100644
--- a/sqllogictest/src/runner.rs
+++ b/sqllogictest/src/runner.rs
@@ -849,7 +849,9 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
                 Record::Statement {
                     sql, expected, loc, ..
                 },
-                RecordOutput::Query { error: None, .. },
+                RecordOutput::Query {
+                    error: None, rows, ..
+                },
             ) => {
                 if let StatementExpect::Error(_) = expected {
                     return Err(TestErrorKind::Ok {
@@ -858,6 +860,16 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
                     }
                     .at(loc));
                 }
+                if let StatementExpect::Count(expected_count) = expected {
+                    if expected_count != rows.len() as u64 {
+                        return Err(TestErrorKind::StatementResultMismatch {
+                            sql,
+                            expected: expected_count,
+                            actual: format!("returned {} rows", rows.len()),
+                        }
+                        .at(loc));
+                    }
+                }
             }
             (
                 Record::Query {
@@ -1383,9 +1395,11 @@ pub fn update_record_with_output<T: ColumnType>(
                 loc,
                 conditions,
                 connection,
-                expected: expected @ (StatementExpect::Ok | StatementExpect::Count(_)),
+                expected: mut expected @ (StatementExpect::Ok | StatementExpect::Count(_)),
+            },
+            RecordOutput::Query {
+                error: None, rows, ..
             },
-            RecordOutput::Query { error: None, .. },
         ) => {
             // statement ok
             // SELECT ...
@@ -1394,6 +1408,10 @@ pub fn update_record_with_output<T: ColumnType>(
             // but don't care about the output.
             // DuckDB has a few of these.
 
+            if let StatementExpect::Count(expected_count) = &mut expected {
+                *expected_count = rows.len() as u64;
+            }
+
             Some(Record::Statement {
                 sql,
                 loc,