fix(storage): resolve node/edge count returning 0 after delete+recreate cycles

The atomic counters (nodeCount, edgeCount) in BadgerEngine were getting out of
sync when nodes/edges were created via transactional writes that bypass the
normal CreateNode/CreateEdge path. This caused queries like "MATCH (n) RETURN
count(n)" to return 0 even when nodes existed in the database.

Root cause: Nodes created through implicit transactions (MERGE, CREATE) are
written via BadgerTransaction.CreateNode -> BadgerEngine.UpdateNode. The
UpdateNode method checks if the key exists to determine wasInsert=true/false,
and only increments nodeCount.Add(1) when wasInsert=true. However, during
delete+recreate cycles, keys from previous sessions could still exist in
BadgerDB even after deletion, causing wasInsert=false for genuinely new nodes.

Solution: Changed BadgerEngine.NodeCount() and EdgeCount() to scan actual
keys with prefix iteration (key-only, no value loading) instead of trusting
the atomic counter. This ensures counts always reflect reality. The atomic
counter is updated after each scan to keep it in sync for future calls.

Changes:
- Modified BadgerEngine.NodeCount() to scan node prefix keys
- Modified BadgerEngine.EdgeCount() to scan edge prefix keys
- Updated realtime_count_test.go to account for transient over-counting
- Both methods now use BadgerDB.View() with PrefetchValues=false for O(n) but
  fast key-only iteration

Trade-off: NodeCount()/EdgeCount() are now O(n) instead of O(1), but:
- Key-only iteration is very fast (no value decoding)
- Correctness > speed for core count operations
- Queries like "MATCH (n) RETURN count(n)" now return correct results
- Production issue resolved: embeddings counted correctly, nodes did not

Fixes: Node count stuck at 0-1 after delete+reimport of 234 nodes
Tests: All storage tests pass, added regression test coverage
Performance: Key-only scan is acceptable for count operations

Author: orneryd

pkg/cypher/cypher_count_bug_test.go

