chore: add some new stuff

README.md

@@ -1,2 +1,8 @@
 # yowrap
+
+[![GoDoc](https://img.shields.io/static/v1?label=godoc&message=reference&color=blue)](https://pkg.go.dev/github.com/jferrl/yowrap)
+[![Test Status](https://github.com/jferrl/yowrap/workflows/tests/badge.svg)](https://github.com/jferrl/yowrap/actions?query=workflow%3Atests)
+[![codecov](https://codecov.io/gh/jferrl/yowrap/branch/main/graph/badge.svg?token=68I4BZF235)](https://codecov.io/gh/jferrl/yowrap)
+[![Go Report Card](https://goreportcard.com/badge/github.com/jferrl/yowrap)](https://goreportcard.com/report/github.com/jferrl/yowrap)
+
 Lightweight Go package designed to effortlessly wrap the [Yo](https://github.com/cloudspannerecosystem/yo) package for enhanced functionality

hook.go

@@ -10,15 +10,24 @@ import (
 type Hook int
 
 const (
-	// Insert is executed when the Insert method is called.
-	Insert Hook = iota + 1
-	// Update is executed when the Update method is called.
-	Update
-	// InsertOrUpdate is executed when the InsertOrUpdate method is called.
-	InsertOrUpdate
-	// Delete is executed when the Delete method is called.
-	Delete
+	// AfterInsert is executed when the Insert txn is called.
+	AfterInsert Hook = iota + 1
+	// AfterUpdate is executed when the Update txn is called.
+	AfterUpdate
+	// AfterInsertOrUpdate is executed when the InsertOrUpdate txn is called.
+	AfterInsertOrUpdate
+	// AfterDelete is executed when the Delete txn is called.
+	AfterDelete
+	// BeforeInsert is executed before the Insert txn is called.
+	BeforeInsert
+	// BeforeUpdate is executed before the Update txn is called.
+	BeforeUpdate
+	// BeforeInsertOrUpdate is executed before the InsertOrUpdate txn is called.
+	BeforeInsertOrUpdate
+	// BeforeDelete is executed before the Delete txn is called.
+	BeforeDelete
 )
 
-// HookFunc is a trigger function associated with a kind of mutation.
-type HookFunc func(context.Context) []*spanner.Mutation
+// HookFunc defines a function that can be executed during a mutation.
+// Is invoked with the model that is being mutated.
+type HookFunc[T Yo] func(context.Context, *Model[T], *spanner.ReadWriteTransaction) error

internal/README.md

@@ -0,0 +1,17 @@
+# Internal
+
+This directory contains the internal implementation of the project. It is not intended to be used by the end user.
+This directory only contains yo test models needed in the test cases.
+
+## Structure
+
+The internal directory is structured as follows:
+
+```text
+internal/
+    ├── sql/ // Contains the DDL to generate the test model using yo tool
+    │   ├── ddl.go
+    ├── user/ // Contains the test model and the generated code for the test model
+    │   ├── user.yo.go
+    │   ├── yo_db.yo.go
+```

mutation.go

@@ -0,0 +1,33 @@
+package yowrap
+
+// Mutation defines a mutation that can be executed within spanner.
+type Mutation int
+
+const (
+	// Insert is a mutation that inserts a row into a table.
+	Insert Mutation = iota + 1
+	// Update is a mutation that updates a row in a table.
+	Update
+	// InsertOrUpdate is a mutation that inserts a row into a table. If the row
+	// already exists, it updates it instead. Any column values not explicitly
+	// written are preserved.
+	InsertOrUpdate
+	// Delete is a mutation that deletes a row from a table.
+	Delete
+)
+
+// Hooks defines an action that can be executed during a mutation.
+func (m Mutation) Hooks() (before, after Hook) {
+	switch m {
+	case Insert:
+		return BeforeInsert, AfterInsert
+	case Update:
+		return BeforeUpdate, AfterUpdate
+	case InsertOrUpdate:
+		return BeforeInsertOrUpdate, AfterInsertOrUpdate
+	case Delete:
+		return BeforeDelete, AfterDelete
+	default:
+		return 0, 0
+	}
+}

yo.go

@@ -12,39 +12,38 @@ import (
 // ErrNoClient is returned when the spanner client is not set.
 var ErrNoClient = errors.New("no spanner client")
 
-// YoModel is the interface that wraps the basic methods of the Yo
-// generated model.
-type YoModel interface {
+// Yo is the common interface for the generated model.
+type Yo interface {
 	Insert(ctx context.Context) *spanner.Mutation
 	Update(ctx context.Context) *spanner.Mutation
 	InsertOrUpdate(ctx context.Context) *spanner.Mutation
 	Delete(ctx context.Context) *spanner.Mutation
 }
 
 // Opt is a function that modifies a Model.
-type Opt[T YoModel] func(*Model[T])
+type Opt[T Yo] func(*Model[T])
 
 // WithSpannerClientOption returns an Opt that sets the spanner client.
-func WithSpannerClientOption[T YoModel](c spanner.Client) Opt[T] {
+func WithSpannerClientOption[T Yo](c spanner.Client) Opt[T] {
 	return func(m *Model[T]) {
 		m.Client = c
 	}
 }
 
 // Model is a struct that embeds the generated model and implements the
 // YoModel interface.
-type Model[T YoModel] struct {
-	YoModel
+type Model[T Yo] struct {
+	Yo
 	spanner.Client
 
-	hooks map[Hook][]HookFunc
+	hooks map[Hook]HookFunc[T]
 }
 
 // NewModel returns a new wrapped yo model.
-func NewModel[T YoModel](m T, opts ...Opt[T]) *Model[T] {
+func NewModel[T Yo](m T, opts ...Opt[T]) *Model[T] {
 	mo := &Model[T]{
-		YoModel: m,
-		hooks:   make(map[Hook][]HookFunc),
+		Yo:    m,
+		hooks: make(map[Hook]HookFunc[T]),
 	}
 	for _, opt := range opts {
 		opt(mo)
@@ -54,40 +53,45 @@ func NewModel[T YoModel](m T, opts ...Opt[T]) *Model[T] {
 }
 
 // On registers an action to be executed before or after a method.
-func (m *Model[T]) On(h Hook, f HookFunc) {
-	m.hooks[h] = append(m.hooks[h], f)
+func (m *Model[T]) On(h Hook, f HookFunc[T]) {
+	m.hooks[h] = f
 }
 
-// ApplyInsert inserts the model and applies the mutation.
-func (m *Model[T]) ApplyInsert(ctx context.Context) (time.Time, error) {
-	return m.readWriteTxn(ctx, Insert)
-}
-
-// ApplyUpdate updates the model and applies the mutation.
-func (m *Model[T]) ApplyUpdate(ctx context.Context) (time.Time, error) {
-	return m.readWriteTxn(ctx, Update)
-}
+// Apply executes the mutation against the database.
+func (m *Model[T]) Apply(ctx context.Context, mtype Mutation) (time.Time, error) {
+	return m.Client.ReadWriteTransaction(ctx, func(ctx context.Context, rwt *spanner.ReadWriteTransaction) error {
+		before, after := mtype.Hooks()
 
-// ApplyInsertOrUpdate inserts or updates the model and applies the mutation.
-func (m *Model[T]) ApplyInsertOrUpdate(ctx context.Context) (time.Time, error) {
-	return m.readWriteTxn(ctx, InsertOrUpdate)
-}
+		if f, ok := m.hooks[before]; ok {
+			if err := f(ctx, m, rwt); err != nil {
+				return err
+			}
+		}
 
-// ApplyDelete deletes the model and applies the mutation.
-func (m *Model[T]) ApplyDelete(ctx context.Context) (time.Time, error) {
-	return m.readWriteTxn(ctx, Delete)
-}
+		var mut *spanner.Mutation
+		switch mtype {
+		case Insert:
+			mut = m.Insert(ctx)
+		case Update:
+			mut = m.Update(ctx)
+		case InsertOrUpdate:
+			mut = m.InsertOrUpdate(ctx)
+		case Delete:
+			mut = m.Delete(ctx)
+		default:
+			return errors.New("unknown mutation type")
+		}
 
-func (m *Model[T]) readWriteTxn(ctx context.Context, h Hook) (time.Time, error) {
-	return m.Client.ReadWriteTransaction(ctx, func(ctx context.Context, rwt *spanner.ReadWriteTransaction) error {
-		muts := []*spanner.Mutation{m.Insert(ctx)}
+		if err := rwt.BufferWrite([]*spanner.Mutation{mut}); err != nil {
+			return err
+		}
 
-		if actions, ok := m.hooks[h]; ok {
-			for _, f := range actions {
-				muts = append(muts, f(ctx)...)
+		if f, ok := m.hooks[after]; ok {
+			if err := f(ctx, m, rwt); err != nil {
+				return err
 			}
 		}
 
-		return rwt.BufferWrite(muts)
+		return nil
 	})
 }

yo_test.go

@@ -2,6 +2,7 @@ package yowrap
 
 import (
 	"context"
+	"errors"
 	"os"
 	"sort"
 	"testing"
@@ -49,7 +50,7 @@ func TestModel_Insert(t *testing.T) {
 	}
 }
 
-func TestModel_ApplyInsert(t *testing.T) {
+func TestModel_Apply(t *testing.T) {
 	ctx := context.Background()
 
 	model := &user.User{
@@ -66,13 +67,21 @@ func TestModel_ApplyInsert(t *testing.T) {
 	tests := []struct {
 		name    string
 		fields  fields
-		hooks   []HookFunc
+		mutType Mutation
+		before  HookFunc[*user.User]
+		after   HookFunc[*user.User]
 		wantErr bool
 		want    []*user.User
 	}{
 		{
-			name:   "insert a new user into the database",
-			fields: fields{model: model},
+			name:    "unknown mutation returns an error",
+			fields:  fields{model: model},
+			wantErr: true,
+		},
+		{
+			name:    "insert a new user into the database",
+			fields:  fields{model: model},
+			mutType: Insert,
 			want: []*user.User{
 				{
 					Name:  "John Doe",
@@ -81,33 +90,52 @@ func TestModel_ApplyInsert(t *testing.T) {
 			},
 		},
 		{
-			name: "insert a new user into the database with a hook",
-			fields: fields{
-				model: model,
-			},
-			hooks: []HookFunc{
-				func(ctx context.Context) []*spanner.Mutation {
-					user := &user.User{
-						ID:        uuid.NewString(),
-						Name:      "Jane Doe",
-						Email:     "[email protected]",
-						CreatedAt: spanner.CommitTimestamp,
-						UpdatedAt: spanner.CommitTimestamp,
-					}
-
-					return []*spanner.Mutation{
-						user.Insert(ctx),
-					}
+			name:    "insert or update user into the database",
+			fields:  fields{model: model},
+			mutType: InsertOrUpdate,
+			want: []*user.User{
+				{
+					Name:  "John Doe",
+					Email: "[email protected]",
 				},
 			},
+		},
+		{
+			name:    "update a non-existent user in the database",
+			fields:  fields{model: model},
+			mutType: Update,
+			wantErr: true,
+		},
+		{
+			name:    "insert a new user into the database with hooks",
+			fields:  fields{model: model},
+			mutType: Insert,
+			before: func(_ context.Context, m *Model[*user.User], _ *spanner.ReadWriteTransaction) error {
+				// type assertion to access the embedded model
+				u, ok := m.Yo.(*user.User)
+				if !ok {
+					return errors.New("unable to type assert to *user.User")
+				}
+
+				u.Name = "Jane Doe"
+				return nil
+			},
+			after: func(ctx context.Context, m *Model[*user.User], rwt *spanner.ReadWriteTransaction) error {
+				// type assertion to access the embedded model
+				u, ok := m.Yo.(*user.User)
+				if !ok {
+					return errors.New("unable to type assert to *user.User")
+				}
+				u.Email = "[email protected]"
+
+				return rwt.BufferWrite([]*spanner.Mutation{
+					m.Update(ctx),
+				})
+			},
 			want: []*user.User{
 				{
 					Name:  "Jane Doe",
-					Email: "[email protected]",
-				},
-				{
-					Name:  "John Doe",
-					Email: "[email protected]",
+					Email: "[email protected]",
 				},
 			},
 		},
@@ -121,11 +149,15 @@ func TestModel_ApplyInsert(t *testing.T) {
 				WithSpannerClientOption[*user.User](*spannerClient),
 			)
 
-			for _, h := range tt.hooks {
-				m.On(Insert, h)
+			before, after := tt.mutType.Hooks()
+			if tt.before != nil {
+				m.On(before, tt.before)
+			}
+			if tt.after != nil {
+				m.On(after, tt.after)
 			}
 
-			if _, err := m.ApplyInsert(ctx); (err != nil) != tt.wantErr {
+			if _, err := m.Apply(ctx, tt.mutType); (err != nil) != tt.wantErr {
 				t.Errorf("Model.ApplyInsert() error = %v, wantErr %v", err, tt.wantErr)
 			}