diff --git a/internal/db/option_test.go b/internal/db/option_test.go index 24e7159c1..60b9c62ab 100644 --- a/internal/db/option_test.go +++ b/internal/db/option_test.go @@ -201,6 +201,7 @@ func Test_getOpts(t *testing.T) { // try setting to true opts = GetOpts(WithDebug(true)) testOpts.withDebug = true + assert.Equal(opts, testOpts) }) t.Run("WithOnConflict", func(t *testing.T) { assert := assert.New(t) diff --git a/internal/db/read_writer.go b/internal/db/read_writer.go index e942b7c1f..a666561f2 100644 --- a/internal/db/read_writer.go +++ b/internal/db/read_writer.go @@ -88,16 +88,18 @@ type Writer interface { // rows updated or an error. Supported options: WithOplog. Update(ctx context.Context, i interface{}, fieldMaskPaths []string, setToNullPaths []string, opt ...Option) (int, error) - // Create an object in the db with options: WithOplog, WithOnConflict. - // The caller is responsible for the transaction life cycle of the writer - // and if an error is returned the caller must decide what to do with - // the transaction, which almost always should be to rollback. + // Create an object in the db with options: WithDebug, WithOplog, NewOplogMsg, + // WithLookup, WithReturnRowsAffected, OnConflict, WithVersion, and + // WithWhere. The caller is responsible for the transaction life cycle of + // the writer and if an error is returned the caller must decide what to do + // with the transaction, which almost always should be to rollback. Create(ctx context.Context, i interface{}, opt ...Option) error // CreateItems will create multiple items of the same type. - // Supported options: WithOplog and WithOplogMsgs. WithOplog and - // WithOplogMsgs may not be used together. WithLookup is not a - // supported option. The caller is responsible for the transaction life + // Supported options: WithDebug, WithOplog, WithOplogMsgs, + // WithReturnRowsAffected, OnConflict, WithVersion, and WithWhere. + /// WithOplog and WithOplogMsgs may not be used together. WithLookup is not + // a supported option. The caller is responsible for the transaction life // cycle of the writer and if an error is returned the caller must decide // what to do with the transaction, which almost always should be to // rollback. @@ -263,12 +265,21 @@ func (rw *Db) lookupAfterWrite(ctx context.Context, i interface{}, opt ...Option return nil } -// Create an object in the db with options: WithOplog, NewOplogMsg, WithLookup -// and OnConflict. WithOplog will write an oplog entry for the create. -// NewOplogMsg will return in-memory oplog message. WithOplog and NewOplogMsg -// cannot be used together. WithLookup with to force a lookup after create. +// Create an object in the db with options: WithDebug, WithOplog, NewOplogMsg, +// WithLookup, WithReturnRowsAffected, OnConflict, WithVersion, and WithWhere. +// +// WithOplog will write an oplog entry for the create. NewOplogMsg will return +// in-memory oplog message. WithOplog and NewOplogMsg cannot be used together. +// WithLookup with to force a lookup after create. +// // OnConflict specifies alternative actions to take when an insert results in a -// unique constraint or exclusion constraint error. +// unique constraint or exclusion constraint error. If WithVersion is used, then +// the update for on conflict will include the version number, which basically +// makes the update use optimistic locking and the update will only succeed if +// the existing rows version matches the WithVersion option. Zero is not a +// valid value for the WithVersion option and will return an error. WithWhere +// allows specifying an additional constraint on the on conflict operation in +// addition to the on conflict target policy (columns or constraint). func (rw *Db) Create(ctx context.Context, i interface{}, opt ...Option) error { const op = "db.Create" if rw.underlying == nil { @@ -345,8 +356,19 @@ func (rw *Db) Create(ctx context.Context, i interface{}, opt ...Option) error { default: return errors.New(ctx, errors.InvalidParameter, op, fmt.Sprintf("invalid conflict action: %v", reflect.TypeOf(opts.withOnConflict.Action))) } + if opts.WithVersion != nil || opts.withWhereClause != "" { + where, args, err := rw.whereClausesFromOpts(ctx, i, opts) + if err != nil { + return errors.Wrap(ctx, err, op) + } + whereConditions := db.Statement.BuildCondition(where, args...) + c.Where = clause.Where{Exprs: whereConditions} + } db = db.Clauses(c) } + if opts.withDebug { + db = db.Debug() + } var ticket *store.Ticket if withOplog { @@ -389,8 +411,34 @@ func (rw *Db) Create(ctx context.Context, i interface{}, opt ...Option) error { return nil } +func (rw *Db) whereClausesFromOpts(ctx context.Context, i interface{}, opts Options) (string, []interface{}, error) { + const op = "db.whereClausesFromOpts" + var where []string + var args []interface{} + if opts.WithVersion != nil { + if *opts.WithVersion == 0 { + return "", nil, errors.New(ctx, errors.InvalidParameter, op, "with version option is zero") + } + mDb := rw.underlying.Model(i) + err := mDb.Statement.Parse(i) + if err != nil && mDb.Statement.Schema == nil { + return "", nil, errors.New(ctx, errors.Unknown, op, "internal error: unable to parse stmt", errors.WithWrap(err)) + } + if !contains(mDb.Statement.Schema.DBNames, "version") { + return "", nil, errors.New(ctx, errors.InvalidParameter, op, fmt.Sprintf("%s does not have a version field", mDb.Statement.Schema.Table)) + } + where = append(where, fmt.Sprintf("%s.version = ?", mDb.Statement.Schema.Table)) // we need to include the table name because of "on conflict" use cases + args = append(args, opts.WithVersion) + } + if opts.withWhereClause != "" { + where, args = append(where, opts.withWhereClause), append(args, opts.withWhereClauseArgs...) + } + return strings.Join(where, " and "), args, nil +} + // CreateItems will create multiple items of the same type. Supported options: -// WithOplog and WithOplogMsgs. WithOplog and WithOplogMsgs may not be used +// WithDebug, WithOplog, WithOplogMsgs, WithReturnRowsAffected, OnConflict, +// WithVersion, and WithWhere WithOplog and WithOplogMsgs may not be used // together. WithLookup is not a supported option. func (rw *Db) CreateItems(ctx context.Context, createItems []interface{}, opt ...Option) error { const op = "db.CreateItems" @@ -433,7 +481,13 @@ func (rw *Db) CreateItems(ctx context.Context, createItems []interface{}, opt .. } } for _, item := range createItems { - if err := rw.Create(ctx, item, WithOnConflict(opts.withOnConflict), WithReturnRowsAffected(opts.withRowsAffected)); err != nil { + if err := rw.Create(ctx, item, + WithOnConflict(opts.withOnConflict), + WithReturnRowsAffected(opts.withRowsAffected), + WithDebug(opts.withDebug), + WithVersion(opts.WithVersion), + WithWhere(opts.withWhereClause, opts.withWhereClauseArgs...), + ); err != nil { return errors.Wrap(ctx, err, op, errors.WithoutEvent()) } } @@ -552,27 +606,11 @@ func (rw *Db) Update(ctx context.Context, i interface{}, fieldMaskPaths []string var underlying *gorm.DB switch { case opts.WithVersion != nil || opts.withWhereClause != "": - var where []string - var args []interface{} - if opts.WithVersion != nil { - if *opts.WithVersion == 0 { - return NoRowsAffected, errors.New(ctx, errors.InvalidParameter, op, "with version option is zero") - } - mDb := rw.underlying.Model(i) - err = mDb.Statement.Parse(i) - if err != nil && mDb.Statement.Schema == nil { - return NoRowsAffected, errors.New(ctx, errors.Unknown, op, "internal error: unable to parse stmt", errors.WithWrap(err)) - } - if !contains(mDb.Statement.Schema.DBNames, "version") { - // if _, ok := stmt.Schema.FieldsByName["version"]; !ok { - return NoRowsAffected, errors.New(ctx, errors.InvalidParameter, op, fmt.Sprintf("%s does not have a version field", mDb.Statement.Schema.Table)) - } - where, args = append(where, "version = ?"), append(args, opts.WithVersion) + where, args, err := rw.whereClausesFromOpts(ctx, i, opts) + if err != nil { + return NoRowsAffected, errors.Wrap(ctx, err, op) } - if opts.withWhereClause != "" { - where, args = append(where, opts.withWhereClause), append(args, opts.withWhereClauseArgs...) - } - underlying = rw.underlying.Model(i).Where(strings.Join(where, " and "), args...).Updates(updateFields) + underlying = rw.underlying.Model(i).Where(where, args...).Updates(updateFields) default: underlying = rw.underlying.Model(i).Updates(updateFields) } diff --git a/internal/db/read_writer_ext_test.go b/internal/db/read_writer_ext_test.go index 6b6a9e340..628c76583 100644 --- a/internal/db/read_writer_ext_test.go +++ b/internal/db/read_writer_ext_test.go @@ -21,25 +21,28 @@ func TestDb_Create_OnConflict(t *testing.T) { rw := db.New(conn) db.TestCreateTables(t, conn) - // create initial user for on conflict tests - id, err := db.NewPublicId("test-user") - require.NoError(t, err) - initialUser, err := db_test.NewTestUser() - require.NoError(t, err) - ts := ×tamp.Timestamp{Timestamp: timestamppb.Now()} - initialUser.CreateTime = ts - initialUser.UpdateTime = ts - initialUser.Name = "foo-" + id - err = rw.Create(ctx, initialUser) - require.NoError(t, err) - assert.NotEmpty(t, initialUser.Id) - + createInitialUser := func() *db_test.TestUser { + // create initial user for on conflict tests + id, err := db.NewPublicId("test-user") + require.NoError(t, err) + initialUser, err := db_test.NewTestUser() + require.NoError(t, err) + ts := ×tamp.Timestamp{Timestamp: timestamppb.Now()} + initialUser.CreateTime = ts + initialUser.UpdateTime = ts + initialUser.Name = "foo-" + id + err = rw.Create(ctx, initialUser) + require.NoError(t, err) + assert.NotEmpty(t, initialUser.Id) + return initialUser + } tests := []struct { - name string - onConflict db.OnConflict - wantUpdate bool - wantEmail string - wantOplog bool + name string + onConflict db.OnConflict + additionalOpts []db.Option + wantUpdate bool + wantEmail string + wantOplog bool }{ { name: "set-columns", @@ -109,11 +112,52 @@ func TestDb_Create_OnConflict(t *testing.T) { wantUpdate: true, wantOplog: true, }, + { + name: "set-columns-with-where-success", + onConflict: db.OnConflict{ + Target: db.Columns{"public_id"}, + Action: db.SetColumns([]string{"name"}), + }, + additionalOpts: []db.Option{db.WithWhere("db_test_user.version = ?", 1)}, + wantUpdate: true, + wantOplog: true, + }, + { + name: "set-columns-with-where-fail", + onConflict: db.OnConflict{ + Target: db.Columns{"public_id"}, + Action: db.SetColumns([]string{"name"}), + }, + additionalOpts: []db.Option{db.WithWhere("db_test_user.version = ?", 100000000000)}, + wantUpdate: false, + wantOplog: false, + }, + { + name: "set-columns-with-version-success", + onConflict: db.OnConflict{ + Target: db.Columns{"public_id"}, + Action: db.SetColumns([]string{"name"}), + }, + additionalOpts: []db.Option{db.WithVersion(func() *uint32 { i := uint32(1); return &i }())}, + wantUpdate: true, + wantOplog: true, + }, + { + name: "set-columns-with-version-fail", + onConflict: db.OnConflict{ + Target: db.Columns{"public_id"}, + Action: db.SetColumns([]string{"name"}), + }, + additionalOpts: []db.Option{db.WithWhere("db_test_user.version = ?", 100000000000)}, + wantUpdate: false, + wantOplog: false, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { assert, require := assert.New(t), require.New(t) + initialUser := createInitialUser() conflictUser, err := db_test.NewTestUser() require.NoError(err) userNameId, err := db.NewPublicId("test-user-name") @@ -125,7 +169,11 @@ func TestDb_Create_OnConflict(t *testing.T) { "op-type": []string{oplog.OpType_OP_TYPE_CREATE.String()}, } var rowsAffected int64 - err = rw.Create(ctx, conflictUser, db.WithOnConflict(&tt.onConflict), db.WithOplog(wrapper, md), db.WithReturnRowsAffected(&rowsAffected)) + opts := []db.Option{db.WithOnConflict(&tt.onConflict), db.WithOplog(wrapper, md), db.WithReturnRowsAffected(&rowsAffected)} + if tt.additionalOpts != nil { + opts = append(opts, tt.additionalOpts...) + } + err = rw.Create(ctx, conflictUser, opts...) require.NoError(err) if tt.wantOplog { err = db.TestVerifyOplog(t, rw, conflictUser.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_CREATE), db.WithCreateNotBefore(10*time.Second)) @@ -136,6 +184,7 @@ func TestDb_Create_OnConflict(t *testing.T) { foundUser.PublicId = conflictUser.PublicId err = rw.LookupByPublicId(context.Background(), foundUser) require.NoError(err) + t.Log(foundUser) if tt.wantUpdate { assert.Equal(int64(1), rowsAffected) assert.Equal(conflictUser.Id, foundUser.Id) @@ -152,6 +201,7 @@ func TestDb_Create_OnConflict(t *testing.T) { } t.Run("CreateItems", func(t *testing.T) { assert, require := assert.New(t), require.New(t) + initialUser := createInitialUser() conflictUser, err := db_test.NewTestUser() require.NoError(err) userNameId, err := db.NewPublicId("test-user-name") diff --git a/internal/db/read_writer_test.go b/internal/db/read_writer_test.go index a45cfa993..ab8ce6bb7 100644 --- a/internal/db/read_writer_test.go +++ b/internal/db/read_writer_test.go @@ -131,7 +131,7 @@ func TestDb_Update(t *testing.T) { }, want: 0, wantErr: true, - wantErrMsg: "db.Update: with version option is zero: parameter violation: error #100", + wantErrMsg: "with version option is zero: parameter violation: error #100", }, { name: "simple-with-version",