go.mercari.io/datastore@v1.8.2/testsuite/tx_batch.go (about)

     1  package testsuite
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sync"
     7  	"testing"
     8  
     9  	"go.mercari.io/datastore"
    10  )
    11  
    12  func transactionBatchPut(ctx context.Context, t *testing.T, client datastore.Client) {
    13  	defer func() {
    14  		err := client.Close()
    15  		if err != nil {
    16  			t.Fatal(err)
    17  		}
    18  	}()
    19  
    20  	type Data struct {
    21  		Str string
    22  	}
    23  
    24  	tx, err := client.NewTransaction(ctx)
    25  	if err != nil {
    26  		t.Fatal(err.Error())
    27  	}
    28  
    29  	// ざっくりした流れ
    30  	//   1. TxBatch で2件Putする
    31  	//   2. PendingKeyをKey化する
    32  	//   3. sync.WaitGroup で待ち合わせする
    33  
    34  	b := tx.Batch()
    35  	wg := sync.WaitGroup{}
    36  	cs := make([]chan datastore.Commit, 0)
    37  	{ // 1st entity
    38  		key := client.IncompleteKey("Data", nil)
    39  		wg.Add(1)
    40  		b.Put(key, &Data{"Hi!"}, func(pKey datastore.PendingKey, err error) error {
    41  			if err != nil {
    42  				return err
    43  			}
    44  			c := make(chan datastore.Commit)
    45  			cs = append(cs, c)
    46  			go func() {
    47  				commit := <-c
    48  				t.Logf("#1: %s", commit.Key(pKey).String())
    49  				wg.Done()
    50  			}()
    51  			return nil
    52  		})
    53  	}
    54  	{ // 2nd entity
    55  		key := client.IncompleteKey("Data", nil)
    56  		wg.Add(1)
    57  		b.Put(key, &Data{"Hi!"}, func(pKey datastore.PendingKey, err error) error {
    58  			if err != nil {
    59  				return err
    60  			}
    61  			c := make(chan datastore.Commit)
    62  			cs = append(cs, c)
    63  			go func() {
    64  				commit := <-c
    65  				t.Logf("#2: %s", commit.Key(pKey).String())
    66  				wg.Done()
    67  			}()
    68  			return nil
    69  		})
    70  	}
    71  
    72  	err = b.Exec()
    73  	if err != nil {
    74  		t.Fatal(err.Error())
    75  	}
    76  
    77  	commit, err := tx.Commit()
    78  	if err != nil {
    79  		t.Fatal(err.Error())
    80  	}
    81  	for _, c := range cs {
    82  		c <- commit
    83  	}
    84  
    85  	wg.Wait()
    86  }
    87  
    88  func transactionBatchPutWithCustomErrHandler(ctx context.Context, t *testing.T, client datastore.Client) {
    89  	defer func() {
    90  		err := client.Close()
    91  		if err != nil {
    92  			t.Fatal(err)
    93  		}
    94  	}()
    95  
    96  	type Data struct {
    97  		Str string
    98  	}
    99  
   100  	tx, err := client.NewTransaction(ctx)
   101  	if err != nil {
   102  		t.Fatal(err.Error())
   103  	}
   104  
   105  	// ざっくりした流れ
   106  	//   1. TxBatch で2件Putする
   107  	//   2. PendingKeyをKey化する
   108  
   109  	b := tx.Batch()
   110  	testErr := errors.New("test")
   111  	{ // 1st entity
   112  		key := client.IncompleteKey("Data", nil)
   113  		b.Put(key, &Data{"Hi!"}, func(pKey datastore.PendingKey, err error) error {
   114  			return testErr
   115  		})
   116  	}
   117  	{ // 2nd entity
   118  		key := client.IncompleteKey("Data", nil)
   119  		b.Put(key, &Data{"Hi!"}, nil)
   120  	}
   121  
   122  	err = b.Exec()
   123  	if err == nil {
   124  		t.Fatal(err.Error())
   125  	}
   126  
   127  	merr, ok := err.(datastore.MultiError)
   128  	if !ok {
   129  		t.Fatalf("unexpected: %v", ok)
   130  	}
   131  	if v := len(merr); v != 1 {
   132  		t.Fatalf("unexpected: %v", ok)
   133  	}
   134  	if v := merr[0]; v != testErr {
   135  		t.Errorf("unexpected: %v", v)
   136  	}
   137  
   138  	_, err = tx.Commit()
   139  	if err != nil {
   140  		t.Fatal(err.Error())
   141  	}
   142  }
   143  
   144  func transactionBatchPutAndAllocateIDs(ctx context.Context, t *testing.T, client datastore.Client) {
   145  	defer func() {
   146  		err := client.Close()
   147  		if err != nil {
   148  			t.Fatal(err)
   149  		}
   150  	}()
   151  
   152  	type Data struct {
   153  		Str string
   154  	}
   155  
   156  	tx, err := client.NewTransaction(ctx)
   157  	if err != nil {
   158  		t.Fatal(err.Error())
   159  	}
   160  
   161  	b := tx.Batch()
   162  	{ // 1st entity
   163  		keys, err := client.AllocateIDs(ctx, []datastore.Key{client.IncompleteKey("Data", nil)})
   164  		if err != nil {
   165  			t.Fatal(err.Error())
   166  		}
   167  		key := keys[0]
   168  		b.Put(key, &Data{"Hi!"}, func(pKey datastore.PendingKey, err error) error {
   169  			if err != nil {
   170  				return err
   171  			}
   172  			t.Logf("#1: %s", key.String())
   173  			return nil
   174  		})
   175  	}
   176  	{ // 2nd entity
   177  		keys, err := client.AllocateIDs(ctx, []datastore.Key{client.IncompleteKey("Data", nil)})
   178  		if err != nil {
   179  			t.Fatal(err.Error())
   180  		}
   181  		key := keys[0]
   182  		b.Put(key, &Data{"Hi!"}, func(pKey datastore.PendingKey, err error) error {
   183  			if err != nil {
   184  				return err
   185  			}
   186  			t.Logf("#2: %s", key.String())
   187  			return nil
   188  		})
   189  	}
   190  
   191  	err = b.Exec()
   192  	if err != nil {
   193  		t.Fatal(err.Error())
   194  	}
   195  
   196  	_, err = tx.Commit()
   197  	if err != nil {
   198  		t.Fatal(err.Error())
   199  	}
   200  }
   201  
   202  func transactionBatchGet(ctx context.Context, t *testing.T, client datastore.Client) {
   203  	defer func() {
   204  		err := client.Close()
   205  		if err != nil {
   206  			t.Fatal(err)
   207  		}
   208  	}()
   209  
   210  	type Data struct {
   211  		Str string
   212  	}
   213  
   214  	key1, err := client.Put(ctx, client.IncompleteKey("Data", nil), &Data{"Data 1"})
   215  	if err != nil {
   216  		t.Fatal(err.Error())
   217  	}
   218  	key2, err := client.Put(ctx, client.IncompleteKey("Data", nil), &Data{"Data 2"})
   219  	if err != nil {
   220  		t.Fatal(err.Error())
   221  	}
   222  
   223  	tx, err := client.NewTransaction(ctx)
   224  	if err != nil {
   225  		t.Fatal(err.Error())
   226  	}
   227  
   228  	b := tx.Batch()
   229  	{ // 1st entity
   230  		dst := &Data{}
   231  		b.Get(key1, dst, func(err error) error {
   232  			if err != nil {
   233  				return err
   234  			}
   235  			t.Logf("#1: %s", dst.Str)
   236  			if v := dst.Str; v != "Data 1" {
   237  				t.Logf("unexpected: %v", v)
   238  			}
   239  			return nil
   240  		})
   241  	}
   242  	{ // 2nd entity
   243  		dst := &Data{}
   244  		b.Get(key2, dst, func(err error) error {
   245  			if err != nil {
   246  				return err
   247  			}
   248  			t.Logf("#2: %s", dst.Str)
   249  			if v := dst.Str; v != "Data 2" {
   250  				t.Logf("unexpected: %v", v)
   251  			}
   252  			return nil
   253  		})
   254  	}
   255  
   256  	err = b.Exec()
   257  	if err != nil {
   258  		t.Fatal(err.Error())
   259  	}
   260  
   261  	err = tx.Rollback()
   262  	if err != nil {
   263  		t.Fatal(err.Error())
   264  	}
   265  }
   266  
   267  func transactionBatchGetWithCustomErrHandler(ctx context.Context, t *testing.T, client datastore.Client) {
   268  	defer func() {
   269  		err := client.Close()
   270  		if err != nil {
   271  			t.Fatal(err)
   272  		}
   273  	}()
   274  
   275  	type Data struct {
   276  		Str string
   277  	}
   278  
   279  	key1, err := client.Put(ctx, client.IncompleteKey("Data", nil), &Data{"Data 1"})
   280  	if err != nil {
   281  		t.Fatal(err.Error())
   282  	}
   283  	key2, err := client.Put(ctx, client.IncompleteKey("Data", nil), &Data{"Data 2"})
   284  	if err != nil {
   285  		t.Fatal(err.Error())
   286  	}
   287  
   288  	tx, err := client.NewTransaction(ctx)
   289  	if err != nil {
   290  		t.Fatal(err.Error())
   291  	}
   292  
   293  	b := tx.Batch()
   294  	testErr := errors.New("test")
   295  	{ // 1st entity
   296  		dst := &Data{}
   297  		b.Get(key1, dst, func(err error) error {
   298  			return testErr
   299  		})
   300  	}
   301  	{ // 2nd entity
   302  		dst := &Data{}
   303  		b.Get(key2, dst, nil)
   304  	}
   305  
   306  	err = b.Exec()
   307  	if err == nil {
   308  		t.Fatal(err.Error())
   309  	}
   310  
   311  	merr, ok := err.(datastore.MultiError)
   312  	if !ok {
   313  		t.Fatalf("unexpected: %v", ok)
   314  	}
   315  	if v := len(merr); v != 1 {
   316  		t.Fatalf("unexpected: %v", ok)
   317  	}
   318  	if v := merr[0]; v != testErr {
   319  		t.Errorf("unexpected: %v", v)
   320  	}
   321  
   322  	err = tx.Rollback()
   323  	if err != nil {
   324  		t.Fatal(err.Error())
   325  	}
   326  }
   327  
   328  func transactionBatchDelete(ctx context.Context, t *testing.T, client datastore.Client) {
   329  	defer func() {
   330  		err := client.Close()
   331  		if err != nil {
   332  			t.Fatal(err)
   333  		}
   334  	}()
   335  
   336  	type Data struct {
   337  		Str string
   338  	}
   339  
   340  	key1, err := client.Put(ctx, client.IncompleteKey("Data", nil), &Data{"Data 1"})
   341  	if err != nil {
   342  		t.Fatal(err.Error())
   343  	}
   344  	key2, err := client.Put(ctx, client.IncompleteKey("Data", nil), &Data{"Data 2"})
   345  	if err != nil {
   346  		t.Fatal(err.Error())
   347  	}
   348  
   349  	tx, err := client.NewTransaction(ctx)
   350  	if err != nil {
   351  		t.Fatal(err.Error())
   352  	}
   353  
   354  	b := tx.Batch()
   355  	{ // 1st entity
   356  		b.Delete(key1, func(err error) error {
   357  			if err != nil {
   358  				return err
   359  			}
   360  			t.Logf("#1: %s", key1.String())
   361  			obj := &Data{}
   362  			// we can get entity!
   363  			err = tx.Get(key1, obj)
   364  			return err
   365  		})
   366  	}
   367  	{ // 2nd entity
   368  		b.Delete(key2, func(err error) error {
   369  			if err != nil {
   370  				return err
   371  			}
   372  			t.Logf("#2: %s", key2.String())
   373  			obj := &Data{}
   374  			// we can get entity!
   375  			err = tx.Get(key2, obj)
   376  			return err
   377  		})
   378  	}
   379  
   380  	err = b.Exec()
   381  	if err != nil {
   382  		t.Fatal(err.Error())
   383  	}
   384  
   385  	err = tx.Rollback()
   386  	if err != nil {
   387  		t.Fatal(err.Error())
   388  	}
   389  }
   390  
   391  func transactionBatchDeleteWithCustomErrHandler(ctx context.Context, t *testing.T, client datastore.Client) {
   392  	defer func() {
   393  		err := client.Close()
   394  		if err != nil {
   395  			t.Fatal(err)
   396  		}
   397  	}()
   398  
   399  	type Data struct {
   400  		Str string
   401  	}
   402  
   403  	key1, err := client.Put(ctx, client.IncompleteKey("Data", nil), &Data{"Data 1"})
   404  	if err != nil {
   405  		t.Fatal(err.Error())
   406  	}
   407  	key2, err := client.Put(ctx, client.IncompleteKey("Data", nil), &Data{"Data 2"})
   408  	if err != nil {
   409  		t.Fatal(err.Error())
   410  	}
   411  
   412  	tx, err := client.NewTransaction(ctx)
   413  	if err != nil {
   414  		t.Fatal(err.Error())
   415  	}
   416  
   417  	b := tx.Batch()
   418  	testErr := errors.New("test")
   419  	{ // 1st entity
   420  		b.Delete(key1, func(err error) error {
   421  			return testErr
   422  		})
   423  	}
   424  	{ // 2nd entity
   425  		b.Delete(key2, nil)
   426  	}
   427  
   428  	err = b.Exec()
   429  	if err == nil {
   430  		t.Fatal(err.Error())
   431  	}
   432  
   433  	merr, ok := err.(datastore.MultiError)
   434  	if !ok {
   435  		t.Fatalf("unexpected: %v", ok)
   436  	}
   437  	if v := len(merr); v != 1 {
   438  		t.Fatalf("unexpected: %v", ok)
   439  	}
   440  	if v := merr[0]; v != testErr {
   441  		t.Errorf("unexpected: %v", v)
   442  	}
   443  
   444  	err = tx.Rollback()
   445  	if err != nil {
   446  		t.Fatal(err.Error())
   447  	}
   448  }