feat: CreateEntity returning *Entity

oom-ai · Nov 10, 2021 · 9964c53 · 9964c53
1 parent d4016f4
commit 9964c53
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/internal/database/metadatav2/mock_metadata/store.go b/internal/database/metadatav2/mock_metadata/store.go
diff --git a/internal/database/metadatav2/postgres/entity.go b/internal/database/metadatav2/postgres/entity.go
@@ -9,15 +9,16 @@ import (
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
-func (db *DB) CreateEntity(ctx context.Context, opt types.CreateEntityOpt) error {
-	query := "insert into feature_entity(name, length, description) values($1, $2, $3)"
-	_, err := db.ExecContext(ctx, query, opt.Name, opt.Length, opt.Description)
+func (db *DB) CreateEntity(ctx context.Context, opt types.CreateEntityOpt) (*types.Entity, error) {
+	var entity types.Entity
+	query := "insert into feature_entity(name, length, description) values($1, $2, $3) returning *"
+	err := db.GetContext(ctx, &entity, query, opt.Name, opt.Length, opt.Description)
 	if er, ok := err.(*pq.Error); ok {
 		if er.Code == pgerrcode.UniqueViolation {
-			return fmt.Errorf("entity %s already exists", opt.Name)
+			return nil, fmt.Errorf("entity %s already exists", opt.Name)
 		}
 	}
-	return err
+	return &entity, err
 }
 
 func (db *DB) UpdateEntity(ctx context.Context, opt types.UpdateEntityOpt) (int64, error) {

diff --git a/internal/database/metadatav2/store.go b/internal/database/metadatav2/store.go
@@ -10,7 +10,7 @@ import (
 
 type Store interface {
 	// entity
-	CreateEntity(ctx context.Context, opt types.CreateEntityOpt) error
+	CreateEntity(ctx context.Context, opt types.CreateEntityOpt) (*types.Entity, error)
 	UpdateEntity(ctx context.Context, opt types.UpdateEntityOpt) (int64, error)
 	GetEntity(ctx context.Context, name string) *typesv2.Entity
 	ListEntity(ctx context.Context) typesv2.EntityList