Sql fixes

crates/storage-query-datafusion/src/deployment/table.rs

@@ -76,7 +76,7 @@ async fn for_each_state(
         append_deployment_row(&mut builder, &mut temp, deployment);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -87,6 +87,6 @@ async fn for_each_state(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/idempotency/table.rs

@@ -80,7 +80,7 @@ async fn for_each_state(
         );
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -91,6 +91,6 @@ async fn for_each_state(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/inbox/table.rs

@@ -73,7 +73,7 @@ async fn for_each_state(
         append_inbox_row(&mut builder, &mut temp, row);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -84,6 +84,6 @@ async fn for_each_state(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/invocation_state/table.rs

@@ -79,7 +79,7 @@ async fn for_each_state<'a, I>(
         append_state_row(&mut builder, &mut temp, row);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -90,6 +90,6 @@ async fn for_each_state<'a, I>(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/invocation_status/table.rs

@@ -70,7 +70,7 @@ async fn for_each_status<'a, I>(
         append_invocation_status_row(&mut builder, &mut temp, row);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -81,6 +81,6 @@ async fn for_each_status<'a, I>(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/journal/table.rs

@@ -70,7 +70,7 @@ async fn for_each_journal<'a, I>(
         append_journal_row(&mut builder, &mut temp, row);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -81,6 +81,6 @@ async fn for_each_journal<'a, I>(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/journal/tests.rs

@@ -11,7 +11,7 @@
 use crate::mocks::*;
 use crate::row;
 use bytes::Bytes;
-use datafusion::arrow::array::{LargeStringArray, UInt32Array};
+use datafusion::arrow::array::{Int64Array, LargeStringArray, UInt32Array};
 use datafusion::arrow::record_batch::RecordBatch;
 use futures::StreamExt;
 use googletest::all;
@@ -130,3 +130,78 @@ async fn get_entries() {
         )
     );
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn select_count_star() {
+    let tc = TaskCenterBuilder::default()
+        .default_runtime_handle(tokio::runtime::Handle::current())
+        .build()
+        .expect("task_center builds");
+    let mut engine = tc
+        .run_in_scope("mock-query-engine", None, MockQueryEngine::create())
+        .await;
+
+    let mut tx = engine.partition_store().transaction();
+    let journal_invocation_id = InvocationId::mock_random();
+    tx.put_journal_entry(
+        &journal_invocation_id,
+        0,
+        JournalEntry::Entry(ProtobufRawEntryCodec::serialize_enriched(Entry::Input(
+            InputEntry {
+                value: Default::default(),
+            },
+        ))),
+    )
+    .await;
+    tx.put_journal_entry(
+        &journal_invocation_id,
+        1,
+        JournalEntry::Entry(EnrichedRawEntry::new(
+            EnrichedEntryHeader::Call {
+                is_completed: false,
+                enrichment_result: Some(CallEnrichmentResult {
+                    invocation_id: InvocationId::mock_random(),
+                    invocation_target: InvocationTarget::mock_virtual_object(),
+                    completion_retention_time: None,
+                    span_context: Default::default(),
+                }),
+            },
+            Bytes::new(),
+        )),
+    )
+    .await;
+    tx.put_journal_entry(
+        &journal_invocation_id,
+        2,
+        JournalEntry::Entry(EnrichedRawEntry::new(
+            EnrichedEntryHeader::Run {},
+            service_protocol::RunEntryMessage {
+                name: "my-side-effect".to_string(),
+                result: None,
+            }
+            .encode_to_vec()
+            .into(),
+        )),
+    )
+    .await;
+    tx.commit().await.unwrap();
+
+    let records = engine
+        .execute("SELECT COUNT(*) AS count FROM sys_journal")
+        .await
+        .unwrap()
+        .collect::<Vec<Result<RecordBatch, _>>>()
+        .await
+        .remove(0)
+        .unwrap();
+
+    assert_that!(
+        records,
+        all!(row!(
+            0,
+            {
+                "count" => Int64Array: eq(3)
+            }
+        ),)
+    );
+}

crates/storage-query-datafusion/src/keyed_service_status/table.rs

@@ -70,7 +70,7 @@ async fn for_each_status<'a, I>(
         append_virtual_object_status_row(&mut builder, &mut temp, row);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -81,6 +81,6 @@ async fn for_each_status<'a, I>(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/promise/table.rs

@@ -70,7 +70,7 @@ async fn for_each_state(
         append_promise_row(&mut builder, &mut temp, owned_promise_row);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -81,6 +81,6 @@ async fn for_each_state(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/service/table.rs

@@ -75,7 +75,7 @@ async fn for_each_state(
         append_service_row(&mut builder, &mut temp, service);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -86,6 +86,6 @@ async fn for_each_state(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/state/table.rs

@@ -69,7 +69,7 @@ async fn for_each_state<'a, I>(
         append_state_row(&mut builder, row);
         if builder.full() {
             let batch = builder.finish();
-            if tx.send(Ok(batch)).await.is_err() {
+            if tx.send(batch).await.is_err() {
                 // not sure what to do here?
                 // the other side has hung up on us.
                 // we probably don't want to panic, is it will cause the entire process to exit
@@ -80,6 +80,6 @@ async fn for_each_state<'a, I>(
     }
     if !builder.empty() {
         let result = builder.finish();
-        let _ = tx.send(Ok(result)).await;
+        let _ = tx.send(result).await;
     }
 }

crates/storage-query-datafusion/src/table_macro.rs

@@ -428,9 +428,12 @@ macro_rules! define_table {
                 self.rows_inserted_so_far == 0
             }
 
-            pub fn finish(self) -> ::datafusion::arrow::record_batch::RecordBatch {
+            pub fn finish(self) -> ::datafusion::common::Result<::datafusion::arrow::record_batch::RecordBatch> {
                 let arrays = self.arrays.finish();
-                ::datafusion::arrow::record_batch::RecordBatch::try_new(self.projected_schema, arrays).unwrap()
+                // We add the row count as it wouldn't otherwise work with queries that
+                // just run aggregate functions (e.g. COUNT(*)) without selecting fields.
+                let options = ::datafusion::arrow::record_batch::RecordBatchOptions::new().with_row_count(Some(self.rows_inserted_so_far));
+                Ok(::datafusion::arrow::record_batch::RecordBatch::try_new_with_options(self.projected_schema, arrays, &options)?)
             }
 
         }