Skip to content

Commit

Permalink
Add support Upsert to model query
Browse files Browse the repository at this point in the history
  • Loading branch information
kolkov committed Jan 17, 2021
1 parent 4d1628e commit 5942e1e
Show file tree
Hide file tree
Showing 3 changed files with 137 additions and 2 deletions.
13 changes: 11 additions & 2 deletions builder_mysql.go
Expand Up @@ -73,8 +73,17 @@ func (b *MysqlBuilder) Upsert(table string, cols Params, constraints ...string)

names := []string{}
for name := range cols {
names = append(names, name)
found := false
for _, pkName := range constraints {
if pkName == name {
found = true
}
}
if !found {
names = append(names, name)
}
}

sort.Strings(names)

lines := []string{}
Expand All @@ -91,7 +100,7 @@ func (b *MysqlBuilder) Upsert(table string, cols Params, constraints ...string)

q.sql += " ON DUPLICATE KEY UPDATE " + strings.Join(lines, ", ")

return q
return b.NewQuery(q.sql).Bind(q.params)
}

var mysqlColumnRegexp = regexp.MustCompile("(?m)^\\s*[`\"](.*?)[`\"]\\s+(.*?),?$")
Expand Down
24 changes: 24 additions & 0 deletions model_query.go
Expand Up @@ -172,3 +172,27 @@ func (q *ModelQuery) Delete() error {
_, err := q.builder.Delete(q.model.tableName, HashExp(pk)).WithContext(q.ctx).Execute()
return err
}

// Upsert creates a Query that represents an UPSERT SQL statement.
// Upsert inserts a row into the table if the primary key or unique index is not found.
// Otherwise it will update the row with the new values.
// The keys of cols are the column names, while the values of cols are the corresponding column
// values to be inserted.
func (q *ModelQuery) Upsert(attrs ...string) error {
if q.lastError != nil {
return q.lastError
}
pk := q.model.pk()
if len(pk) == 0 {
return MissingPKError
}
var pks []string

cols := q.model.columns(attrs, q.exclude)
for name := range pk {
cols[name] = pk[name]
pks = append(pks, name)
}
_, err := q.builder.Upsert(q.model.tableName, Params(cols), pks...).WithContext(q.ctx).Execute()
return err
}
102 changes: 102 additions & 0 deletions model_query_test.go
Expand Up @@ -236,3 +236,105 @@ func TestModelQuery_Delete(t *testing.T) {
var a int
assert.NotNil(t, db.Model(&a).Delete())
}

func TestModelQuery_Upsert(t *testing.T) {
db := getPreparedDB()
defer db.Close()

id := 2
name := "test"
email := "test@example.com"
{
// updating normally
customer := Customer{
ID: id,
Name: name,
Email: email,
}
err := db.Model(&customer).Upsert()
if assert.Nil(t, err) {
var c Customer
db.Select().From("customer").Where(HashExp{"ID": id}).One(&c)
assert.Equal(t, name, c.Name)
assert.Equal(t, email, c.Email)
assert.Equal(t, 0, c.Status)
}
}

{
// updating without primary keys
item2 := Item{
Name: name,
}
err := db.Model(&item2).Upsert()
assert.Equal(t, MissingPKError, err)
}

{
// updating all fields
customer := CustomerPtr{
ID: &id,
Name: name,
Email: &email,
}
err := db.Model(&customer).Upsert()
if assert.Nil(t, err) {
assert.Equal(t, id, *customer.ID)
var c CustomerPtr
db.Select().From("customer").Where(HashExp{"ID": id}).One(&c)
assert.Equal(t, name, c.Name)
if assert.NotNil(t, c.Email) {
assert.Equal(t, email, *c.Email)
}
assert.Nil(t, c.Status)
}
}

{
// updating selected fields only
id = 3
customer := CustomerPtr{
ID: &id,
Name: name,
Email: &email,
}
err := db.Model(&customer).Upsert("Name", "Email")
if assert.Nil(t, err) {
assert.Equal(t, id, *customer.ID)
var c CustomerPtr
db.Select().From("customer").Where(HashExp{"ID": id}).One(&c)
assert.Equal(t, name, c.Name)
if assert.NotNil(t, c.Email) {
assert.Equal(t, email, *c.Email)
}
if assert.NotNil(t, c.Status) {
assert.Equal(t, 2, *c.Status)
}
}
}

{
// inserting normally
customer := Customer{
ID: 5,
Name: name,
Email: email,
}
err := db.Model(&customer).Upsert()
if assert.Nil(t, err) {
assert.Equal(t, 5, customer.ID)
var c Customer
db.Select().From("customer").Where(HashExp{"ID": 5}).One(&c)
assert.Equal(t, name, c.Name)
assert.Equal(t, email, c.Email)
assert.Equal(t, 0, c.Status)
assert.False(t, c.Address.Valid)
}
}

{
// updating non-struct
var a int
assert.NotNil(t, db.Model(&a).Upsert())
}
}

0 comments on commit 5942e1e

Please sign in to comment.