You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
67 lines
1.3 KiB
67 lines
1.3 KiB
package persist |
|
|
|
import ( |
|
"context" |
|
"sig-pub/pkg/utils/exit" |
|
|
|
"github.com/jackc/pgx/v5" |
|
"github.com/jackc/pgx/v5/pgxpool" |
|
) |
|
|
|
type PGBatchWriter struct { |
|
pool *pgxpool.Pool |
|
} |
|
|
|
func NewPGBatchWriter() *PGBatchWriter { |
|
return &PGBatchWriter{} |
|
} |
|
|
|
func (w *PGBatchWriter) Init(ctx context.Context, connString string) (err error) { |
|
w.pool, err = pgxpool.New(ctx, connString) |
|
if err != nil { |
|
return |
|
} |
|
exit.AddHook(w.pool.Close, exit.WithOrderTail()) |
|
|
|
return |
|
} |
|
|
|
func InsertBatch[T any](w *PGBatchWriter, ctx context.Context, datas []T) (err error) { |
|
if len(datas) == 0 { |
|
return |
|
} |
|
|
|
tableColumns := make(map[string][]string, 3) |
|
tableValues := make(map[string][][]any, 3) |
|
for _, data := range datas { |
|
table, columns, valuesM, e := ReflectGormData(data) |
|
if e != nil { |
|
err = e |
|
return |
|
} |
|
if cs, ok := tableColumns[table]; !ok { |
|
tableColumns[table] = columns |
|
} else { |
|
columns = cs |
|
} |
|
var values []any |
|
for _, c := range columns { |
|
v := valuesM[c] |
|
values = append(values, v) |
|
} |
|
tableValues[table] = append(tableValues[table], values) |
|
} |
|
tx, err := w.pool.Begin(ctx) |
|
if err != nil { |
|
return |
|
} |
|
defer tx.Rollback(ctx) |
|
for table, values := range tableValues { |
|
_, err = tx.CopyFrom(ctx, pgx.Identifier{table}, tableColumns[table], pgx.CopyFromRows(values)) |
|
if err != nil { |
|
return |
|
} |
|
} |
|
err = tx.Commit(ctx) |
|
return |
|
}
|
|
|