Merge pull request weaviate#4967 from weaviate/merge-stable-v1.24-int…

…o-stable-v1.25 Merge stable/v1.24 into stable/v1.25
TraceMachina · May 17, 2024 · e1add93 · e1add93
2 parents b6ec268 + 5edf354
commit e1add93
Show file tree

Hide file tree

Showing 10 changed files with 136 additions and 10 deletions.
diff --git a/adapters/repos/db/lsmkv/bucket.go b/adapters/repos/db/lsmkv/bucket.go
@@ -359,6 +359,47 @@ func (b *Bucket) Get(key []byte) ([]byte, error) {
 	return b.disk.get(key)
 }
 
+func (b *Bucket) GetErrDeleted(key []byte) ([]byte, error) {
+	b.flushLock.RLock()
+	defer b.flushLock.RUnlock()
+
+	v, err := b.active.get(key)
+	if err == nil {
+		// item found and no error, return and stop searching, since the strategy
+		// is replace
+		return v, nil
+	}
+	if errors.Is(err, lsmkv.Deleted) {
+		// deleted in the mem-table (which is always the latest) means we don't
+		// have to check the disk segments, return nil now
+		return nil, err
+	}
+
+	if err != lsmkv.NotFound {
+		panic(fmt.Sprintf("unsupported error in bucket.Get: %v\n", err))
+	}
+
+	if b.flushing != nil {
+		v, err := b.flushing.get(key)
+		if err == nil {
+			// item found and no error, return and stop searching, since the strategy
+			// is replace
+			return v, nil
+		}
+		if errors.Is(err, lsmkv.Deleted) {
+			// deleted in the now most recent memtable  means we don't have to check
+			// the disk segments, return nil now
+			return nil, err
+		}
+
+		if err != lsmkv.NotFound {
+			panic("unsupported error in bucket.Get")
+		}
+	}
+
+	return b.disk.getErrDeleted(key)
+}
+
 // GetBySecondary retrieves an object using one of its secondary keys. A bucket
 // can have an infinite number of secondary keys. Specify the secondary key
 // position as the first argument.

diff --git a/adapters/repos/db/lsmkv/segment_group.go b/adapters/repos/db/lsmkv/segment_group.go
@@ -321,6 +321,38 @@ func (sg *SegmentGroup) getWithUpperSegmentBoundary(key []byte, topMostSegment i
 	return nil, nil
 }
 