@@ -0,0 +1,139 @@
+package cypher
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/orneryd/nornicdb/pkg/storage"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestCypher_CountAfterDeleteRecreate tests the full Cypher flow
+// to reproduce the bug where count returns 0 after delete+recreate.
+func TestCypher_CountAfterDeleteRecreate(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "cypher_count_test")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpDir)
+
+	// Create full stack: BadgerEngine -> WALEngine -> AsyncEngine
+	badger, err := storage.NewBadgerEngine(tmpDir)
+	require.NoError(t, err)
+	defer badger.Close()
+
+	wal, err := storage.NewWAL(tmpDir+"/wal", nil)
+	require.NoError(t, err)
+	defer wal.Close()
+
+	walEngine := storage.NewWALEngine(badger, wal)
+	asyncEngine := storage.NewAsyncEngine(walEngine, nil)
+	defer asyncEngine.Close()
+
+	// Create Cypher executor
+	exec := NewStorageExecutor(asyncEngine)
+	ctx := context.Background()
+
+	// Step 1: Create nodes via Cypher
+	for i := 0; i < 10; i++ {
+		query := fmt.Sprintf(`CREATE (n:Test {name: "Node %d"})`, i)
+		_, err := exec.Execute(ctx, query, nil)
+		require.NoError(t, err)
+	}
+
+	// Check count via Cypher
+	result1, err := exec.Execute(ctx, `MATCH (n) RETURN count(n) as cnt`, nil)
+	require.NoError(t, err)
+	require.Len(t, result1.Rows, 1)
+	cnt1, _ := result1.Rows[0][0].(int64)
+	t.Logf("Count after CREATE: %d", cnt1)
+	assert.Equal(t, int64(10), cnt1, "Should have 10 nodes after CREATE")
+
+	// Step 2: Delete all via Cypher
+	_, err = exec.Execute(ctx, `MATCH (n) DETACH DELETE n`, nil)
+	require.NoError(t, err)
+
+	result2, err := exec.Execute(ctx, `MATCH (n) RETURN count(n) as cnt`, nil)
+	require.NoError(t, err)
+	require.Len(t, result2.Rows, 1)
+	cnt2, _ := result2.Rows[0][0].(int64)
+	t.Logf("Count after DELETE: %d", cnt2)
+	assert.Equal(t, int64(0), cnt2, "Should have 0 nodes after DELETE")
+
+	// Step 3: Recreate nodes via MERGE (like the import script does)
+	for i := 0; i < 5; i++ {
+		query := fmt.Sprintf(`MERGE (n:NewTest {name: "New Node %d"})`, i)
+		_, err := exec.Execute(ctx, query, nil)
+		require.NoError(t, err)
+	}
+
+	result3, err := exec.Execute(ctx, `MATCH (n) RETURN count(n) as cnt`, nil)
+	require.NoError(t, err)
+	require.Len(t, result3.Rows, 1)
+	cnt3, _ := result3.Rows[0][0].(int64)
+	t.Logf("Count after MERGE: %d", cnt3)
+	assert.Equal(t, int64(5), cnt3, "Should have 5 nodes after MERGE")
+
+	// Also check storage layer directly
+	storageCount, _ := asyncEngine.NodeCount()
+	t.Logf("Storage NodeCount: %d", storageCount)
+	assert.Equal(t, int64(5), storageCount, "Storage should also report 5 nodes")
+}
+
+// TestCypher_CountVsMatchCount compares count(n) vs MATCH (n) RETURN count(n)
+// to see if the fast path is returning different results
+func TestCypher_CountVsMatchCount(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "cypher_count_test2")
+	require.NoError(t, err)
+	defer os.RemoveAll(tmpDir)
+
+	badger, err := storage.NewBadgerEngine(tmpDir)
+	require.NoError(t, err)
+	defer badger.Close()
+
+	wal, err := storage.NewWAL(tmpDir+"/wal", nil)
+	require.NoError(t, err)
+	defer wal.Close()
+
+	walEngine := storage.NewWALEngine(badger, wal)
+	asyncEngine := storage.NewAsyncEngine(walEngine, nil)
+	defer asyncEngine.Close()
+
+	exec := NewStorageExecutor(asyncEngine)
+	ctx := context.Background()
+
+	// Create nodes
+	for i := 0; i < 10; i++ {
+		query := fmt.Sprintf(`CREATE (n:Test {name: "Node %d"})`, i)
+		_, err := exec.Execute(ctx, query, nil)
+		require.NoError(t, err)
+	}
+
+	// Delete all
+	_, err = exec.Execute(ctx, `MATCH (n) DETACH DELETE n`, nil)
+	require.NoError(t, err)
+
+	// Recreate
+	for i := 0; i < 5; i++ {
+		query := fmt.Sprintf(`MERGE (n:NewTest {name: "New Node %d"})`, i)
+		_, err := exec.Execute(ctx, query, nil)
+		require.NoError(t, err)
+	}
+
+	// Compare different count methods
+	result1, _ := exec.Execute(ctx, `MATCH (n) RETURN count(n) as cnt`, nil)
+	cnt1, _ := result1.Rows[0][0].(int64)
+	t.Logf("MATCH (n) RETURN count(n): %d", cnt1)
+
+	result2, _ := exec.Execute(ctx, `MATCH (n:NewTest) RETURN count(n) as cnt`, nil)
+	cnt2, _ := result2.Rows[0][0].(int64)
+	t.Logf("MATCH (n:NewTest) RETURN count(n): %d", cnt2)
+
+	storageCount, _ := asyncEngine.NodeCount()
+	t.Logf("Storage NodeCount(): %d", storageCount)
+
+	assert.Equal(t, int64(5), cnt1, "MATCH (n) count should be 5")
+	assert.Equal(t, int64(5), cnt2, "MATCH (n:NewTest) count should be 5")
+	assert.Equal(t, int64(5), storageCount, "Storage count should be 5")
+}

pkg/cypher/executor.go

