add Handle interface

2 files changed, 22 insertions, 4 deletions

diff --git a/oblast.go b/oblast.go
index 9ed60a4..be1571e 100644
--- a/oblast.go
+++ b/oblast.go
@@ -40,7 +40,11 @@
 //	}
 package oblast // import "go.xyrillian.de/oblast"
 
-import "go.xyrillian.de/oblast/internal"
+import (
+	"database/sql"
+
+	"go.xyrillian.de/oblast/internal"
+)
 
 // PlanOption is an option that can be given to NewStore() to influence query planning for a certain type of record.
 type PlanOption func(*internal.PlanOpts)
@@ -56,3 +60,18 @@ func TableNameIs(name string) PlanOption {
 func PrimaryKeyIs(columnNames ...string) PlanOption {
 	return func(opts *internal.PlanOpts) { opts.PrimaryKeyColumnNames = columnNames }
 }
+
+// Handle is an interface for functions providing direct DB access.
+// It covers methods provided by both *sql.DB and *sql.Tx, thus allowing functions using it to be used both within and outside of transactions.
+type Handle interface {
+	Exec(query string, args ...any) (sql.Result, error)
+	Prepare(query string) (*sql.Stmt, error)
+	Query(query string, args ...any) (*sql.Rows, error)
+	QueryRow(query string, args ...any) *sql.Row
+}
+
+// static assertion that the respective types implement the interface
+var (
+	_ Handle = &sql.DB{}
+	_ Handle = &sql.Tx{}
+)
diff --git a/query.go b/query.go
index e47feb6..fd80f56 100644
--- a/query.go
+++ b/query.go
@@ -11,8 +11,7 @@ import (
 	"go.xyrillian.de/oblast/internal"
 )
 
-// TODO: allow taking *sql.Tx in addition to *sql.DB
-func (s Store[R]) Select(db *sql.DB, query string, args ...any) (result []R, returnedError error) {
+func (s Store[R]) Select(db Handle, query string, args ...any) (result []R, returnedError error) {
 	// NOTE: This function body should be as short as possible to reduce the binary size after monomorphization.
 	//       Any expression that does not depend on type R should be factored out into a reusable function.
 
@@ -37,7 +36,7 @@ func (s Store[R]) Select(db *sql.DB, query string, args ...any) (result []R, ret
 	return result, nil
 }
 
-func startQuery(db *sql.DB, plan internal.Plan, query string, args ...any) (rows *sql.Rows, indexes [][]int, err error) {
+func startQuery(db Handle, plan internal.Plan, query string, args ...any) (rows *sql.Rows, indexes [][]int, err error) {
 	rows, err = db.Query(query, args...)
 	if err != nil {
 		return nil, nil, fmt.Errorf("during Query(): %w", err)