+func (sg *SegmentGroup) getErrDeleted(key []byte) ([]byte, error) {
+	sg.maintenanceLock.RLock()
+	defer sg.maintenanceLock.RUnlock()
+
+	return sg.getWithUpperSegmentBoundaryErrDeleted(key, len(sg.segments)-1)
+}
+
+func (sg *SegmentGroup) getWithUpperSegmentBoundaryErrDeleted(key []byte, topMostSegment int) ([]byte, error) {
+	// assumes "replace" strategy
+
+	// start with latest and exit as soon as something is found, thus making sure
+	// the latest takes presence
+	for i := topMostSegment; i >= 0; i-- {
+		v, err := sg.segments[i].get(key)
+		if err != nil {
+			if errors.Is(err, lsmkv.NotFound) {
+				continue
+			}
+
+			if errors.Is(err, lsmkv.Deleted) {
+				return nil, err
+			}
+
+			panic(fmt.Sprintf("unsupported error in segmentGroup.get(): %v", err))
+		}
+
+		return v, nil
+	}
+
+	return nil, nil
+}
+
 func (sg *SegmentGroup) getBySecondaryIntoMemory(pos int, key []byte, buffer []byte) ([]byte, []byte, error) {
 	sg.maintenanceLock.RLock()
 	defer sg.maintenanceLock.RUnlock()

diff --git a/adapters/repos/db/replication.go b/adapters/repos/db/replication.go
@@ -22,6 +22,7 @@ import (
 	"github.com/go-openapi/strfmt"
 	"github.com/pkg/errors"
 	"github.com/weaviate/weaviate/entities/additional"
+	"github.com/weaviate/weaviate/entities/lsmkv"
 	"github.com/weaviate/weaviate/entities/multi"
 	"github.com/weaviate/weaviate/entities/schema"
 	"github.com/weaviate/weaviate/entities/storobj"
@@ -323,7 +324,10 @@ func (i *Index) OverwriteObjects(ctx context.Context,
 			continue
 		}
 		// valid update
-		found, err := s.ObjectByID(ctx, data.ID, nil, additional.Properties{})
+		found, err := s.ObjectByIDErrDeleted(ctx, data.ID, nil, additional.Properties{})
+		if err != nil && errors.Is(err, lsmkv.Deleted) {
+			continue
+		}
 		var curUpdateTime int64 // 0 means object doesn't exist on this node
 		if found != nil {
 			curUpdateTime = found.LastUpdateTimeUnix()

diff --git a/adapters/repos/db/shard.go b/adapters/repos/db/shard.go
@@ -83,6 +83,7 @@ type ShardLike interface {
 	PutObject(context.Context, *storobj.Object) error
 	PutObjectBatch(context.Context, []*storobj.Object) []error
 	ObjectByID(ctx context.Context, id strfmt.UUID, props search.SelectProperties, additional additional.Properties) (*storobj.Object, error)
+	ObjectByIDErrDeleted(ctx context.Context, id strfmt.UUID, props search.SelectProperties, additional additional.Properties) (*storobj.Object, error)
 	Exists(ctx context.Context, id strfmt.UUID) (bool, error)
 	ObjectSearch(ctx context.Context, limit int, filters *filters.LocalFilter, keywordRanking *searchparams.KeywordRanking, sort []filters.Sort, cursor *filters.Cursor, additional additional.Properties) ([]*storobj.Object, []float32, error)
 	ObjectVectorSearch(ctx context.Context, searchVector []float32, targetVector string, targetDist float32, limit int, filters *filters.LocalFilter, sort []filters.Sort, groupBy *searchparams.GroupBy, additional additional.Properties) ([]*storobj.Object, []float32, error)

diff --git a/adapters/repos/db/shard_lazyloader.go b/adapters/repos/db/shard_lazyloader.go
@@ -208,6 +208,13 @@ func (l *LazyLoadShard) ObjectByID(ctx context.Context, id strfmt.UUID, props se
 	return l.shard.ObjectByID(ctx, id, props, additional)
 }
 
+func (l *LazyLoadShard) ObjectByIDErrDeleted(ctx context.Context, id strfmt.UUID, props search.SelectProperties, additional additional.Properties) (*storobj.Object, error) {
+	if err := l.Load(ctx); err != nil {
+		return nil, err
+	}
+	return l.shard.ObjectByIDErrDeleted(ctx, id, props, additional)
+}
+
 func (l *LazyLoadShard) Exists(ctx context.Context, id strfmt.UUID) (bool, error) {
 	if err := l.Load(ctx); err != nil {
 		return false, err

diff --git a/adapters/repos/db/shard_read.go b/adapters/repos/db/shard_read.go
@@ -35,6 +35,29 @@ import (
 	"github.com/weaviate/weaviate/entities/storobj"
 )
 
+func (s *Shard) ObjectByIDErrDeleted(ctx context.Context, id strfmt.UUID, props search.SelectProperties, additional additional.Properties) (*storobj.Object, error) {
+	idBytes, err := uuid.MustParse(id.String()).MarshalBinary()
+	if err != nil {
+		return nil, err
+	}
+
+	bytes, err := s.store.Bucket(helpers.ObjectsBucketLSM).GetErrDeleted(idBytes)
+	if err != nil {
+		return nil, err
+	}
+
+	if bytes == nil {
+		return nil, nil
+	}
+
+	obj, err := storobj.FromBinary(bytes)
+	if err != nil {
+		return nil, errors.Wrap(err, "unmarshal object")
+	}
+
+	return obj, nil
+}
+
 func (s *Shard) ObjectByID(ctx context.Context, id strfmt.UUID, props search.SelectProperties, additional additional.Properties) (*storobj.Object, error) {
 	s.activityTracker.Add(1)
 	idBytes, err := uuid.MustParse(id.String()).MarshalBinary()

diff --git a/modules/generative-palm/clients/palm.go b/modules/generative-palm/clients/palm.go
@@ -188,7 +188,7 @@ func (v *palm) Generate(ctx context.Context, cfg moduletools.ClassConfig, prompt
 		return nil, errors.Wrap(err, "create POST request")
 	}
 
-	apiKey, err := v.getApiKey(ctx)
+	apiKey, err := v.getApiKey(ctx, useGenerativeAIEndpoint)
 	if err != nil {
 		return nil, errors.Wrapf(err, "Google API Key")
 	}
@@ -413,7 +413,16 @@ func (v *palm) generateForPrompt(textProperties map[string]string, prompt string
 	return prompt, nil
 }
 
-func (v *palm) getApiKey(ctx context.Context) (string, error) {
+func (v *palm) getApiKey(ctx context.Context, useGenerativeAIEndpoint bool) (string, error) {
+	if useGenerativeAIEndpoint {
+		if apiKeyValue := v.getValueFromContext(ctx, "X-Google-Studio-Api-Key"); apiKeyValue != "" {
+			return apiKeyValue, nil
+		}
+	} else {
+		if apiKeyValue := v.getValueFromContext(ctx, "X-Google-Vertex-Api-Key"); apiKeyValue != "" {
+			return apiKeyValue, nil
+		}
+	}
 	if apiKeyValue := v.getValueFromContext(ctx, "X-Google-Api-Key"); apiKeyValue != "" {
 		return apiKeyValue, nil
 	}
@@ -424,7 +433,7 @@ func (v *palm) getApiKey(ctx context.Context) (string, error) {
 		return v.apiKey, nil
 	}
 	return "", errors.New("no api key found " +
-		"neither in request header: X-Palm-Api-Key or X-Google-Api-Key " +
+		"neither in request header: X-Palm-Api-Key or X-Google-Api-Key or X-Google-Vertex-Api-Key or X-Google-Studio-Api-Key " +
 		"nor in environment variable under PALM_APIKEY or GOOGLE_APIKEY")
 }
 

diff --git a/modules/text2vec-palm/clients/palm.go b/modules/text2vec-palm/clients/palm.go
@@ -99,7 +99,7 @@ func (v *palm) vectorize(ctx context.Context, input []string, taskType taskType,
 		return nil, errors.Wrap(err, "create POST request")
 	}
 
-	apiKey, err := v.getApiKey(ctx)
+	apiKey, err := v.getApiKey(ctx, useGenerativeAIEndpoint)
 	if err != nil {
 		return nil, errors.Wrapf(err, "Google API Key")
 	}
@@ -179,7 +179,16 @@ func (v *palm) checkResponse(statusCode int, palmApiError *palmApiError) error {
 	return nil
 }
 
-func (v *palm) getApiKey(ctx context.Context) (string, error) {
+func (v *palm) getApiKey(ctx context.Context, useGenerativeAIEndpoint bool) (string, error) {
+	if useGenerativeAIEndpoint {
+		if apiKeyValue := v.getValueFromContext(ctx, "X-Google-Studio-Api-Key"); apiKeyValue != "" {
+			return apiKeyValue, nil
+		}
+	} else {
+		if apiKeyValue := v.getValueFromContext(ctx, "X-Google-Vertex-Api-Key"); apiKeyValue != "" {
+			return apiKeyValue, nil
+		}
+	}
 	if apiKeyValue := v.getValueFromContext(ctx, "X-Google-Api-Key"); apiKeyValue != "" {
 		return apiKeyValue, nil
 	}
@@ -190,7 +199,7 @@ func (v *palm) getApiKey(ctx context.Context) (string, error) {
 		return v.apiKey, nil
 	}
 	return "", errors.New("no api key found " +
-		"neither in request header: X-Palm-Api-Key or X-Google-Api-Key " +
+		"neither in request header: X-Palm-Api-Key or X-Google-Api-Key or X-Google-Vertex-Api-Key or X-Google-Studio-Api-Key " +
 		"nor in environment variable under PALM_APIKEY or GOOGLE_APIKEY")
 }
 

diff --git a/modules/text2vec-palm/clients/palm_test.go b/modules/text2vec-palm/clients/palm_test.go
@@ -143,7 +143,7 @@ func TestClient(t *testing.T) {
 
 		require.NotNil(t, err)
 		assert.Equal(t, err.Error(), "Google API Key: no api key found "+
-			"neither in request header: X-Palm-Api-Key or X-Google-Api-Key "+
+			"neither in request header: X-Palm-Api-Key or X-Google-Api-Key or X-Google-Vertex-Api-Key or X-Google-Studio-Api-Key "+
 			"nor in environment variable under PALM_APIKEY or GOOGLE_APIKEY")
 	})
 
@@ -163,7 +163,7 @@ func TestClient(t *testing.T) {
 
 		require.NotNil(t, err)
 		assert.Equal(t, err.Error(), "Google API Key: no api key found "+
-			"neither in request header: X-Palm-Api-Key or X-Google-Api-Key "+
+			"neither in request header: X-Palm-Api-Key or X-Google-Api-Key or X-Google-Vertex-Api-Key or X-Google-Studio-Api-Key "+
 			"nor in environment variable under PALM_APIKEY or GOOGLE_APIKEY")
 	})
 }

diff --git a/usecases/config/config_handler.go b/usecases/config/config_handler.go
@@ -285,7 +285,7 @@ type CORS struct {
 const (
 	DefaultCORSAllowOrigin  = "*"
 	DefaultCORSAllowMethods = "*"
-	DefaultCORSAllowHeaders = "Content-Type, Authorization, Batch, X-Openai-Api-Key, X-Openai-Organization, X-Openai-Baseurl, X-Anyscale-Baseurl, X-Anyscale-Api-Key, X-Cohere-Api-Key, X-Cohere-Baseurl, X-Huggingface-Api-Key, X-Azure-Api-Key, X-Google-Api-Key, X-Palm-Api-Key, X-Jinaai-Api-Key, X-Aws-Access-Key, X-Aws-Secret-Key, X-Voyageai-Baseurl, X-Voyageai-Api-Key, X-Mistral-Baseurl, X-Mistral-Api-Key, X-OctoAI-Api-Key"
+	DefaultCORSAllowHeaders = "Content-Type, Authorization, Batch, X-Openai-Api-Key, X-Openai-Organization, X-Openai-Baseurl, X-Anyscale-Baseurl, X-Anyscale-Api-Key, X-Cohere-Api-Key, X-Cohere-Baseurl, X-Huggingface-Api-Key, X-Azure-Api-Key, X-Google-Api-Key, X-Google-Vertex-Api-Key, X-Google-Studio-Api-Key, X-Palm-Api-Key, X-Jinaai-Api-Key, X-Aws-Access-Key, X-Aws-Secret-Key, X-Voyageai-Baseurl, X-Voyageai-Api-Key, X-Mistral-Baseurl, X-Mistral-Api-Key, X-OctoAI-Api-Key"
 )
 
 func (r ResourceUsage) Validate() error {