Merge pull request #7 from randlabs/improvements

Improvements and new features

Author: mxmauro

.gitignore

@@ -23,3 +23,4 @@ app
 demo
 
 vendor/*
+qodana.yml

common_test.go

@@ -0,0 +1,85 @@
+package postgres_test
+
+import (
+	"crypto/rand"
+	"encoding/json"
+	"flag"
+	"testing"
+)
+
+// -----------------------------------------------------------------------------
+
+var (
+	pgUrl          string
+	pgHost         string
+	pgPort         uint
+	pgUsername     string
+	pgPassword     string
+	pgDatabaseName string
+)
+
+var (
+	testJSON      TestJSON
+	testBLOB      []byte
+	testJSONBytes []byte
+)
+
+// -----------------------------------------------------------------------------
+
+func init() {
+	flag.StringVar(&pgUrl, "url", "", "Specifies the Postgres URL.")
+	flag.StringVar(&pgHost, "host", "127.0.0.1", "Specifies the Postgres server host. (Defaults to '127.0.0.1')")
+	flag.UintVar(&pgPort, "port", 5432, "Specifies the Postgres server port. (Defaults to 5432)")
+	flag.StringVar(&pgUsername, "user", "postgres", "Specifies the user name. (Defaults to 'postgres')")
+	flag.StringVar(&pgPassword, "password", "", "Specifies the user password.")
+	flag.StringVar(&pgDatabaseName, "db", "", "Specifies the database name.")
+
+	testJSON = TestJSON{
+		Id:   1,
+		Text: "demo",
+	}
+
+	testBLOB = make([]byte, 1024)
+	_, _ = rand.Read(testBLOB)
+
+	testJSONBytes, _ = json.Marshal(testJSON)
+}
+
+// -----------------------------------------------------------------------------
+
+func checkSettings(t *testing.T) {
+	if len(pgHost) == 0 {
+		t.Fatalf("Server host not specified")
+	}
+	if pgPort > 65535 {
+		t.Fatalf("Server port not specified or invalid")
+	}
+	if len(pgUsername) == 0 {
+		t.Fatalf("User name to access database server not specified")
+	}
+	if len(pgPassword) == 0 {
+		t.Fatalf("User password to access database server not specified")
+	}
+	if len(pgDatabaseName) == 0 {
+		t.Fatalf("Database name not specified")
+	}
+}
+
+func addressOf[T any](x T) *T {
+	return &x
+}
+
+func jsonReEncode(src string) (string, error) {
+	var v interface{}
+
+	err := json.Unmarshal([]byte(src), &v)
+	if err == nil {
+		var reencoded []byte
+
+		reencoded, err = json.Marshal(v)
+		if err == nil {
+			return string(reencoded), nil
+		}
+	}
+	return "", err
+}

connection.go

@@ -0,0 +1,95 @@
+package postgres
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// -----------------------------------------------------------------------------
+
+// Conn encloses a single connection object.
+type Conn struct {
+	db   *Database
+	conn *pgxpool.Conn
+}
+
+// -----------------------------------------------------------------------------
+
+// DB returns the underlying database driver.
+func (c *Conn) DB() *Database {
+	return c.db
+}
+
+// Exec executes an SQL statement within the single connection.
+func (c *Conn) Exec(ctx context.Context, sql string, args ...interface{}) (int64, error) {
+	affectedRows := int64(0)
+	ct, err := c.conn.Exec(ctx, sql, args...)
+	if err == nil {
+		affectedRows = ct.RowsAffected()
+	}
+	return affectedRows, c.db.processError(err)
+}
+
+// QueryRow executes a SQL query within the single connection.
+func (c *Conn) QueryRow(ctx context.Context, sql string, args ...interface{}) Row {
+	return &rowGetter{
+		db:  c.db,
+		row: c.conn.QueryRow(ctx, sql, args...),
+	}
+}
+
+// QueryRows executes a SQL query within the single connection.
+func (c *Conn) QueryRows(ctx context.Context, sql string, args ...interface{}) Rows {
+	rows, err := c.conn.Query(ctx, sql, args...)
+	return &rowsGetter{
+		db:   c.db,
+		ctx:  ctx,
+		rows: rows,
+		err:  err,
+	}
+}
+
+// Copy executes a SQL copy query within the single connection.
+func (c *Conn) Copy(ctx context.Context, tableName string, columnNames []string, callback CopyCallback) (int64, error) {
+	n, err := c.conn.CopyFrom(
+		ctx,
+		pgx.Identifier{tableName},
+		columnNames,
+		&copyWithCallback{
+			ctx:      ctx,
+			callback: callback,
+		},
+	)
+
+	// Done
+	return n, c.db.processError(err)
+}
+
+// WithinTx executes a callback function within the context of a single connection.
+func (c *Conn) WithinTx(ctx context.Context, cb WithinTxCallback) error {
+	innerTx, err := c.conn.BeginTx(ctx, pgx.TxOptions{
+		IsoLevel:       pgx.ReadCommitted, //pgx.Serializable,
+		AccessMode:     pgx.ReadWrite,
+		DeferrableMode: pgx.NotDeferrable,
+	})
+	if err == nil {
+		err = cb(ctx, Tx{
+			db: c.db,
+			tx: innerTx,
+		})
+		if err == nil {
+			err = innerTx.Commit(ctx)
+			if err != nil {
+				err = newError(err, "unable to commit db transaction")
+			}
+		}
+		if err != nil {
+			_ = innerTx.Rollback(context.Background()) // Using context.Background() on purpose
+		}
+	} else {
+		err = newError(err, "unable to start transaction")
+	}
+	return c.db.processError(err)
+}

helpers.go

@@ -48,3 +48,7 @@ func newError(wrappedErr error, text string) *Error {
 func encodeDSN(s string) string {
 	return strings.ReplaceAll(s, "'", "\\'")
 }
+
+func quoteIdentifier(s string) string {
+	return "\"" + strings.ReplaceAll(s, "\"", "\"\"") + "\""
+}

internal.go

@@ -1,7 +1,6 @@
 package postgres
 
 import (
-	"context"
 	"errors"
 
 	"github.com/jackc/pgx/v5"
@@ -13,31 +12,18 @@ var errNoRows = &NoRowsError{}
 
 // -----------------------------------------------------------------------------
 
-// Gets a connection from the pool and initiates a transaction.
-func (db *Database) getTx(ctx context.Context) (pgx.Tx, error) {
-	tx, err := db.pool.BeginTx(ctx, pgx.TxOptions{
-		IsoLevel:       pgx.ReadCommitted, //pgx.Serializable,
-		AccessMode:     pgx.ReadWrite,
-		DeferrableMode: pgx.NotDeferrable,
-	})
-	if err != nil {
-		return nil, newError(err, "unable to start transaction")
-	}
-
-	//Done
-	return tx, nil
-}
-
 func (db *Database) processError(err error) error {
+	isNoRows := false
 	if errors.Is(err, pgx.ErrNoRows) {
 		err = errNoRows
+		isNoRows = true
 	}
 
 	// Only deal with fatal database errors. Cancellation, timeouts and empty result sets are not considered fatal.
 	db.err.mutex.Lock()
 	defer db.err.mutex.Unlock()
 
-	if err != nil && IsDatabaseError(err) && err != errNoRows {
+	if err != nil && (!isNoRows) && IsDatabaseError(err) {
 		if db.err.last == nil {
 			db.err.last = err
 			if db.err.handler != nil {