Skip to content

Commit

Permalink
Introduce DB#InsertObtainID() function
Browse files Browse the repository at this point in the history
  • Loading branch information
yhabteab committed Aug 14, 2024
1 parent 715b56e commit 2045ba4
Show file tree
Hide file tree
Showing 3 changed files with 55 additions and 0 deletions.
12 changes: 12 additions & 0 deletions database/contracts.go
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
}
6 changes: 6 additions & 0 deletions database/db.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
)
37 changes: 37 additions & 0 deletions database/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 2045ba4

Please sign in to comment.