Merge pull request #47 from basenana/feature/namespace

feat: add namespace
basenana · May 14, 2024 · 962b40b · 962b40b
2 parents fd27195 + 2e13804
commit 962b40b
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 5 deletions.
diff --git a/cmd/apps/ingest.go b/cmd/apps/ingest.go
@@ -41,7 +41,7 @@ var IngestCmd = &cobra.Command{
 }
 
 func ingest(ps string) error {
-	f := friday.Fri.WithContext(context.TODO()).OriginFile(&ps)
+	f := friday.Fri.WithContext(context.TODO()).Namespace("test").OriginFile(&ps)
 	res := &friday.IngestState{}
 	f = f.Ingest(res)
 	if f.Error != nil {

diff --git a/pkg/friday/friday.go b/pkg/friday/friday.go
@@ -116,3 +116,8 @@ func (f *Friday) WithContext(ctx context.Context) *Friday {
 	}
 	return t
 }
+
+func (f *Friday) Namespace(namespace string) *Friday {
+	f.statement.context = context.WithValue(f.statement.context, "namespace", namespace)
+	return f
+}
diff --git a/pkg/vectorstore/postgres/migrate.go b/pkg/vectorstore/postgres/migrate.go
@@ -34,6 +34,22 @@ func buildMigrations() []*gormigrate.Migration {
 				return nil
 			},
 		},
+		{
+			ID: "2024051400",
+			Migrate: func(db *gorm.DB) error {
+				err := db.AutoMigrate(
+					&Index{},
+				)
+
+				if err != nil {
+					return err
+				}
+				// init namespace
+				_ = db.Exec(`UPDATE friday_idx SET namespace='global' WHERE 1=1;`)
+				return nil
+			},
+			Rollback: func(db *gorm.DB) error { return nil },
+		},
 	}
 }
 

diff --git a/pkg/vectorstore/postgres/model.go b/pkg/vectorstore/postgres/model.go
@@ -26,6 +26,7 @@ import (
 type Index struct {
 	ID        string `gorm:"column:id;type:varchar(256);primaryKey"`
 	Name      string `gorm:"column:name;index:index_name"`
+	Namespace string `gorm:"column:namespace;index:index_namespace"`
 	OID       int64  `gorm:"column:oid;index:index_oid"`
 	Group     int    `gorm:"column:idx_group;index:index_group"`
 	ParentID  int64  `gorm:"column:parent_entry_id;index:index_parent_id"`

diff --git a/pkg/vectorstore/postgres/postgres.go b/pkg/vectorstore/postgres/postgres.go
@@ -33,6 +33,8 @@ import (
 	"github.com/basenana/friday/pkg/vectorstore/db"
 )
 
+const defaultNamespace = "global"
+
 type PostgresClient struct {
 	log     logger.Logger
 	dEntity *db.Entity
@@ -72,6 +74,10 @@ func NewPostgresClient(log logger.Logger, postgresUrl string) (*PostgresClient,
 }
 
 func (p *PostgresClient) Store(ctx context.Context, element *models.Element, extra map[string]any) error {
+	namespace := ctx.Value("namespace")
+	if namespace == nil {
+		namespace = defaultNamespace
+	}
 	return p.dEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
 		if extra == nil {
 			extra = make(map[string]interface{})
@@ -91,9 +97,10 @@ func (p *PostgresClient) Store(ctx context.Context, element *models.Element, ext
 		}
 
 		v.Extra = string(b)
+		v.Namespace = namespace.(string)
 
 		vModel := &Index{}
-		res := tx.Where("name = ? AND idx_group = ?", element.Name, element.Group).First(vModel)
+		res := tx.Where("namespace = ? AND name = ? AND idx_group = ?", namespace, element.Name, element.Group).First(vModel)
 		if res.Error != nil && res.Error != gorm.ErrRecordNotFound {
 			return res.Error
 		}
@@ -109,7 +116,7 @@ func (p *PostgresClient) Store(ctx context.Context, element *models.Element, ext
 		}
 
 		vModel.Update(v)
-		res = tx.Where("name = ? AND idx_group = ?", element.Name, element.Group).Updates(vModel)
+		res = tx.Where("namespace = ? AND name = ? AND idx_group = ?", namespace, element.Name, element.Group).Updates(vModel)
 		if res.Error != nil || res.RowsAffected == 0 {
 			if res.RowsAffected == 0 {
 				return errors.New("operation conflict")
@@ -121,6 +128,10 @@ func (p *PostgresClient) Store(ctx context.Context, element *models.Element, ext
 }
 
 func (p *PostgresClient) Search(ctx context.Context, query models.DocQuery, vectors []float32, k int) ([]*models.Doc, error) {
+	namespace := ctx.Value("namespace")
+	if namespace == nil {
+		namespace = defaultNamespace
+	}
 	vectors64 := make([]float64, 0)
 	for _, v := range vectors {
 		vectors64 = append(vectors64, float64(v))
@@ -130,6 +141,7 @@ func (p *PostgresClient) Search(ctx context.Context, query models.DocQuery, vect
 	var res *gorm.DB
 
 	res = p.dEntity.WithContext(ctx)
+	res = res.Where("namespace = ?", namespace)
 	if query.ParentId != 0 {
 		res = res.Where("parent_entry_id = ?", query.ParentId)
 	}
@@ -171,12 +183,16 @@ func (p *PostgresClient) Search(ctx context.Context, query models.DocQuery, vect
 }
 
 func (p *PostgresClient) Get(ctx context.Context, oid int64, name string, group int) (*models.Element, error) {
+	namespace := ctx.Value("namespace")
+	if namespace == nil {
+		namespace = defaultNamespace
+	}
 	vModel := &Index{}
 	var res *gorm.DB
 	if oid == 0 {
-		res = p.dEntity.WithContext(ctx).Where("name = ? AND idx_group = ?", name, group).First(vModel)
+		res = p.dEntity.WithContext(ctx).Where("namespace = ? AND name = ? AND idx_group = ?", namespace, name, group).First(vModel)
 	} else {
-		res = p.dEntity.WithContext(ctx).Where("name = ? AND oid = ? AND idx_group = ?", name, oid, group).First(vModel)
+		res = p.dEntity.WithContext(ctx).Where("namespace = ? AND name = ? AND oid = ? AND idx_group = ?", namespace, name, oid, group).First(vModel)
 	}
 	if res.Error != nil {
 		if res.Error == gorm.ErrRecordNotFound {