Make GetDBTransaction shareable across multiple finisher methods

- Added support to share same transcation across multiple finisher methods like Update/Find/Delete in GetDBTransaction Fixes #49
vmware-labs · Sep 19, 2023 · 6d33919 · 6d33919
1 parent 34094f9
commit 6d33919
Show file tree

Hide file tree

Showing 6 changed files with 113 additions and 31 deletions.
diff --git a/pkg/datastore/database.go b/pkg/datastore/database.go
@@ -109,7 +109,8 @@ func (db *relationalDb) getDBTransaction(ctx context.Context, tableName string,
 			return nil, err
 		}
 	}
-	tx = tx.Table(tableName)
+	// Refer Example 3 here: https://gorm.io/docs/method_chaining.html#New-Session-Method
+	tx = tx.Debug().Table(tableName)
 	if err = tx.Error; err != nil {
 		err = ErrStartingTx.Wrap(err).WithMap(map[ErrorContextKey]string{
 			DB_NAME:      db.dbName,

diff --git a/pkg/datastore/datastore_with_txcontext_test.go b/pkg/datastore/datastore_with_txcontext_test.go
@@ -41,22 +41,27 @@ func rollbackTx(t *testing.T, tx *gorm.DB) {
 	}
 }
 
-func testTxCrud(t *testing.T, ds datastore.DataStore, ctx context.Context, myCokeApp *App, user1, user2 *AppUser) {
+func testTxCrud(t *testing.T, ds datastore.DataStore, ctx context.Context, user1, user2 *AppUser) {
 	t.Helper()
 	assert := assert.New(t)
 	var err error
 
 	txFetcher := authorizer.SimpleTransactionFetcher{}
 
-	// Querying of previously inserted records should succeed
+	t.Log("Querying created records in single transaction")
+	tx, err := ds.Helper().GetDBTransaction(ctx, datastore.GetTableName(user1), user1)
+	assert.NoError(err)
+	defer rollbackTx(t, tx)
+	txCtx := txFetcher.WithTransactionCtx(ctx, tx)
 	for _, record := range []*AppUser{user1, user2} {
 		queryResult := AppUser{Id: record.Id}
-		err = ds.Find(ctx, &queryResult)
+		err = ds.Find(txCtx, &queryResult)
 		assert.NoError(err)
 		assert.Equal(record, &queryResult)
 	}
+	assert.NoError(tx.Commit().Error)
 
-	// Updating non-key fields in a record should succeed
+	t.Log("Verifying update and find in single transaction")
 	user1.Name = "Jeyhun G."
 	user1.Email = "jeyhun111@mail.com"
 	user1.EmailConfirmed = !user1.EmailConfirmed
@@ -65,58 +70,56 @@ func testTxCrud(t *testing.T, ds datastore.DataStore, ctx context.Context, myCok
 	user2.Email = "jahangir111@mail.com"
 	user2.EmailConfirmed = !user2.EmailConfirmed
 	user2.NumFollowers--
-
-	// TODO: Find, Update and Delete are still not supported due to appending where clauses,
-	// after each call to Find/Update/Delete methods
-
-	tx, err := ds.Helper().GetDBTransaction(ctx, datastore.GetTableName(user1), user1)
+	tx, err = ds.Helper().GetDBTransaction(ctx, datastore.GetTableName(user1), user1)
 	assert.NoError(err)
 	defer rollbackTx(t, tx)
-	txCtx := txFetcher.WithTransactionCtx(ctx, tx)
+	txCtx = txFetcher.WithTransactionCtx(ctx, tx)
 	for _, record := range []*AppUser{user1, user2} {
-		rowsAffected, err := ds.Upsert(txCtx, record)
+		rowsAffected, err := ds.Update(txCtx, record)
 		assert.NoError(err)
 		assert.EqualValues(1, rowsAffected)
-	}
-	assert.NoError(tx.Commit().Error)
 
-	for _, record := range []*AppUser{user1, user2} {
 		queryResult := &AppUser{Id: record.Id}
-		err = ds.Find(ctx, queryResult)
+		err = ds.Find(txCtx, queryResult)
 		assert.NoError(err)
 		assert.Equal(record, queryResult)
 	}
+	assert.NoError(tx.Commit().Error)
 
+	t.Log("Verifying upsert and find in single transaction")
+	user1.NumFollowers++
+	user2.NumFollowers--
 	tx, err = ds.Helper().GetDBTransaction(ctx, datastore.GetTableName(user1), user1)
 	assert.NoError(err)
 	defer rollbackTx(t, tx)
 	txCtx = txFetcher.WithTransactionCtx(ctx, tx)
-	// Upsert operation should be an update for already existing records
-	user1.NumFollowers++
-	user2.NumFollowers--
 	for _, record := range []*AppUser{user1, user2} {
 		rowsAffected, err := ds.Upsert(txCtx, record)
 		assert.NoError(err)
 		assert.EqualValues(1, rowsAffected)
-	}
-	assert.NoError(tx.Commit().Error)
-	for _, record := range []*AppUser{user1, user2} {
+
 		queryResult := &AppUser{Id: record.Id}
-		err = ds.Find(ctx, queryResult)
+		err = ds.Find(txCtx, queryResult)
 		assert.NoError(err)
 		assert.Equal(record, queryResult)
 	}
+	assert.NoError(tx.Commit().Error)
 
-	// Deletion of existing records should not fail, and the records should no longer be found in the DB
+	t.Log("Verifying deletion and find in single transaction")
+	tx, err = ds.Helper().GetDBTransaction(ctx, datastore.GetTableName(user1), user1)
+	assert.NoError(err)
+	defer rollbackTx(t, tx)
+	txCtx = txFetcher.WithTransactionCtx(ctx, tx)
 	for _, record := range []*AppUser{user1, user2} {
-		rowsAffected, err := ds.Delete(ctx, record)
+		rowsAffected, err := ds.Delete(txCtx, record)
 		assert.NoError(err)
 		assert.EqualValues(1, rowsAffected)
 		queryResult := AppUser{Id: record.Id}
-		err = ds.Find(ctx, &queryResult)
+		err = ds.Find(txCtx, &queryResult)
 		assert.ErrorIs(err, ErrRecordNotFound)
 		assert.True(queryResult.AreNonKeyFieldsEmpty())
 	}
+	assert.NoError(tx.Commit().Error)
 }
 
 func BenchmarkTxCrud(b *testing.B) {
@@ -126,17 +129,91 @@ func BenchmarkTxCrud(b *testing.B) {
 	LOG = logger.WithField(datastore.COMP, datastore.SAAS_PERSISTENCE)
 
 	var t testing.T
-	ds, _ := SetupDataStore("BenchmarkCrud")
+	ds, _ := SetupDataStore("BenchmarkTxCrud")
 	defer ds.Reset()
-	myCokeApp, user1, user2 := SetupDbTables(ds)
+	_, user1, user2 := SetupDbTables(ds)
 	for n := 0; n < b.N; n++ {
-		testCrud(&t, ds, CokeAdminCtx, myCokeApp, user1, user2)
+		testTxCrud(&t, ds, CokeAdminCtx, user1, user2)
 	}
 }
 
 func TestTxCrud(t *testing.T) {
 	ds, _ := SetupDataStore("TestTxCrud")
 	defer ds.Reset()
-	myCokeApp, user1, user2 := SetupDbTables(ds)
-	testTxCrud(t, ds, CokeAdminCtx, myCokeApp, user1, user2)
+	_, user1, user2 := SetupDbTables(ds)
+	testTxCrud(t, ds, CokeAdminCtx, user1, user2)
+}
+
+func TestSingleTxCrud(t *testing.T) {
+	assert := assert.New(t)
+	var err error
+
+	ctx := CokeAdminCtx
+	ds, _ := SetupDataStore("TestSingleTxCrud")
+	defer ds.Reset()
+	_, user1, user2 := SetupDbTables(ds)
+
+	txFetcher := authorizer.SimpleTransactionFetcher{}
+	tx, err := ds.Helper().GetDBTransaction(ctx, datastore.GetTableName(user1), user1)
+	assert.NoError(err)
+	defer rollbackTx(t, tx)
+	txCtx := txFetcher.WithTransactionCtx(ctx, tx)
+
+	t.Log("Querying created records in same transaction")
+	for _, record := range []*AppUser{user1, user2} {
+		queryResult := AppUser{Id: record.Id}
+		err = ds.Find(txCtx, &queryResult)
+		assert.NoError(err)
+		assert.Equal(record, &queryResult)
+	}
+
+	t.Log("Verifying update and find in same transaction")
+	user1.Name = "Jeyhun G."
+	user1.Email = "jeyhun111@mail.com"
+	user1.EmailConfirmed = !user1.EmailConfirmed
+	user1.NumFollowers++
+	user2.Name = "Jahangir G."
+	user2.Email = "jahangir111@mail.com"
+	user2.EmailConfirmed = !user2.EmailConfirmed
+	user2.NumFollowers--
+	tx, err = ds.Helper().GetDBTransaction(ctx, datastore.GetTableName(user1), user1)
+	assert.NoError(err)
+	defer rollbackTx(t, tx)
+	txCtx = txFetcher.WithTransactionCtx(ctx, tx)
+	for _, record := range []*AppUser{user1, user2} {
+		rowsAffected, err := ds.Update(txCtx, record)
+		assert.NoError(err)
+		assert.EqualValues(1, rowsAffected)
+
+		queryResult := &AppUser{Id: record.Id}
+		err = ds.Find(txCtx, queryResult)
+		assert.NoError(err)
+		assert.Equal(record, queryResult)
+	}
+
+	t.Log("Verifying upsert and find in same transaction")
+	user1.NumFollowers++
+	user2.NumFollowers--
+	for _, record := range []*AppUser{user1, user2} {
+		rowsAffected, err := ds.Upsert(txCtx, record)
+		assert.NoError(err)
+		assert.EqualValues(1, rowsAffected)
+
+		queryResult := &AppUser{Id: record.Id}
+		err = ds.Find(txCtx, queryResult)
+		assert.NoError(err)
+		assert.Equal(record, queryResult)
+	}
+
+	t.Log("Verifying deletion and find in same transaction")
+	for _, record := range []*AppUser{user1, user2} {
+		rowsAffected, err := ds.Delete(txCtx, record)
+		assert.NoError(err)
+		assert.EqualValues(1, rowsAffected)
+		queryResult := AppUser{Id: record.Id}
+		err = ds.Find(txCtx, &queryResult)
+		assert.ErrorIs(err, ErrRecordNotFound)
+		assert.True(queryResult.AreNonKeyFieldsEmpty())
+	}
+	assert.NoError(tx.Commit().Error)
 }
diff --git a/pkg/datastore/example_multiinstance_test.go b/pkg/datastore/example_multiinstance_test.go
@@ -40,6 +40,7 @@ func ExampleDataStore_multiInstance() {
 
 	// Initializes the Datastore using the metadata authorizer and connection details obtained from the ENV variables.
 	ds, err := datastore.FromEnvWithDB(datastore.GetCompLogger(), mdAuthorizer, instancer, "ExampleDataStore_multiInstance")
+	defer ds.Reset()
 	if err != nil {
 		log.Fatalf("datastore initialization from env errored: %s", err)
 	}

diff --git a/pkg/datastore/example_multitenancy_test.go b/pkg/datastore/example_multitenancy_test.go
@@ -37,6 +37,7 @@ func ExampleDataStore_multiTenancy() {
 
 	// Initializes the Datastore using the metadata authorizer and connection details obtained from the ENV variables.
 	ds, err := datastore.FromEnvWithDB(datastore.GetCompLogger(), mdAuthorizer, nil, "ExampleDataStore_multiTenancy")
+	defer ds.Reset()
 	if err != nil {
 		log.Fatalf("datastore initialization from env errored: %s", err)
 	}

diff --git a/pkg/datastore/noinstancer_test.go b/pkg/datastore/noinstancer_test.go
@@ -42,6 +42,7 @@ func TestDataStoreWithoutInstancer(t *testing.T) {
 	ProdInstanceCtx := instancer.WithInstanceId(ServiceAdminCtx, "Prod")
 
 	ds, _ := datastore.FromEnv(datastore.GetCompLogger(), mdAuthorizer, nil)
+	defer ds.Reset()
 	roleMapping := map[string]dbrole.DbRole{
 		SERVICE_AUDITOR: dbrole.READER,
 		SERVICE_ADMIN:   dbrole.WRITER,

diff --git a/pkg/protostore/example_protostore_test.go b/pkg/protostore/example_protostore_test.go
@@ -16,6 +16,7 @@ func ExampleProtoStore() {
 	myLogger := datastore.GetCompLogger()
 	mdAuthorizer := authorizer.MetadataBasedAuthorizer{}
 	myDatastore, _ := datastore.FromEnvWithDB(myLogger, mdAuthorizer, nil, "ExampleProtoStore")
+	defer myDatastore.Reset()
 	ctx := mdAuthorizer.GetAuthContext("Coke", "service_admin")
 	myProtostore := protostore.GetProtoStore(myLogger, myDatastore)