Fix dropping key range in MemQDB (#783)

* Fix dropping key range in MemQDB * More fixes * Change for explicit error
pg-sharding · Oct 2, 2024 · 54f3a6f · 54f3a6f
1 parent 2a91357
commit 54f3a6f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/qdb/memqdb.go b/qdb/memqdb.go
@@ -200,10 +200,15 @@ func (q *MemQDB) DropKeyRange(_ context.Context, id string) error {
 	q.mu.Lock()
 	defer q.mu.Unlock()
 
-	lock, ok := q.Locks[id]
+	_, ok := q.Krs[id]
 	if !ok {
 		return nil
 	}
+
+	lock, ok := q.Locks[id]
+	if !ok {
+		return spqrerror.New(spqrerror.SPQR_METADATA_CORRUPTION, fmt.Sprintf("no lock in MemQDB for key range \"%s\"", id))
+	}
 	if !lock.TryLock() {
 		return spqrerror.Newf(spqrerror.SPQR_KEYRANGE_ERROR, "key range \"%s\" is locked", id)
 	}

diff --git a/qdb/memqdb_test.go b/qdb/memqdb_test.go
@@ -204,6 +204,7 @@ func TestKeyRanges(t *testing.T) {
 		DistributionId: "dserr",
 	}))
 
+	assert.NoError(memqdb.DropKeyRange(ctx, "nonexistentKeyRange"))
 }
 
 func Test_MemQDB_GetKeyRange(t *testing.T) {