diff --git a/go.sum b/go.sum index 23a74ae..db0aa3d 100644 --- a/go.sum +++ b/go.sum @@ -247,8 +247,8 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/ github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69rRypqCw= -github.com/xDarkicex/memory v1.0.36 h1:O6kUm798Q67JKq4gFeiA4ZGb6kbf69SDHAcgy7Ww5AI= -github.com/xDarkicex/memory v1.0.36/go.mod h1:ucTTiUZMrWXY/nDFkyLMlQ7BnO3qCmG2P2BMJKOSdGc= +github.com/xDarkicex/memory v1.0.37 h1:ge+8VVZwBzVJOXtti3PY7N/BZoE1paOd8P45kxCijDk= +github.com/xDarkicex/memory v1.0.37/go.mod h1:ucTTiUZMrWXY/nDFkyLMlQ7BnO3qCmG2P2BMJKOSdGc= github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= diff --git a/internal/storage/singlefile/engine.go b/internal/storage/singlefile/engine.go index 188a3cf..cd7d4c8 100644 --- a/internal/storage/singlefile/engine.go +++ b/internal/storage/singlefile/engine.go @@ -965,7 +965,7 @@ func (e *Engine) applyCommittedFrames(frames []walRecord) error { } func (e *Engine) applyCreateCollection(name string, config storage.CollectionConfig, lsn uint64) { - if _, exists := e.state.Collections[name]; exists { + if collection := e.state.Collections[name]; collection != nil && !collection.Deleted { return } e.state.Collections[name] = &persistedCollection{ diff --git a/libravdb/collection_management_test.go b/libravdb/collection_management_test.go index c9f08c1..c0dad31 100644 --- a/libravdb/collection_management_test.go +++ b/libravdb/collection_management_test.go @@ -2,6 +2,7 @@ package libravdb import ( "context" + "errors" "fmt" "testing" "time" @@ -9,6 +10,108 @@ import ( "github.com/xDarkicex/libravdb/internal/memory" ) +func TestEnsureCollectionDimensionMismatchIsNonDestructive(t *testing.T) { + ctx := context.Background() + db, err := Open(WithStoragePath(testDBPath(t))) + if err != nil { + t.Fatalf("Failed to create database: %v", err) + } + defer db.Close() + + collection, err := db.EnsureCollection(ctx, "embeddings", 3, WithFlat()) + if err != nil { + t.Fatalf("EnsureCollection create: %v", err) + } + + if err := collection.Insert(ctx, "vec_1", []float32{0.1, 0.2, 0.3}, map[string]interface{}{"source": "original"}); err != nil { + t.Fatalf("Insert: %v", err) + } + + _, err = db.EnsureCollection(ctx, "embeddings", 4, WithFlat()) + if err == nil { + t.Fatal("Expected dimension mismatch error, got nil") + } + if !errors.Is(err, ErrDimensionMismatch) { + t.Fatalf("Expected ErrDimensionMismatch, got %v", err) + } + + var mismatch *CollectionDimensionMismatchError + if !errors.As(err, &mismatch) { + t.Fatalf("Expected CollectionDimensionMismatchError, got %T: %v", err, err) + } + if mismatch.Collection != "embeddings" || mismatch.ExistingDimension != 3 || mismatch.RequestedDimension != 4 { + t.Fatalf("Unexpected mismatch details: %+v", mismatch) + } + + surviving, err := db.GetCollection("embeddings") + if err != nil { + t.Fatalf("GetCollection after mismatch: %v", err) + } + if surviving != collection { + t.Fatal("Expected original collection instance to remain loaded") + } + if surviving.Dimension() != 3 { + t.Fatalf("Expected original dimension 3 to survive, got %d", surviving.Dimension()) + } + + record, err := surviving.Get(ctx, "vec_1") + if err != nil { + t.Fatalf("Expected original vector to survive mismatch: %v", err) + } + if len(record.Vector) != 3 || record.Vector[0] != 0.1 || record.Vector[1] != 0.2 || record.Vector[2] != 0.3 { + t.Fatalf("Unexpected surviving vector: %#v", record.Vector) + } + + stats := surviving.Stats(ctx) + if stats.VectorCount != 1 { + t.Fatalf("Expected one vector after mismatch, got %d", stats.VectorCount) + } +} + +func TestEnsureCollectionRecreateOnDimensionMismatchOptIn(t *testing.T) { + ctx := context.Background() + path := testDBPath(t) + db, err := Open(WithStoragePath(path)) + if err != nil { + t.Fatalf("Failed to create database: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + collection, err := db.EnsureCollection(ctx, "embeddings", 3, WithFlat()) + if err != nil { + t.Fatalf("EnsureCollection create: %v", err) + } + if err := collection.Insert(ctx, "vec_1", []float32{0.1, 0.2, 0.3}, nil); err != nil { + t.Fatalf("Insert: %v", err) + } + + recreated, err := db.EnsureCollectionRecreateOnDimensionMismatch(ctx, "embeddings", 4, WithFlat()) + if err != nil { + t.Fatalf("EnsureCollectionRecreateOnDimensionMismatch: %v", err) + } + if recreated == collection { + t.Fatal("Expected collection to be recreated through explicit opt-in") + } + if err := db.Close(); err != nil { + t.Fatalf("Close before reopen: %v", err) + } + + db, err = Open(WithStoragePath(path)) + if err != nil { + t.Fatalf("Reopen after recreate: %v", err) + } + recreated, err = db.GetCollection("embeddings") + if err != nil { + t.Fatalf("GetCollection after reopen: %v", err) + } + if recreated.Dimension() != 4 { + t.Fatalf("Expected recreated dimension 4 after reopen, got %d", recreated.Dimension()) + } + if stats := recreated.Stats(ctx); stats.VectorCount != 0 { + t.Fatalf("Expected recreated collection to start empty, got %d vectors", stats.VectorCount) + } +} + func TestCollectionMemoryManagement(t *testing.T) { // Create database db, err := Open(WithStoragePath(testDBPath(t))) diff --git a/libravdb/database.go b/libravdb/database.go index cdbb06b..db75d4e 100644 --- a/libravdb/database.go +++ b/libravdb/database.go @@ -187,11 +187,23 @@ func (db *Database) CreateCollection(ctx context.Context, name string, opts ...C // EnsureCollection gets an existing collection, or creates it with the given options. // If the collection exists but its dimension differs from the requested dimension, -// it is dropped and recreated atomically before being returned. -// This prevents corrupted-DB scenarios where EnsureAuthoredCollections ran before -// the embedding model loaded, creating a dim=1 collection instead of the model's -// actual dimension. +// it returns a CollectionDimensionMismatchError without modifying the existing +// collection. func (db *Database) EnsureCollection(ctx context.Context, name string, dimension int, opts ...CollectionOption) (*Collection, error) { + return db.ensureCollection(ctx, name, dimension, false, opts...) +} + +// EnsureCollectionRecreateOnDimensionMismatch gets an existing collection, or +// creates it with the given options. If the collection exists but its dimension +// differs from the requested dimension, it is dropped and recreated. +// +// Prefer EnsureCollection unless destructive recovery from a known-bad schema is +// explicitly intended. +func (db *Database) EnsureCollectionRecreateOnDimensionMismatch(ctx context.Context, name string, dimension int, opts ...CollectionOption) (*Collection, error) { + return db.ensureCollection(ctx, name, dimension, true, opts...) +} + +func (db *Database) ensureCollection(ctx context.Context, name string, dimension int, recreateOnDimensionMismatch bool, opts ...CollectionOption) (*Collection, error) { if dimension <= 0 { return nil, ErrInvalidDimension } @@ -208,14 +220,16 @@ func (db *Database) EnsureCollection(ctx context.Context, name string, dimension if col.Dimension() == dimension { return col, nil } - // Dimension mismatch — drop and recreate atomically. + if !recreateOnDimensionMismatch { + return nil, newCollectionDimensionMismatchError(name, col.Dimension(), dimension) + } if err := db.deleteCollectionLocked(col, name); err != nil { return nil, fmt.Errorf("failed to drop mismatched collection %q: %w", name, err) } // Fall through to create. } - col, err := db.createCollectionLocked(ctx, name, opts...) + col, err := db.createCollectionLocked(ctx, name, ensureCollectionOptions(dimension, opts)...) if err == nil { return col, nil } @@ -225,12 +239,27 @@ func (db *Database) EnsureCollection(ctx context.Context, name string, dimension if col.Dimension() == dimension { return col, nil } - return nil, fmt.Errorf("%w: collection %q has dimension %d, want %d", ErrDimensionMismatch, name, col.Dimension(), dimension) + return nil, newCollectionDimensionMismatchError(name, col.Dimension(), dimension) } } return nil, err } +func newCollectionDimensionMismatchError(name string, existing, requested int) error { + return &CollectionDimensionMismatchError{ + Collection: name, + ExistingDimension: existing, + RequestedDimension: requested, + } +} + +func ensureCollectionOptions(dimension int, opts []CollectionOption) []CollectionOption { + createOpts := make([]CollectionOption, 0, len(opts)+1) + createOpts = append(createOpts, opts...) + createOpts = append(createOpts, WithDimension(dimension)) + return createOpts +} + // createCollectionLocked creates a collection. Caller must hold db.mu. func (db *Database) createCollectionLocked(ctx context.Context, name string, opts ...CollectionOption) (*Collection, error) { collection, err := newCollection(ctx, name, db.storage, db.metrics, db.newWriteController(), opts...) diff --git a/libravdb/errors.go b/libravdb/errors.go index 16bf69c..1577e4c 100644 --- a/libravdb/errors.go +++ b/libravdb/errors.go @@ -23,6 +23,23 @@ var ( ErrEmptyIndex = errors.New("index is empty") ) +// CollectionDimensionMismatchError describes an existing collection whose +// stored vector dimension is incompatible with the requested dimension. +type CollectionDimensionMismatchError struct { + Collection string + ExistingDimension int + RequestedDimension int +} + +func (e *CollectionDimensionMismatchError) Error() string { + return fmt.Sprintf("%v: collection %q has dimension %d, want %d", + ErrDimensionMismatch, e.Collection, e.ExistingDimension, e.RequestedDimension) +} + +func (e *CollectionDimensionMismatchError) Unwrap() error { + return ErrDimensionMismatch +} + // Streaming errors var ( ErrBackpressureActive = errors.New("backpressure is active, cannot send more data") diff --git a/tests/main_test.go b/tests/main_test.go index 4238e69..4c03b3b 100644 --- a/tests/main_test.go +++ b/tests/main_test.go @@ -1,8 +1,9 @@ package tests import ( - "go.uber.org/goleak" "testing" + + "go.uber.org/goleak" ) func TestMain(m *testing.M) {