Using RowToStructByName does not seem to work ByName with embedded structs

[svennis94]

Hej!

First thank you for the great work on PGx! Very nice library and it is easy to switch from sqlx to pgx.

In our project we are using array_agg to fill our embedded structs. There is one thing we are running into when using the pgx.CollectRows() with pgx.RowToStructByName with embedded structs. It seems that the RowToStructByName strategy is not used for the embedded struct, only for the parent struct, as when the order of the columns in the array_agg does not match the order of the struct fields of the embedded struct there will be a scan error if the types don't match. If the types match there is no issue, but you'll have data in the wrong fields of the Go struct.

It might be that we are doing something wrong here, but it seems that when it comes to the array_agg embedded structs it is using the RowToStructByPos strategy rather than the RowToStructByName. We can work around this by specifying the columns in the query in the same order as the Go structs, but the point of using the RowToStructByName strategy was to try to avoid this 🙂

I have created a testcase for it, it seems that because the highlighted column is at the end of the of the table definition it is not able to map the aggregated array data into the []Tag embedded struct. Since in the []Tag embedded struct the Highlighted field is after the Name field. When specifying the order of the columns in the array_agg method the query succeeds.

The test case is best run on an empty DB. It will create the DB schema and insert mock data to run the selects against.

Is there something we are doing wrong as the RowToStructByName doesn't seem to map by name for the []Tag struct? Or might this be something that could be implemented in pgx? (I could write a PR if that is needed).

Thanks in advance!

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Our base post table
var postsTable = `
	CREATE TABLE IF NOT EXISTS posts (
		id SERIAL PRIMARY KEY,
		title TEXT,
		body TEXT,
		created_at TIMESTAMPTZ,
		updated_at TIMESTAMPTZ
	);
`

// Our base tags table
var tagsTable = `
	CREATE TABLE IF NOT EXISTS tags (
		id SERIAL PRIMARY KEY,
		name TEXT,
		created_at TIMESTAMPTZ,
		updated_at TIMESTAMPTZ
	);
`

// Our base pivot table
var pivotTable = `
	CREATE TABLE IF NOT EXISTS posts_tags (
		post_id INTEGER NOT NULL,
		tag_id INTEGER NOT NULL,
		PRIMARY KEY (post_id, tag_id)
	);
`

// Append extra column to tags table
var alterTags = `
	ALTER TABLE tags ADD COLUMN IF NOT EXISTS highlighted BOOL;
`

type Post struct {
	ID        int
	Title     string
	Body      string
	Tags      []Tag
	CreatedAt *time.Time
	UpdatedAt *time.Time
}

type Tag struct {
	ID          int
	Name        string
	Highlighted bool // Added later to our struct but not in the same order as the table
	CreatedAt   *time.Time
	UpdatedAt   *time.Time
}

