diff --git a/database/contracts.go b/database/contracts.go index bf55d320..6da723a8 100644 --- a/database/contracts.go +++ b/database/contracts.go @@ -1,5 +1,10 @@ package database +import ( + "context" + "github.com/jmoiron/sqlx" +) + // Entity is implemented by each type that works with the database package. type Entity interface { Fingerprinter @@ -54,3 +59,10 @@ type PgsqlOnConflictConstrainter interface { // PgsqlOnConflictConstraint returns the primary or unique key constraint name of the PostgreSQL table. PgsqlOnConflictConstraint() string } + +// TxOrDB is just a helper interface that can represent a *[sqlx.Tx] or *[DB] instance. +type TxOrDB interface { + sqlx.ExtContext + + PrepareNamedContext(ctx context.Context, query string) (*sqlx.NamedStmt, error) +} diff --git a/database/db.go b/database/db.go index d024e421..e623d062 100644 --- a/database/db.go +++ b/database/db.go @@ -804,3 +804,9 @@ func (db *DB) Log(ctx context.Context, query string, counter *com.Counter) perio db.logger.Debugf("Finished executing %q with %d rows in %s", query, counter.Total(), tick.Elapsed) })) } + +var ( + // Assert TxOrDB interface compliance of the DB and sqlx.Tx types. + _ TxOrDB = (*DB)(nil) + _ TxOrDB = (*sqlx.Tx)(nil) +) diff --git a/database/utils.go b/database/utils.go index 55c9264d..aa1f4bc9 100644 --- a/database/utils.go +++ b/database/utils.go @@ -7,6 +7,7 @@ import ( "github.com/icinga/icinga-go-library/com" "github.com/icinga/icinga-go-library/strcase" "github.com/icinga/icinga-go-library/types" + "github.com/jmoiron/sqlx" "github.com/pkg/errors" "strconv" ) @@ -44,6 +45,42 @@ func SplitOnDupId[T IDer]() com.BulkChunkSplitPolicy[T] { } } +// InsertObtainID executes the given query and fetches the last inserted ID. +// +// Using this method for database tables that don't define an auto-incrementing ID, or none at all, +// will not work. The only supported column that can be retrieved with this method is id. +// +// This function expects [TxOrDB] as an executor of the provided query, and is usually a *[sqlx.Tx] or *[DB] instance. +// +// Returns the retrieved ID on success and error on any database inserting/retrieving failure. +func InsertObtainID(ctx context.Context, conn TxOrDB, stmt string, arg any) (int64, error) { + var resultID int64 + switch conn.DriverName() { + case PostgreSQL: + stmt = stmt + " RETURNING id" + query, args, err := conn.BindNamed(stmt, arg) + if err != nil { + return 0, errors.Wrapf(err, "cannot bind named query %q", stmt) + } + + if err := sqlx.GetContext(ctx, conn, &resultID, query, args...); err != nil { + return 0, CantPerformQuery(err, query) + } + default: + result, err := sqlx.NamedExecContext(ctx, conn, stmt, arg) + if err != nil { + return 0, CantPerformQuery(err, stmt) + } + + resultID, err = result.LastInsertId() + if err != nil { + return 0, errors.Wrap(err, "cannot retrieve last inserted ID") + } + } + + return resultID, nil +} + // setGaleraOpts sets the "wsrep_sync_wait" variable for each session ensures that causality checks are performed // before execution and that each statement is executed on a fully synchronized node. Doing so prevents foreign key // violation when inserting into dependent tables on different MariaDB/MySQL nodes. When using MySQL single nodes,