sqlc-dev / sqlc

Generate type-safe code from SQL
https://sqlc.dev
MIT License
12.33k stars 780 forks source link

Enum is not working with CopyFrom #2116

Open zhenik123 opened 1 year ago

zhenik123 commented 1 year ago

Version

1.16.0

What happened?

Postgres copyfrom query in a table which has a column with Enum type fails. I am using pgx5.

Relevant log output

ERROR:  COPY from stdin failed: unable to encode "ENUM_VALUE" into binary format for unknown type (OID 16393): cannot find encode plan

Database schema

CREATE TABLE violation (
    slug VARCHAR(64) PRIMARY KEY UNIQUE,
    parent VARCHAR(64),
    FOREIGN KEY (parent) REFERENCES violation(slug) ON DELETE SET NULL
);

CREATE TYPE reference_types AS ENUM (
    'DEFAULT',
    'ENUM_VALUE',
    'OTHER_VALUE'
);

CREATE TABLE user_violation (
    id serial PRIMARY KEY,
    violation_slug VARCHAR(64) NOT NULL,
    FOREIGN KEY (violation_slug) REFERENCES violation(slug) ON DELETE SET NULL,
    reference_type reference_types NOT NULL,
);

SQL queries

-- name: BulkInsertUserViolations :copyfrom
INSERT INTO user_violation (
    violation_slug,
    reference_type
) VALUES (
    $1,
    $2
);

Configuration

version: "1"
packages:
  - name: "violation"
    path: "../internal/violation"
    queries: "../internal/violation/queries"
    schema: "../internal/violation/migrations"
    engine: "postgresql"
    sql_package: "pgx/v5"
    emit_json_tags: false
    emit_prepared_queries: false
    emit_interface: true
    emit_exact_table_names: false
overrides:
  - go_type: "github.com/jackc/pgx/v5/pgtype.UUID"
    db_type: "uuid"

Playground URL

https://play.sqlc.dev/p/70883aaa9d4b68218cca526275348874560f43aa01ffa362d8364516f5104c16

What operating system are you using?

Linux

What database engines are you using?

PostgreSQL

What type of code are you generating?

Go

zhenik123 commented 1 year ago

According to pgx documentation CopyFrom requires all values use the binary format. A pgtype.Type that supports the binary format must be registered for the type of each column. Almost all types implemented by pgx support the binary format. Even though enum types appear to be strings they still must be registered to use with CopyFrom. This can be done with Conn.LoadType and pgtype.Map.RegisterType.

I used AfterConnect hook in pgx pool config which is called after a connection is established, but before it is added to the pool. and used Conn.LoadType and pgtype.Map.RegisterType. It solved my problem. but I guess it would be better if it was handled in sqlc auto generated codes.

getSuitableType := func(enumType string) (*pgtype.Type, error) {
    conn, err := pool.Acquire(ctx)
    defer conn.Release()
    if err != nil {
    return nil, err
    }
    return conn.Conn().LoadType(ctx, enumType)
}

var registerSuitableTypes []*pgtype.Type
for _, enumName := range enumTypes {
    dt, err := getSuitableType(enumName)
    if err != nil {
        return fmt.Errorf("failed to load suitable enum type: %w", err)
    }
    registerSuitableTypes = append(registerSuitableTypes, dt)
}

pgxPoolConfig.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
    for _, dt := range registerSuitableTypes {
        conn.TypeMap().RegisterType(dt)
    }
    return nil
}
louisrli commented 1 year ago

Hey, thanks a lot for the code above. I used it as a starting point, I'm going to leave my code here in case it helps anyone else. In my case, I had an enum called lang_tag and also needed to use the array version (_lang_tag in the pgtypes table, as the convention is for array types to be prefixed with an underscore).

In case you also need to use the array type of the enum, you actually have to register the enum type before loading the array type; thus my code has a slight extra line of code from the one above to register the type (for the temporary connection only). If you don't register the type immediately, then when you load the array type, you'll run into array element oid not registered (in the linked source you'll see that it tries to load type on the enum_type before loading the array _enum_type)

func GetDbPool() (*pgxpool.Pool, error) {
        // Set up a new pool with the custom types and the config.
        config, err := pgxpool.ParseConfig(MakeConnectionString())
        if err != nil {
                return nil, err
        }
        dbpool, err := pgxpool.NewWithConfig(context.Background(), config)

        // Collect the custom data types once, store them in memory, and register them for every future connection.
        customTypes, err := getCustomDataTypes(context.Background(), dbpool)
        config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
                for _, t := range customTypes {
                        conn.TypeMap().RegisterType(t)
                }
                return nil
        }
        // Immediately close the old pool and open a new one with the new config.
        dbpool.Close()
        dbpool, err = pgxpool.NewWithConfig(context.Background(), config)
        return dbpool, err
}

// Any custom DB types made with CREATE TYPE need to be registered with pgx.
// https://github.com/kyleconroy/sqlc/issues/2116
// https://stackoverflow.com/questions/75658429/need-to-update-psql-row-of-a-composite-type-in-golang-with-jack-pgx
// https://pkg.go.dev/github.com/jackc/pgx/v5/pgtype
func getCustomDataTypes(ctx context.Context, pool *pgxpool.Pool) ([]*pgtype.Type, error) {
        // Get a single connection just to load type information.
        conn, err := pool.Acquire(ctx)
        defer conn.Release()
        if err != nil {
                return nil, err
        }

        dataTypeNames := []string{
                "lang_tag",
                // An underscore prefix is an array type in pgtypes.
                "_lang_tag",
        }

        var typesToRegister []*pgtype.Type
        for _, typeName := range dataTypeNames {
                dataType, err := conn.Conn().LoadType(ctx, typeName)
                if err != nil {
                        return nil, status.Errorf(codes.Internal, "failed to load type %s: %v", typeName, err)
                }
                // You need to register only for this connection too, otherwise the array type will look for the register element type.
                conn.Conn().TypeMap().RegisterType(dataType)
                typesToRegister = append(typesToRegister, dataType)
        }
        return typesToRegister, nil
}
louisrli commented 1 year ago

Also worth noting that I ran into this issue not with COPYFROM, but an UPDATE query (in either case it has to do with pgx5 + enum types needing to be registered, though I don't know enough about how sqlc works to know if sqlc is trying/failing to automatically register them somehow)

yeletisrujan commented 8 months ago

I had the exact same problem. Thanks to this code snippet, it helped me!