From bbe2653bc51361e5d7607729356344ef979a9f5a Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 23 Sep 2023 08:40:06 -0500
Subject: [PATCH] Prepare chooses statement name based on sql if name == sql

This makes it easier to explicitly manage prepared statements.

refs #1716
---
 conn.go      | 43 ++++++++++++++++++++++++++++++++-----------
 conn_test.go | 22 ++++++++++++++++++++++
 2 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/conn.go b/conn.go
index b867acfe6..489f32fda 100644
--- a/conn.go
+++ b/conn.go
@@ -2,6 +2,8 @@ package pgx
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"strconv"
@@ -284,12 +286,15 @@ func (c *Conn) Close(ctx context.Context) error {
 	return err
 }
 
-// Prepare creates a prepared statement with name and sql. sql can contain placeholders
-// for bound parameters. These placeholders are referenced positional as $1, $2, etc.
+// Prepare creates a prepared statement with name and sql. sql can contain placeholders for bound parameters. These
+// placeholders are referenced positionally as $1, $2, etc. name can be used instead of sql with Query, QueryRow, and
+// Exec to execute the statement. It can also be used with Batch.Queue.
 //
-// Prepare is idempotent; i.e. it is safe to call Prepare multiple times with the same
-// name and sql arguments. This allows a code path to Prepare and Query/Exec without
-// concern for if the statement has already been prepared.
+// The underlying PostgreSQL identifier for the prepared statement will be name if name != sql or a digest of sql if
+// name == sql.
+//
+// Prepare is idempotent; i.e. it is safe to call Prepare multiple times with the same name and sql arguments. This
+// allows a code path to Prepare and Query/Exec without concern for if the statement has already been prepared.
 func (c *Conn) Prepare(ctx context.Context, name, sql string) (sd *pgconn.StatementDescription, err error) {
 	if c.prepareTracer != nil {
 		ctx = c.prepareTracer.TracePrepareStart(ctx, c, TracePrepareStartData{Name: name, SQL: sql})
@@ -311,22 +316,38 @@ func (c *Conn) Prepare(ctx context.Context, name, sql string) (sd *pgconn.Statem
 		}()
 	}
 
-	sd, err = c.pgConn.Prepare(ctx, name, sql, nil)
+	var psName, psKey string
+	if name == sql {
+		digest := sha256.Sum256([]byte(sql))
+		psName = "stmt_" + hex.EncodeToString(digest[0:24])
+		psKey = sql
+	} else {
+		psName = name
+		psKey = name
+	}
+
+	sd, err = c.pgConn.Prepare(ctx, psName, sql, nil)
 	if err != nil {
 		return nil, err
 	}
 
-	if name != "" {
-		c.preparedStatements[name] = sd
+	if psKey != "" {
+		c.preparedStatements[psKey] = sd
 	}
 
 	return sd, nil
 }
 
-// Deallocate released a prepared statement
+// Deallocate releases a prepared statement.
 func (c *Conn) Deallocate(ctx context.Context, name string) error {
-	delete(c.preparedStatements, name)
-	_, err := c.pgConn.Exec(ctx, "deallocate "+quoteIdentifier(name)).ReadAll()
+	var psName string
+	if sd, ok := c.preparedStatements[name]; ok {
+		delete(c.preparedStatements, name)
+		psName = sd.Name
+	} else {
+		psName = name
+	}
+	_, err := c.pgConn.Exec(ctx, "deallocate "+quoteIdentifier(psName)).ReadAll()
 	return err
 }
 
diff --git a/conn_test.go b/conn_test.go
index f988ccb31..739b6619e 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -526,6 +526,28 @@ func TestPrepareStatementCacheModes(t *testing.T) {
 	})
 }
 
+func TestPrepareWithDigestedName(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	pgxtest.RunWithQueryExecModes(ctx, t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		sql := "select $1::text"
+		sd, err := conn.Prepare(ctx, sql, sql)
+		require.NoError(t, err)
+		require.Equal(t, "stmt_2510cc7db17de3f42758a2a29c8b9ef8305d007b997ebdd6", sd.Name)
+
+		var s string
+		err = conn.QueryRow(ctx, sql, "hello").Scan(&s)
+		require.NoError(t, err)
+		require.Equal(t, "hello", s)
+
+		err = conn.Deallocate(ctx, sql)
+		require.NoError(t, err)
+	})
+}
+
 func TestListenNotify(t *testing.T) {
 	t.Parallel()