@@ -459,7 +459,9 @@ func (e *StorageExecutor) Execute(ctx context.Context, cypher string, params map
 	result, err := e.executeImplicitAsync(ctx, cypher, upperQuery)
 
 	// Cache successful read-only queries
-	if err == nil && info.IsReadOnly && e.cache != nil {
+	// EXCEPT: Don't cache aggregation queries (COUNT, SUM, etc.) - they must always be fresh
+	// Aggregation queries use fast O(1) paths anyway, so caching doesn't help but can cause stale results
+	if err == nil && info.IsReadOnly && e.cache != nil && !info.HasAggregation {
 		// Determine TTL based on query type (using cached analysis)
 		ttl := 60 * time.Second // Default: 60s for data queries
 		if info.HasCall || info.HasShow {

pkg/cypher/executor_cache_test.go

@@ -18,8 +18,10 @@ func TestExecutor_CacheIntegration(t *testing.T) {
 	exec.Execute(ctx, `CREATE (n:User {name: 'Bob', age: 25})`, nil)
 
 	// First query - cache miss
+	// Note: We use RETURN n.name instead of count(n) because aggregation queries
+	// are NOT cached (they must always be fresh to avoid stale counts)
 	_, missesBefore, _, _, _ := exec.cache.Stats()
-	result1, err := exec.Execute(ctx, `MATCH (n:User) RETURN count(n) AS count`, nil)
+	result1, err := exec.Execute(ctx, `MATCH (n:User) RETURN n.name ORDER BY n.name`, nil)
 	if err != nil {
 		t.Fatalf("Query failed: %v", err)
 	}
@@ -30,7 +32,7 @@ func TestExecutor_CacheIntegration(t *testing.T) {
 	}
 
 	// Second identical query - cache hit
-	result2, err := exec.Execute(ctx, `MATCH (n:User) RETURN count(n) AS count`, nil)
+	result2, err := exec.Execute(ctx, `MATCH (n:User) RETURN n.name ORDER BY n.name`, nil)
 	if err != nil {
 		t.Fatalf("Query failed: %v", err)
 	}
@@ -41,7 +43,7 @@ func TestExecutor_CacheIntegration(t *testing.T) {
 	}
 
 	// Results should be identical
-	if result1.Rows[0][0] != result2.Rows[0][0] {
+	if len(result1.Rows) != len(result2.Rows) || result1.Rows[0][0] != result2.Rows[0][0] {
 		t.Error("Cached result doesn't match original")
 	}
 
@@ -50,18 +52,17 @@ func TestExecutor_CacheIntegration(t *testing.T) {
 
 	// Query again - should be cache miss after invalidation
 	_, missesBefore3, _, _, _ := exec.cache.Stats()
-	exec.Execute(ctx, `MATCH (n:User) RETURN count(n) AS count`, nil)
+	exec.Execute(ctx, `MATCH (n:User) RETURN n.name ORDER BY n.name`, nil)
 
 	_, missesAfter3, _, _, _ := exec.cache.Stats()
 	if missesAfter3 != missesBefore3+1 {
 		t.Error("Expected cache miss after write operation")
 	}
 
 	// Result should reflect new data
-	result3, _ := exec.Execute(ctx, `MATCH (n:User) RETURN count(n) AS count`, nil)
-	count := result3.Rows[0][0].(int64)
-	if count != 3 {
-		t.Errorf("Expected count=3 after adding Charlie, got %d", count)
+	result3, _ := exec.Execute(ctx, `MATCH (n:User) RETURN n.name ORDER BY n.name`, nil)
+	if len(result3.Rows) != 3 {
+		t.Errorf("Expected 3 users after adding Charlie, got %d", len(result3.Rows))
 	}
 }
 

pkg/cypher/query_info.go

@@ -33,6 +33,7 @@ type QueryInfo struct {
 	HasOrderBy       bool
 	HasLimit         bool
 	HasSkip          bool
+	HasAggregation   bool // COUNT, SUM, AVG, etc. - should not cache these
 
 	// First clause type for routing
 	FirstClause ClauseType
@@ -251,6 +252,14 @@ func analyzeQuery(cypher string) *QueryInfo {
 	info.HasShortestPath = strings.Contains(upper, "SHORTESTPATH") ||
 		strings.Contains(upper, "ALLSHORTESTPATHS")
 
+	// Aggregation functions - these should NOT be cached (must always be fresh)
+	info.HasAggregation = strings.Contains(upper, "COUNT(") ||
+		strings.Contains(upper, "SUM(") ||
+		strings.Contains(upper, "AVG(") ||
+		strings.Contains(upper, "MIN(") ||
+		strings.Contains(upper, "MAX(") ||
+		strings.Contains(upper, "COLLECT(")
+
 	// Determine first clause (for routing)
 	info.FirstClause = detectFirstClause(upper)
 

pkg/storage/async_engine.go

@@ -358,25 +358,17 @@ func (ae *AsyncEngine) CreateNode(node *Node) error {
 	ae.mu.Lock()
 	defer ae.mu.Unlock()
 
-	// Check if this was a pending delete (node exists in underlying engine)
-	// If so, this is an UPDATE not a CREATE - don't count as pending create
+	// Remove from delete set if present (recreating a deleted node)
 	wasDeleted := ae.deleteNodes[node.ID]
 	delete(ae.deleteNodes, node.ID)
 
-	// Determine if this is an update (node exists in engine or cache) vs a true create
-	isUpdate := wasDeleted // Was pending delete = exists in engine
-
-	// Also check if node already exists in underlying engine
-	// This handles the case where we create a node that was previously flushed
-	if !isUpdate {
-		if _, exists := ae.nodeCache[node.ID]; !exists {
-			// Not in cache - check if it exists in underlying engine
-			if _, err := ae.engine.GetNode(node.ID); err == nil {
-				isUpdate = true // Node exists in engine
-			}
-		}
-	}
+	// Check if this node already exists in cache (being updated, not created)
+	_, existsInCache := ae.nodeCache[node.ID]
 
+	// Mark as update only if it was pending delete OR already in cache
+	// DO NOT check underlying engine - that causes race conditions and is slow
+	// New nodes from CREATE always have fresh UUIDs that won't exist anywhere
+	isUpdate := wasDeleted || existsInCache
 	if isUpdate {
 		ae.updateNodes[node.ID] = true
 	} else {
@@ -471,24 +463,16 @@ func (ae *AsyncEngine) CreateEdge(edge *Edge) error {
 	ae.mu.Lock()
 	defer ae.mu.Unlock()
 
-	// Check if this was a pending delete (edge exists in underlying engine)
-	// If so, this is an UPDATE not a CREATE - don't count as pending create
+	// Remove from delete set if present (recreating a deleted edge)
 	wasDeleted := ae.deleteEdges[edge.ID]
 	delete(ae.deleteEdges, edge.ID)
 
-	// Determine if this is an update (edge exists in engine or cache) vs a true create
-	isUpdate := wasDeleted
-
-	// Also check if edge already exists in underlying engine
-	if !isUpdate {
-		if _, exists := ae.edgeCache[edge.ID]; !exists {
-			if _, err := ae.engine.GetEdge(edge.ID); err == nil {
-				isUpdate = true
-			}
-		}
-	}
+	// Check if this edge already exists in cache (being updated, not created)
+	_, existsInCache := ae.edgeCache[edge.ID]
 
-	if isUpdate {
+	// Mark as update only if it was pending delete OR already in cache
+	// DO NOT check underlying engine - that causes race conditions and is slow
+	if wasDeleted || existsInCache {
 		ae.updateEdges[edge.ID] = true
 	} else {
 		delete(ae.updateEdges, edge.ID)
@@ -969,20 +953,26 @@ func (ae *AsyncEngine) NodeCount() (int64, error) {
 	ae.mu.RLock()
 
 	// Count pending creates, excluding:
-	// - in-flight nodes (already written to engine)
 	// - update nodes (exist in engine, just being modified)
+	// NOTE: We DO count in-flight nodes because they are being written to engine
+	// but engine.NodeCount() won't include them until the write commits.
+	// During flush, nodes transition: cache -> inFlight -> engine
+	// If we skip inFlight nodes AND engine hasn't committed, count = 0 (BUG!)
 	pendingCreates := int64(0)
 	pendingUpdates := int64(0)
+	inFlightCreates := int64(0)
 	for id := range ae.nodeCache {
-		if ae.inFlightNodes[id] {
-			continue // Don't count in-flight (already in engine)
-		}
 		if ae.updateNodes[id] {
 			pendingUpdates++ // Exists in engine, just updating
 			continue
 		}
+		if ae.inFlightNodes[id] {
+			inFlightCreates++ // Being written to engine right now
+			continue
+		}
 		pendingCreates++
 	}
+	// Also count nodes that are in-flight but NOT updates (they're being created)
 	pendingDeletes := int64(len(ae.deleteNodes))
 
 	engineCount, err := ae.engine.NodeCount()
@@ -994,7 +984,8 @@ func (ae *AsyncEngine) NodeCount() (int64, error) {
 
 	// Adjust for pending creates and deletes
 	// Note: pendingUpdates don't change count (already counted in engineCount)
-	count := engineCount + pendingCreates - pendingDeletes
+	// Include inFlightCreates because they're being written but not yet in engineCount
+	count := engineCount + pendingCreates + inFlightCreates - pendingDeletes
 
 	// Clamp to zero if negative (should never happen, log for debugging)
 	if count < 0 {
@@ -1011,16 +1002,19 @@ func (ae *AsyncEngine) EdgeCount() (int64, error) {
 	ae.mu.RLock()
 
 	// Count pending creates, excluding:
-	// - in-flight edges (already written to engine)
 	// - update edges (exist in engine, just being modified)
+	// NOTE: We DO count in-flight edges because they are being written to engine
+	// but engine.EdgeCount() won't include them until the write commits.
 	pendingCreates := int64(0)
+	inFlightCreates := int64(0)
 	for id := range ae.edgeCache {
-		if ae.inFlightEdges[id] {
-			continue // Don't count in-flight (already in engine)
-		}
 		if ae.updateEdges[id] {
 			continue // Exists in engine, just updating
 		}
+		if ae.inFlightEdges[id] {
+			inFlightCreates++ // Being written to engine right now
+			continue
+		}
 		pendingCreates++
 	}
 	pendingDeletes := int64(len(ae.deleteEdges))
@@ -1034,7 +1028,8 @@ func (ae *AsyncEngine) EdgeCount() (int64, error) {
 
 	// Adjust for pending creates and deletes
 	// Note: updates don't change count (already counted in engineCount)
-	count := engineCount + pendingCreates - pendingDeletes
+	// Include inFlightCreates because they're being written but not yet in engineCount
+	count := engineCount + pendingCreates + inFlightCreates - pendingDeletes
 
 	// Clamp to zero if negative (should never happen, log for debugging)
 	if count < 0 {