From ebcf22108b238c1c7734f5f48ff1c182f2e253a0 Mon Sep 17 00:00:00 2001
From: asabya <asabya@gmail.com>
Date: Thu, 11 May 2023 10:33:54 +0530
Subject: [PATCH] feat: #149, add count in kv root manifest

---
 pkg/collection/entry.go     |  1 +
 pkg/collection/index.go     |  2 +-
 pkg/collection/index_api.go |  8 ++++++++
 pkg/collection/kv.go        | 12 ++----------
 pkg/collection/kv_test.go   | 35 +++++++++++++++++++++++++++++++++++
 pkg/dfs/kv_api.go           |  1 +
 6 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/pkg/collection/entry.go b/pkg/collection/entry.go
index 13af4d62..c942dd5d 100644
--- a/pkg/collection/entry.go
+++ b/pkg/collection/entry.go
@@ -24,6 +24,7 @@ type Manifest struct {
 	IdxType      IndexType `json:"index_type"`
 	CreationTime int64     `json:"creation_time"`
 	Entries      []*Entry  `json:"entries,omitempty"`
+	Count        uint64    `json:"count,omitempty"` // number of entries in the kv table, this should be updated on root manifest
 	dirtyFlag    bool
 }
 
diff --git a/pkg/collection/index.go b/pkg/collection/index.go
index 73106bb3..9eea7689 100644
--- a/pkg/collection/index.go
+++ b/pkg/collection/index.go
@@ -168,7 +168,7 @@ func OpenIndex(podName, collectionName, indexName, podPassword string, fd *feed.
 		accountInfo:        ai,
 		feed:               fd,
 		client:             client,
-		count:              0,
+		count:              manifest.Count,
 		memDB:              manifest,
 		logger:             logger,
 	}
diff --git a/pkg/collection/index_api.go b/pkg/collection/index_api.go
index cd5c7b9e..2d3addbe 100644
--- a/pkg/collection/index_api.go
+++ b/pkg/collection/index_api.go
@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"path/filepath"
 	"strings"
+	"sync/atomic"
 	"time"
 )
 
@@ -120,6 +121,9 @@ func (idx *Index) Delete(key string) ([][]byte, error) {
 	}
 
 	manifest.Entries = append(manifest.Entries[:i], manifest.Entries[i+1:]...)
+	var delta int = -1
+	atomic.AddUint64(&idx.count, uint64(delta))
+	manifest.Count = idx.count
 	err = idx.updateManifest(manifest, idx.encryptionPassword)
 	if err != nil {
 		return nil, err
@@ -315,6 +319,10 @@ func (idx *Index) addOrUpdateStringEntry(ctx context.Context, manifest *Manifest
 	}
 
 	if entryAdded && !memory {
+		// update the count
+		atomic.AddUint64(&idx.count, 1)
+		manifest.Count = idx.count
+
 		return idx.updateManifest(manifest, idx.encryptionPassword)
 	}
 	return nil // skipcq: TCV-001
diff --git a/pkg/collection/kv.go b/pkg/collection/kv.go
index 39c16268..cb2a5ca8 100644
--- a/pkg/collection/kv.go
+++ b/pkg/collection/kv.go
@@ -223,12 +223,8 @@ func (kv *KeyValue) KVCount(name, encryptionPassword string) (*TableKeyCount, er
 	kv.openKVTMu.Lock()
 	defer kv.openKVTMu.Unlock()
 	if table, ok := kv.openKVTables[name]; ok {
-		count, err := table.index.CountIndex(table.index.encryptionPassword)
-		if err != nil {
-			return nil, err
-		}
 		return &TableKeyCount{
-			Count:     count,
+			Count:     table.index.count,
 			TableName: name,
 		}, nil
 	} else {
@@ -236,12 +232,8 @@ func (kv *KeyValue) KVCount(name, encryptionPassword string) (*TableKeyCount, er
 		if err != nil {
 			return nil, err
 		}
-		count, err := idx.CountIndex(idx.encryptionPassword)
-		if err != nil {
-			return nil, err
-		}
 		return &TableKeyCount{
-			Count:     count,
+			Count:     idx.count,
 			TableName: name,
 		}, nil
 	}
diff --git a/pkg/collection/kv_test.go b/pkg/collection/kv_test.go
index 42d24aa2..e5f99e1b 100644
--- a/pkg/collection/kv_test.go
+++ b/pkg/collection/kv_test.go
@@ -457,6 +457,15 @@ func TestKeyValueStore(t *testing.T) {
 			t.Fatal("values do not match", string(value), "value1")
 		}
 
+		// test count
+		countObject, err := kvStore.KVCount("kv_table_11", podPassword)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if countObject.Count != 2 {
+			t.Fatal("kv count value should be two")
+		}
 		// delete the key
 		_, err = kvStore.KVDelete("kv_table_11", "key1")
 		if err != nil {
@@ -558,6 +567,15 @@ func TestKeyValueStore(t *testing.T) {
 		if !bytes.Equal(value, gotValue) {
 			t.Fatal("values do not match", string(value), string(gotValue))
 		}
+
+		// check the count
+		countObject, err := kvStore.KVCount("kv_table_batch_2", podPassword)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if countObject.Count != 1 {
+			t.Fatal("kv count value should be one")
+		}
 	})
 
 	t.Run("batch_put_columns_and_get_values", func(t *testing.T) {
@@ -638,6 +656,15 @@ func TestKeyValueStore(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
+
+		// check the count
+		countObject, err := kvStore.KVCount("kv_table_Itr_0", podPassword)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if countObject.Count != 100 {
+			t.Fatal("kv count value should be 100")
+		}
 		sortedKeys, sortedValues := sortLexicographically(t, keys, values)
 
 		itr, err := kvStore.KVSeek("kv_table_Itr_0", "", "", -1)
@@ -816,6 +843,14 @@ func TestKeyValueStore(t *testing.T) {
 			}
 		}
 
+		// check the count
+		count, err := kvStore.KVCount(fmt.Sprintf("kv_table_Itr_1%d", tableNo), podPassword)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if count.Count != uint64(len(list)) {
+			t.Fatal("count mismatch", count, len(list))
+		}
 		startIndex := 0
 		endIndex := 0
 
diff --git a/pkg/dfs/kv_api.go b/pkg/dfs/kv_api.go
index 91254b59..3c1e06cd 100644
--- a/pkg/dfs/kv_api.go
+++ b/pkg/dfs/kv_api.go
@@ -22,6 +22,7 @@ import (
 
 type KVGetter interface {
 	KVGet(name, key string) ([]string, []byte, error)
+	OpenKVTable(name, encryptionPassword string) error
 }
 
 // KVCreate does validation checks and calls the create KVtable function.