From cfb45a9e3db2df2bcc06857eadd48311654ec99a Mon Sep 17 00:00:00 2001
From: Yonas Habteab <yonas.habteab@icinga.com>
Date: Thu, 19 Oct 2023 15:33:44 +0200
Subject: [PATCH] Introduce `DB#InsertObtainID()` method

---
 database/contracts.go | 13 +++++++++++++
 database/utils.go     | 45 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/database/contracts.go b/database/contracts.go
index bf55d320..097ec03c 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,11 @@ 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
+	sqlx.PreparerContext
+
+	PrepareNamedContext(ctx context.Context, query string) (*sqlx.NamedStmt, error)
+}
diff --git a/database/utils.go b/database/utils.go
index 2ae372ce..6391e320 100644
--- a/database/utils.go
+++ b/database/utils.go
@@ -2,11 +2,13 @@ package database
 
 import (
 	"context"
+	"database/sql"
 	"database/sql/driver"
 	"github.com/go-sql-driver/mysql"
 	"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"
 )
 
@@ -43,6 +45,49 @@ 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 wrapped in [types.Int] on success and error on any database inserting/retrieving failure.
+func InsertObtainID(ctx context.Context, conn TxOrDB, stmt string, arg any) (types.Int, error) {
+	var resultID int64
+	switch conn.DriverName() {
+	case PostgreSQL:
+		query := stmt + " RETURNING id"
+		ps, err := conn.PrepareNamedContext(ctx, query)
+		if err != nil {
+			return types.Int{}, errors.Wrapf(err, "cannot prepare %q", query)
+		}
+		// We're deferring the ps#Close invocation here just to be on the safe side, otherwise it's
+		// closed manually later on and the error is handled gracefully (if any).
+		defer func() { _ = ps.Close() }()
+
+		if err = ps.GetContext(ctx, &resultID, arg); err != nil {
+			return types.Int{}, CantPerformQuery(err, query)
+		}
+
+		if err = ps.Close(); err != nil {
+			return types.Int{}, errors.Wrapf(err, "cannot close prepared statement %q", query)
+		}
+	case MySQL:
+		result, err := sqlx.NamedExecContext(ctx, conn, stmt, arg)
+		if err != nil {
+			return types.Int{}, CantPerformQuery(err, stmt)
+		}
+
+		resultID, err = result.LastInsertId()
+		if err != nil {
+			return types.Int{}, errors.Wrap(err, "cannot retrieve last inserted ID")
+		}
+	default:
+		return types.Int{}, errors.Errorf("unsupported driver: %s", conn.DriverName())
+	}
+
+	return types.Int{NullInt64: sql.NullInt64{Int64: resultID, Valid: true}}, 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,