add AllocateIncompleteID and remove batch 500 save models

ds.go

@@ -11,6 +11,7 @@ import (
 // Client type
 type Client struct {
 	*datastore.Client
+	AllocateIncompleteID bool
 }
 
 // NewClient creates new ds client which wrap datastore client
@@ -19,5 +20,5 @@ func NewClient(ctx context.Context, projectID string, opts ...option.ClientOptio
 	if err != nil {
 		return nil, err
 	}
-	return &Client{client}, nil
+	return &Client{Client: client}, nil
 }

save.go

@@ -2,6 +2,8 @@ package ds
 
 import (
 	"context"
+
+	"cloud.google.com/go/datastore"
 )
 
 func beforeSave(kind string, src interface{}) {
@@ -21,38 +23,20 @@ func beforeSave(kind string, src interface{}) {
 // SaveModel saves model to datastore
 // kind is optional, if key already set
 // if key was not set in model, will call NewKey with given kind
+// if key is incomplete key and AllocateIncompleteID is true, will call allocate ids before put model to datastore
 func (client *Client) SaveModel(ctx context.Context, kind string, src interface{}) error {
 	beforeSave(kind, src)
 
 	x := src.(KeyGetSetter)
-	key, err := client.Put(ctx, x.GetKey(), x)
-	x.SetKey(key)
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
-const maxPutBatchSize = 500
-
-func (client *Client) saveModels(ctx context.Context, src interface{}) error {
-	xs := valueOf(src)
-	if xs.Len() > maxPutBatchSize {
-		// TODO: refactor error
-		errChan := make(chan error)
-		go func() { errChan <- client.saveModels(ctx, xs.Slice(0, maxPutBatchSize).Interface()) }()
-		go func() { errChan <- client.saveModels(ctx, xs.Slice(maxPutBatchSize, xs.Len()).Interface()) }()
-		err := <-errChan
-		if err != nil {
-			return err
-		}
-		err = <-errChan
+	if key := x.GetKey(); key.Incomplete() && client.AllocateIncompleteID {
+		keys, err := client.AllocateIDs(ctx, []*datastore.Key{key})
 		if err != nil {
 			return err
 		}
-		return nil
+		x.SetKey(keys[0])
 	}
-	err := client.PutModels(ctx, src)
+	key, err := client.Put(ctx, x.GetKey(), x)
+	x.SetKey(key)
 	if err != nil {
 		return err
 	}
@@ -67,7 +51,7 @@ func (client *Client) SaveModels(ctx context.Context, kind string, src interface
 		x := xs.Index(i).Interface()
 		beforeSave(kind, x)
 	}
-	err := client.saveModels(ctx, src)
+	err := client.PutModels(ctx, src)
 	if err != nil {
 		return err
 	}