From 01d864f3c3bbb4f2ec299d8b401f8f953cf58106 Mon Sep 17 00:00:00 2001 From: "Roman A. Grigorovich" Date: Tue, 17 Oct 2023 12:58:16 +0300 Subject: [PATCH] Ability to extend upsert expression with options --- .../driver/override/main/17_upsert.go.tpl | 16 +++--- .../main/singleton/psql_upsert.go.tpl | 57 ++++++++++++++----- 2 files changed, 52 insertions(+), 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 index 7e6dca33a..04b552c89 100644 --- 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 index edb1b09a1..5ea9a6127 100644 --- 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) } }