Repository: jackc/pgx
Pull request #2077 (2024-07-07 01:22:56)

Fix bug relating to reuse of types
When `LoadTypes` is being called, it does not include the namespace-qualified types in its result. While these namespaces are visible to `LoadTypes` itself, `RegisterTypes` will not recognise this form of the types, only allowing them to be used if they are on the schema path, and referred to without their namespace component.


Patch:
diff --git a/batch.go b/batch.go
index c3c2834f..3540f57f 100644
--- a/batch.go
+++ b/batch.go
@@ -60,13 +60,9 @@ type Batch struct {
 	QueuedQueries []*QueuedQuery
 }
 
-// Queue queues a query to batch b. query can be an SQL query or the name of a prepared statement. The only pgx option
-// argument that is supported is QueryRewriter. Queries are executed using the connection's DefaultQueryExecMode.
-//
-// While query can contain multiple statements if the connection's DefaultQueryExecMode is QueryModeSimple, this should
-// be avoided. QueuedQuery.Fn must not be set as it will only be called for the first query. That is, QueuedQuery.Query,
-// QueuedQuery.QueryRow, and QueuedQuery.Exec must not be called. In addition, any error messages or tracing that
-// include the current query may reference the wrong query.
+// Queue queues a query to batch b. query can be an SQL query or the name of a prepared statement.
+// The only pgx option argument that is supported is QueryRewriter. Queries are executed using the
+// connection's DefaultQueryExecMode.
 func (b *Batch) Queue(query string, arguments ...any) *QueuedQuery {
 	qq := &QueuedQuery{
 		SQL:       query,
@@ -132,7 +128,7 @@ func (br *batchResults) Exec() (pgconn.CommandTag, error) {
 	if !br.mrr.NextResult() {
 		err := br.mrr.Close()
 		if err == nil {
-			err = errors.New("no more results in batch")
+			err = errors.New("no result")
 		}
 		if br.conn.batchTracer != nil {
 			br.conn.batchTracer.TraceBatchQuery(br.ctx, br.conn, TraceBatchQueryData{
@@ -184,7 +180,7 @@ func (br *batchResults) Query() (Rows, error) {
 	if !br.mrr.NextResult() {
 		rows.err = br.mrr.Close()
 		if rows.err == nil {
-			rows.err = errors.New("no more results in batch")
+			rows.err = errors.New("no result")
 		}
 		rows.closed = true
 
@@ -291,10 +287,7 @@ func (br *pipelineBatchResults) Exec() (pgconn.CommandTag, error) {
 		return pgconn.CommandTag{}, br.err
 	}
 
-	query, arguments, err := br.nextQueryAndArgs()
-	if err != nil {
-		return pgconn.CommandTag{}, err
-	}
+	query, arguments, _ := br.nextQueryAndArgs()
 
 	results, err := br.pipeline.GetResults()
 	if err != nil {
@@ -337,9 +330,9 @@ func (br *pipelineBatchResults) Query() (Rows, error) {
 		return &baseRows{err: br.err, closed: true}, br.err
 	}
 
-	query, arguments, err := br.nextQueryAndArgs()
-	if err != nil {
-		return &baseRows{err: err, closed: true}, err
+	query, arguments, ok := br.nextQueryAndArgs()
+	if !ok {
+		query = "batch query"
 	}
 
 	rows := br.conn.getRows(br.ctx, query, arguments)
@@ -428,16 +421,13 @@ func (br *pipelineBatchResults) earlyError() error {
 	return br.err
 }
 
-func (br *pipelineBatchResults) nextQueryAndArgs() (query string, args []any, err error) {
-	if br.b == nil {
-		return "", nil, errors.New("no reference to batch")
-	}
-
-	if br.qqIdx >= len(br.b.QueuedQueries) {
-		return "", nil, errors.New("no more results in batch")
+func (br *pipelineBatchResults) nextQueryAndArgs() (query string, args []any, ok bool) {
+	if br.b != nil && br.qqIdx < len(br.b.QueuedQueries) {
+		bi := br.b.QueuedQueries[br.qqIdx]
+		query = bi.SQL
+		args = bi.Arguments
+		ok = true
+		br.qqIdx++
 	}
-
-	bi := br.b.QueuedQueries[br.qqIdx]
-	br.qqIdx++
-	return bi.SQL, bi.Arguments, nil
+	return
 }
diff --git a/derived_types.go b/derived_types.go
index 3df45e48..22ab069c 100644
--- a/derived_types.go
+++ b/derived_types.go
@@ -235,7 +235,9 @@ func (c *Conn) LoadTypes(ctx context.Context, typeNames []string) ([]*pgtype.Typ
 		if type_ != nil {
 			m.RegisterType(type_)
 			if ti.NspName != "" {
-				m.RegisterType(&pgtype.Type{Name: ti.NspName + "." + type_.Name, OID: type_.OID, Codec: type_.Codec})
+				nspType := &pgtype.Type{Name: ti.NspName + "." + type_.Name, OID: type_.OID, Codec: type_.Codec}
+				m.RegisterType(nspType)
+				result = append(result, nspType)
 			}
 			result = append(result, type_)
 		}


Test patch:
diff --git a/batch_test.go b/batch_test.go
index eb560e06..9ff2417f 100644
--- a/batch_test.go
+++ b/batch_test.go
@@ -290,45 +290,6 @@ func TestConnSendBatchMany(t *testing.T) {
 	})
 }
 
-// https://github.com/jackc/pgx/issues/1801#issuecomment-2203784178
-func TestConnSendBatchReadResultsWhenNothingQueued(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) {
-		batch := &pgx.Batch{}
-		br := conn.SendBatch(ctx, batch)
-		commandTag, err := br.Exec()
-		require.Equal(t, "", commandTag.String())
-		require.EqualError(t, err, "no more results in batch")
-		err = br.Close()
-		require.NoError(t, err)
-	})
-}
-
-func TestConnSendBatchReadMoreResultsThanQueriesSent(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) {
-		batch := &pgx.Batch{}
-		batch.Queue("select 1")
-		br := conn.SendBatch(ctx, batch)
-		commandTag, err := br.Exec()
-		require.Equal(t, "SELECT 1", commandTag.String())
-		require.NoError(t, err)
-		commandTag, err = br.Exec()
-		require.Equal(t, "", commandTag.String())
-		require.EqualError(t, err, "no more results in batch")
-		err = br.Close()
-		require.NoError(t, err)
-	})
-}
-
 func TestConnSendBatchWithPreparedStatement(t *testing.T) {
 	t.Parallel()
 
diff --git a/derived_types_test.go b/derived_types_test.go
index e6c1c869..6fb6e1d3 100644
--- a/derived_types_test.go
+++ b/derived_types_test.go
@@ -29,9 +29,12 @@ create type dtype_test as (
 
 		types, err := conn.LoadTypes(ctx, []string{"dtype_test"})
 		require.NoError(t, err)
-		require.Len(t, types, 3)
-		require.Equal(t, types[0].Name, "anotheruint64")
-		require.Equal(t, types[1].Name, "_anotheruint64")
-		require.Equal(t, types[2].Name, "dtype_test")
+		require.Len(t, types, 6)
+		require.Equal(t, types[0].Name, "public.anotheruint64")
+		require.Equal(t, types[1].Name, "anotheruint64")
+		require.Equal(t, types[2].Name, "public._anotheruint64")
+		require.Equal(t, types[3].Name, "_anotheruint64")
+		require.Equal(t, types[4].Name, "public.dtype_test")
+		require.Equal(t, types[5].Name, "dtype_test")
 	})
 }
diff --git a/pgtype/derived_types_test.go b/pgtype/derived_types_test.go
index 4f716ea9..05d109bd 100644
--- a/pgtype/derived_types_test.go
+++ b/pgtype/derived_types_test.go
@@ -28,10 +28,13 @@ create type dt_test as (
 		defer conn.Exec(ctx, "drop type dt_test")
 
 		dtypes, err := conn.LoadTypes(ctx, []string{"dt_test"})
-		require.Len(t, dtypes, 3)
-		require.Equal(t, dtypes[0].Name, "dt_uint64")
-		require.Equal(t, dtypes[1].Name, "_dt_uint64")
-		require.Equal(t, dtypes[2].Name, "dt_test")
+		require.Len(t, dtypes, 6)
+		require.Equal(t, dtypes[0].Name, "public.dt_uint64")
+		require.Equal(t, dtypes[1].Name, "dt_uint64")
+		require.Equal(t, dtypes[2].Name, "public._dt_uint64")
+		require.Equal(t, dtypes[3].Name, "_dt_uint64")
+		require.Equal(t, dtypes[4].Name, "public.dt_test")
+		require.Equal(t, dtypes[5].Name, "dt_test")
 		require.NoError(t, err)
 		conn.TypeMap().RegisterTypes(dtypes)