func main() {
	ctx := context.Background()

	// Connect to DB
	host := os.Getenv("PG_HOST")
	port := os.Getenv("PG_PORT")
	user := os.Getenv("PG_USER")
	password := os.Getenv("PG_PASSWORD")
	dbName := os.Getenv("PG_DATABASE")

	connStr := fmt.Sprintf("host=%s port=%s user=%s dbname=%s password=%s", host, port, user, dbName, password)

	config, err := pgxpool.ParseConfig(connStr)
	if err != nil {
		log.Fatal(err)
	}

	db, err := pgxpool.NewWithConfig(ctx, config)
	if err != nil {
		log.Fatal(err)
	}

	// We build up our table structure and append a column to the tags table in a later version
	// of our code base
	dbStruct := []string{
		postsTable,
		tagsTable,
		pivotTable,
		alterTags,
	}

	for _, mig := range dbStruct {
		_, err = db.Exec(ctx, mig)
		if err != nil {
			log.Fatal(err)
		}
	}

	mockData := []string{
		`INSERT INTO posts (title, body, created_at, updated_at) VALUES ('Test post 1', 'Very very nice post', NOW(), NOW()), ('Test post 1', 'Very very nice post', NOW(), NOW());`,
		`INSERT INTO tags (name, highlighted, created_at, updated_at) VALUES ('tag1', TRUE, NOW(), NOW()), ('tag2', TRUE, NOW(), NOW()), ('tag3', TRUE, NOW(), NOW()), ('tag4', TRUE, NOW(), NOW());`,
		`INSERT INTO posts_tags (post_id, tag_id) VALUES (1, 1), (1, 2), (2, 3), (2, 4);`,
	}

	for _, mock := range mockData {
		_, err = db.Exec(ctx, mock)
		if err != nil {
			log.Fatal(err)
		}
	}

	// We query our result with specifying the order of columns to match the Go struct,
	// so basically it is using the pgx.RowToStructByPos for the tags.
	fmt.Println(">>> Working:")
	stmt := `
		SELECT
			posts.*,
			(
				SELECT array_agg(
					row(
						tags.id,
						tags.name,
						tags.highlighted,
						tags.created_at,
						tags.updated_at
					)
				)
				FROM tags
				INNER JOIN posts_tags ON posts_tags.tag_id = tags.id AND posts_tags.post_id = posts.id
			) AS tags
		FROM posts;
	`

	rows, err := db.Query(ctx, stmt)
	if err != nil {
		log.Fatal(err)
	}

	posts, err := pgx.CollectRows(rows, pgx.RowToStructByName[Post])
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("%#v\n", posts)
	fmt.Println(">>> Failing:")

	// This testcase fails as it is trying to map the created_at column into the highlighted column
	// as the struct fields do not match the order of the columns.
	stmt = `
		SELECT
			posts.*,
			(
				SELECT array_agg(
					row(
						tags.*
					)
				)
				FROM tags
				INNER JOIN posts_tags ON posts_tags.tag_id = tags.id AND posts_tags.post_id = posts.id
			) AS tags
		FROM posts;
	`

	rows, err = db.Query(ctx, stmt)
	if err != nil {
		log.Fatal(err)
	}

	posts, err = pgx.CollectRows(rows, pgx.RowToStructByName[Post])
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("%#v\n", posts)
}

Output:

>>> Working:
[]main.Post{main.Post{ID:1, Title:"Test post 1", Body:"Very very nice post", Tags:[]main.Tag{main.Tag{ID:1, Name:"tag1", Highlighted:true, CreatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823874000, time.Local), UpdatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823874000, time.Local)}, main.Tag{ID:2, Name:"tag2", Highlighted:true, CreatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823874000, time.Local), UpdatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823874000, time.Local)}}, CreatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823513000, time.Local), UpdatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823513000, time.Local)}, main.Post{ID:2, Title:"Test post 1", Body:"Very very nice post", Tags:[]main.Tag{main.Tag{ID:3, Name:"tag3", Highlighted:true, CreatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823874000, time.Local), UpdatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823874000, time.Local)}, main.Tag{ID:4, Name:"tag4", Highlighted:true, CreatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823874000, time.Local), UpdatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823874000, time.Local)}}, CreatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823513000, time.Local), UpdatedAt:time.Date(2024, time.February, 20, 10, 30, 5, 823513000, time.Local)}}

>>> Failing:
2024/02/20 10:30:05 can't scan into dest[5]: failed to scan array element 0: cannot scan timestamptz (OID 1184) in binary format into *bool
exit status 1

[jackc]

It potentially is possible, but is a little tricky. At the top level of a query, RowToStructByName can work because the query results include the names of the columns. However, composite values do not include the names of the fields. They only have the field values (and field type OIDs in the binary format). That is why composites are always read positionally.

In theory, it should be possible to map by name if the composite types have been registered so the type definition is available in the *pgtype.Map. I'm not sure how much plumbing would be required to implement that though.

[svennis94]

Good to know! That was what I was suspecting.

What I am trying now is to write a custom Codec for the tags and _tags OID that maps it by name. I'll reply here with the codec that might work for other persons that want to do the same.

To be continued...

[svennis94]

So after some tinkering I have build a PoC, this is by no means perfect. But it works without any changes to PGx itself. We have decided to go a different route and we will defined the columns in the correct order. But I'll leave the PoC here in case anybody else is looking to retrieve the same result.

There was one small issue that I ran into and that is that the pgtype.ptrStructWrapper has the exportedFields field, sadly that field is not exported. If it was exported then I could use the index from the reflect.Value based on the column and field name. For the PoC I generated my own mapping (column name -> field index in struct, as a map[string]int). Other than that I have used the default CompositeCodec.

I started by adding the types to the query:

