Merge pull request #1317 from atzedus/upsert-options-feature

Ability to extend upsert expression with options
volatiletech · Oct 18, 2023 · 6c91199 · 6c91199
2 parents 63485f5 + 01d864f
commit 6c91199
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 21 deletions.
diff --git a/drivers/sqlboiler-psql/driver/override/main/17_upsert.go.tpl b/drivers/sqlboiler-psql/driver/override/main/17_upsert.go.tpl
@@ -3,16 +3,16 @@
 {{- $schemaTable := .Table.Name | .SchemaTable}}
 {{if .AddGlobal -}}
 // UpsertG attempts an insert, and does an update or ignore on conflict.
-func (o *{{$alias.UpSingular}}) UpsertG({{if not .NoContext}}ctx context.Context, {{end -}} updateOnConflict bool, conflictColumns []string, updateColumns, insertColumns boil.Columns) error {
-	return o.Upsert({{if .NoContext}}boil.GetDB(){{else}}ctx, boil.GetContextDB(){{end}}, updateOnConflict, conflictColumns, updateColumns, insertColumns)
+func (o *{{$alias.UpSingular}}) UpsertG({{if not .NoContext}}ctx context.Context, {{end -}} updateOnConflict bool, conflictColumns []string, updateColumns, insertColumns boil.Columns, opts ...UpsertOptionFunc) error {
+	return o.Upsert({{if .NoContext}}boil.GetDB(){{else}}ctx, boil.GetContextDB(){{end}}, updateOnConflict, conflictColumns, updateColumns, insertColumns, opts...)
 }
 
 {{end -}}
 
 {{if and .AddGlobal .AddPanic -}}
 // UpsertGP attempts an insert, and does an update or ignore on conflict. Panics on error.
-func (o *{{$alias.UpSingular}}) UpsertGP({{if not .NoContext}}ctx context.Context, {{end -}} updateOnConflict bool, conflictColumns []string, updateColumns, insertColumns boil.Columns) {
-	if err := o.Upsert({{if .NoContext}}boil.GetDB(){{else}}ctx, boil.GetContextDB(){{end}}, updateOnConflict, conflictColumns, updateColumns, insertColumns); err != nil {
+func (o *{{$alias.UpSingular}}) UpsertGP({{if not .NoContext}}ctx context.Context, {{end -}} updateOnConflict bool, conflictColumns []string, updateColumns, insertColumns boil.Columns, opts ...UpsertOptionFunc) {
+	if err := o.Upsert({{if .NoContext}}boil.GetDB(){{else}}ctx, boil.GetContextDB(){{end}}, updateOnConflict, conflictColumns, updateColumns, insertColumns, opts...); err != nil {
 		panic(boil.WrapErr(err))
 	}
 }
@@ -22,8 +22,8 @@ func (o *{{$alias.UpSingular}}) UpsertGP({{if not .NoContext}}ctx context.Contex
 {{if .AddPanic -}}
 // UpsertP attempts an insert using an executor, and does an update or ignore on conflict.
 // UpsertP panics on error.
-func (o *{{$alias.UpSingular}}) UpsertP({{if .NoContext}}exec boil.Executor{{else}}ctx context.Context, exec boil.ContextExecutor{{end}}, updateOnConflict bool, conflictColumns []string, updateColumns, insertColumns boil.Columns) {
-	if err := o.Upsert({{if not .NoContext}}ctx, {{end -}} exec, updateOnConflict, conflictColumns, updateColumns, insertColumns); err != nil {
+func (o *{{$alias.UpSingular}}) UpsertP({{if .NoContext}}exec boil.Executor{{else}}ctx context.Context, exec boil.ContextExecutor{{end}}, updateOnConflict bool, conflictColumns []string, updateColumns, insertColumns boil.Columns, opts ...UpsertOptionFunc) {
+	if err := o.Upsert({{if not .NoContext}}ctx, {{end -}} exec, updateOnConflict, conflictColumns, updateColumns, insertColumns, opts...); err != nil {
 		panic(boil.WrapErr(err))
 	}
 }
@@ -32,7 +32,7 @@ func (o *{{$alias.UpSingular}}) UpsertP({{if .NoContext}}exec boil.Executor{{els
 
 // Upsert attempts an insert using an executor, and does an update or ignore on conflict.
 // See boil.Columns documentation for how to properly use updateColumns and insertColumns.
-func (o *{{$alias.UpSingular}}) Upsert({{if .NoContext}}exec boil.Executor{{else}}ctx context.Context, exec boil.ContextExecutor{{end}}, updateOnConflict bool, conflictColumns []string, updateColumns, insertColumns boil.Columns) error {
+func (o *{{$alias.UpSingular}}) Upsert({{if .NoContext}}exec boil.Executor{{else}}ctx context.Context, exec boil.ContextExecutor{{end}}, updateOnConflict bool, conflictColumns []string, updateColumns, insertColumns boil.Columns, opts ...UpsertOptionFunc) error {
 	if o == nil {
 		return errors.New("{{.PkgName}}: no {{.Table.Name}} provided for upsert")
 	}
@@ -111,7 +111,7 @@ func (o *{{$alias.UpSingular}}) Upsert({{if .NoContext}}exec boil.Executor{{else
 			conflict = make([]string, len({{$alias.DownSingular}}PrimaryKeyColumns))
 			copy(conflict, {{$alias.DownSingular}}PrimaryKeyColumns)
 		}
-		cache.query = buildUpsertQueryPostgres(dialect, "{{$schemaTable}}", updateOnConflict, ret, update, conflict, insert)
+		cache.query = buildUpsertQueryPostgres(dialect, "{{$schemaTable}}", updateOnConflict, ret, update, conflict, insert, opts...)
 
 		cache.valueMapping, err = queries.BindMapping({{$alias.DownSingular}}Type, {{$alias.DownSingular}}Mapping, insert)
 		if err != nil {

diff --git a/drivers/sqlboiler-psql/driver/override/main/singleton/psql_upsert.go.tpl b/drivers/sqlboiler-psql/driver/override/main/singleton/psql_upsert.go.tpl
@@ -1,9 +1,33 @@
+type UpsertOptions struct {
+	conflictObject string
+	updateSet string
+}
+
+type UpsertOptionFunc func(o *UpsertOptions)
+
+func UpsertConflictObject(conflictObject string) UpsertOptionFunc {
+	return func(o *UpsertOptions) {
+		o.conflictObject = conflictObject
+	}
+}
+
+func UpsertUpdateSet(updateSet string) UpsertOptionFunc {
+	return func(o *UpsertOptions) {
+		o.updateSet = updateSet
+	}
+}
+
 // buildUpsertQueryPostgres builds a SQL statement string using the upsertData provided.
-func buildUpsertQueryPostgres(dia drivers.Dialect, tableName string, updateOnConflict bool, ret, update, conflict, whitelist []string) string {
+func buildUpsertQueryPostgres(dia drivers.Dialect, tableName string, updateOnConflict bool, ret, update, conflict, whitelist []string, opts ...UpsertOptionFunc) string {
 	conflict = strmangle.IdentQuoteSlice(dia.LQ, dia.RQ, conflict)
 	whitelist = strmangle.IdentQuoteSlice(dia.LQ, dia.RQ, whitelist)
 	ret = strmangle.IdentQuoteSlice(dia.LQ, dia.RQ, ret)
 
+	upsertOpts := &UpsertOptions{}
+	for _, o := range opts {
+		o(upsertOpts)
+	}
+
 	buf := strmangle.GetBuffer()
 	defer strmangle.PutBuffer(buf)
 
@@ -21,28 +45,35 @@ func buildUpsertQueryPostgres(dia drivers.Dialect, tableName string, updateOnCon
 		columns,
 	)
 
-	if len(conflict) != 0 {
+	if upsertOpts.conflictObject != "" {
+		buf.WriteString(upsertOpts.conflictObject)
+	} else if len(conflict) != 0 {
 		buf.WriteByte('(')
 		buf.WriteString(strings.Join(conflict, ", "))
-		buf.WriteString(") ")
+		buf.WriteByte(')')
 	}
+	buf.WriteByte(' ')
 
 	if !updateOnConflict || len(update) == 0 {
 		buf.WriteString("DO NOTHING")
 	} else {
 		buf.WriteString("DO UPDATE SET ")
 
-		for i, v := range update {
-		    if len(v) == 0 {
-		        continue
-		    }
-			if i != 0 {
-				buf.WriteByte(',')
+		if upsertOpts.updateSet != "" {
+			buf.WriteString(upsertOpts.updateSet)
+		} else {
+			for i, v := range update {
+				if len(v) == 0 {
+					continue
+				}
+				if i != 0 {
+					buf.WriteByte(',')
+				}
+				quoted := strmangle.IdentQuote(dia.LQ, dia.RQ, v)
+				buf.WriteString(quoted)
+				buf.WriteString(" = EXCLUDED.")
+				buf.WriteString(quoted)
 			}
-			quoted := strmangle.IdentQuote(dia.LQ, dia.RQ, v)
-			buf.WriteString(quoted)
-			buf.WriteString(" = EXCLUDED.")
-			buf.WriteString(quoted)
 		}
 	}