SELECT
    posts.*,
    (
        SELECT array_agg(
            row(
                tags.*
            )::tags
        )::tags[]
        FROM tags
        INNER JOIN posts_tags ON posts_tags.tag_id = tags.id AND posts_tags.post_id = posts.id
    ) AS tags
FROM posts;

Then in the AfterConnect of the database pool I added the LoadType for those types to the database pool config. Be aware that you'd already need to have these types/tables created in the database. Otherwise the LoadType will throw an error that the type does not exists. In other words, don't register these types when you are using migrations.

config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
    tagDataType, err := conn.LoadType(ctx, "tags")
    if err != nil {
        return err
    }

    cCodec, ok := tagDataType.Codec.(*pgtype.CompositeCodec)
    if ok {
        dt := &CustomCodec{}
        dt.Fields = cCodec.Fields
        // Define the column names -> struct field index for the type
        // This could be made generic by using reflection instead of hard coding column names and indexes, but this is a PoC
        dt.StrctMap = map[string]int{
            "id":          0,
            "name":        1,
            "highlighted": 2,
            "created_at":  3,
            "updated_at":  4,
        }

        tagDataType.Codec = dt
    }

    conn.TypeMap().RegisterType(tagDataType)

    dataType, err := conn.LoadType(ctx, "_tags")
    if err != nil {
        return err
    }

    dataType.Codec = &pgtype.ArrayCodec{
        ElementType: tagDataType,
    }

    conn.TypeMap().RegisterType(dataType)

    return nil
}

In the code you can see that the CustomCodec is being used this is a codec that extends the CompositeCodec. It will retrieve the column names from the binary data and use the column name to look up the struct index and use that index to scan. So all of the logic is still the same as the CompositeCodec that scans based on the indexes. The CustomCodec is just using a mapping of the column names to struct field index. It also sets the preferred format to binary format to retrieve the column names of the composite type.

type CustomCodec struct {
    pgtype.CompositeCodec
    StrctMap map[string]int
}

func (c *CustomCodec) PreferredFormat() int16 {
    return pgtype.BinaryFormatCode
}

func (c *CustomCodec) PlanScan(m *pgtype.Map, oid uint32, format int16, target any) pgtype.ScanPlan {
    switch format {
    case pgtype.BinaryFormatCode:
        switch target.(type) {
        case pgtype.CompositeIndexScanner:
            return &scanPlanBinaryCompositeToCompositeIndexScanner{cc: c, m: m}
        }
    }

    return nil
}

type scanPlanBinaryCompositeToCompositeIndexScanner struct {
    cc *CustomCodec
    m  *pgtype.Map
}

func (plan *scanPlanBinaryCompositeToCompositeIndexScanner) Scan(src []byte, target any) error {
    targetScanner := (target).(pgtype.CompositeIndexScanner)

    if src == nil {
        return targetScanner.ScanNull()
    }

    scanner := pgtype.NewCompositeBinaryScanner(plan.m, src)
    for _, field := range plan.cc.Fields {
        if scanner.Next() {
            fieldTarget := targetScanner.ScanIndex(plan.cc.StrctMap[field.Name])
            if fieldTarget != nil {
                fieldPlan := plan.m.PlanScan(field.Type.OID, pgtype.BinaryFormatCode, fieldTarget)
                if fieldPlan == nil {
                    return fmt.Errorf("unable to encode %v into OID %d in binary format", field, field.Type.OID)
                }

                err := fieldPlan.Scan(scanner.Bytes(), fieldTarget)
                if err != nil {
                    return err
                }
            }
        } else {
            return errors.New("read past end of composite")
        }
    }

    if err := scanner.Err(); err != nil {
        return err
    }

    return nil
}

If the CompositeIndexScanner would expose the exportedFields you could remove the CustomCodec.StrctMap and use the exportedFields to retrieve the index in the struct based on column names.

One idea is to create a PR (if there is interest for it and it benefits PGx) to create a pgtype.CompositeByNameCodec. You would still need to define your custom types in the AfterConnect to be able to retrieve the column names etc. But then you would be able to define to use the CompositeByNameCodec from PGx as the codec and you don't need to define the codec yourself. But not sure if this is interesting to add to PGx.

We have decided to not continue this route but instead just specify the columns, as it is easier to understand and there is less code to maintain. This is a PoC so I would not recommend to use this in production, we have not written tests for this and have not tested it with for example nested composite types. Use at your own risk! 😄

Thanks for the help!