From db7df79e1061144d3aad7ec264c2c24774e3c217 Mon Sep 17 00:00:00 2001 From: Jack Christensen Date: Sat, 7 Apr 2018 14:03:29 -0500 Subject: [PATCH] Add domain support fixes #407 --- CHANGELOG.md | 1 + README.md | 1 + conn.go | 51 +++++++++++++++++++++++++++++++++++++++ conn_test.go | 25 +++++++++++++++++++ pgmock/pgmock.go | 49 +++++++++++++++++++++++++++++++++++++ travis/before_script.bash | 1 + 6 files changed, 128 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 78ee6932..a4ff956d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ * Support sslkey, sslcert, and sslrootcert URI params (Sean Chittenden) * Allow any scheme in ParseURI (for convenience with cockroachdb) (Sean Chittenden) +* Add support for domain types ## Fixes diff --git a/README.md b/README.md index 1acaabff..d6499ba4 100644 --- a/README.md +++ b/README.md @@ -100,6 +100,7 @@ Then run the following SQL: Connect to database pgx_test and run: create extension hstore; + create domain uint64 as numeric(20,0); Next open conn_config_test.go.example and make a copy without the .example. If your PostgreSQL server is accepting connections on 127.0.0.1, diff --git a/conn.go b/conn.go index 9cb325fe..2bdb7181 100644 --- a/conn.go +++ b/conn.go @@ -15,6 +15,7 @@ import ( "os" "os/user" "path/filepath" + "reflect" "regexp" "strconv" "strings" @@ -428,6 +429,10 @@ where ( return nil, err } + if err = c.initConnInfoDomains(cinfo); err != nil { + return nil, err + } + return cinfo, nil } @@ -494,6 +499,52 @@ where t.typtype = 'b' return nil } +// initConnInfoDomains introspects for domains and registers a data type for them. +func (c *Conn) initConnInfoDomains(cinfo *pgtype.ConnInfo) error { + type domain struct { + oid pgtype.OID + name string + baseOID pgtype.OID + } + + domains := make([]*domain, 0, 16) + + rows, err := c.Query(`select t.oid, t.typname, t.typbasetype +from pg_type t + join pg_type base_type on t.typbasetype=base_type.oid +where t.typtype = 'd' + and base_type.typtype = 'b'`) + if err != nil { + return err + } + + for rows.Next() { + var d domain + if err := rows.Scan(&d.oid, &d.name, &d.baseOID); err != nil { + return err + } + + domains = append(domains, &d) + } + + if rows.Err() != nil { + return rows.Err() + } + + for _, d := range domains { + baseDataType, ok := cinfo.DataTypeForOID(d.baseOID) + if ok { + cinfo.RegisterDataType(pgtype.DataType{ + Value: reflect.New(reflect.ValueOf(baseDataType.Value).Elem().Type()).Interface().(pgtype.Value), + Name: d.name, + OID: d.oid, + }) + } + } + + return nil +} + // crateDBTypesQuery checks if the given err is likely to be the result of // CrateDB not implementing the pg_types table correctly. If yes, a CrateDB // specific query against pg_types is executed and its results are returned. If diff --git a/conn_test.go b/conn_test.go index 6f1d41ea..df57b670 100644 --- a/conn_test.go +++ b/conn_test.go @@ -2073,3 +2073,28 @@ func TestConnInitConnInfo(t *testing.T) { ensureConnValid(t, conn) } + +func TestDomainType(t *testing.T) { + conn := mustConnect(t, *defaultConnConfig) + defer closeConn(t, conn) + + dt, ok := conn.ConnInfo.DataTypeForName("uint64") + if !ok { + t.Fatal("Expected data type for domain uint64 to be present") + } + if dt, ok := dt.Value.(*pgtype.Numeric); !ok { + t.Fatal("Expected data type value for domain uint64 to be *pgtype.Numeric, but it was %T", dt) + } + + var n uint64 + err := conn.QueryRow("select $1::uint64", uint64(42)).Scan(&n) + if err != nil { + t.Fatal(err) + } + + if n != 42 { + t.Fatalf("Expected n to be 42, but was %v", n) + } + + ensureConnValid(t, conn) +} diff --git a/pgmock/pgmock.go b/pgmock/pgmock.go index 71f18852..fec0d3f6 100644 --- a/pgmock/pgmock.go +++ b/pgmock/pgmock.go @@ -479,6 +479,55 @@ where ( SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}), }...) + steps = append(steps, []Step{ + ExpectMessage(&pgproto3.Parse{ + Query: "select t.oid, t.typname, t.typbasetype\nfrom pg_type t\n join pg_type base_type on t.typbasetype=base_type.oid\nwhere t.typtype = 'd'\n and base_type.typtype = 'b'", + }), + ExpectMessage(&pgproto3.Describe{ + ObjectType: 'S', + }), + ExpectMessage(&pgproto3.Sync{}), + SendMessage(&pgproto3.ParseComplete{}), + SendMessage(&pgproto3.ParameterDescription{}), + SendMessage(&pgproto3.RowDescription{ + Fields: []pgproto3.FieldDescription{ + {Name: "oid", + TableOID: 1247, + TableAttributeNumber: 65534, + DataTypeOID: 26, + DataTypeSize: 4, + TypeModifier: 4294967295, + Format: 0, + }, + {Name: "typname", + TableOID: 1247, + TableAttributeNumber: 1, + DataTypeOID: 19, + DataTypeSize: 64, + TypeModifier: 4294967295, + Format: 0, + }, + {Name: "typbasetype", + TableOID: 1247, + TableAttributeNumber: 65534, + DataTypeOID: 26, + DataTypeSize: 4, + TypeModifier: 4294967295, + Format: 0, + }, + }, + }), + SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}), + ExpectMessage(&pgproto3.Bind{ + ResultFormatCodes: []int16{1, 1, 1}, + }), + ExpectMessage(&pgproto3.Execute{}), + ExpectMessage(&pgproto3.Sync{}), + SendMessage(&pgproto3.BindComplete{}), + SendMessage(&pgproto3.CommandComplete{CommandTag: "SELECT 0"}), + SendMessage(&pgproto3.ReadyForQuery{TxStatus: 'I'}), + }...) + return steps } diff --git a/travis/before_script.bash b/travis/before_script.bash index 7e206e7a..e9c0f865 100755 --- a/travis/before_script.bash +++ b/travis/before_script.bash @@ -9,6 +9,7 @@ then # of aclitem formatting. It turns out aclitems cannot contain non-existing users/roles. psql -U postgres -c 'create database pgx_test' psql -U postgres pgx_test -c 'create extension hstore' + psql -U postgres pgx_test -c 'create domain uint64 as numeric(20,0)' psql -U postgres -c "create user pgx_ssl SUPERUSER PASSWORD 'secret'" psql -U postgres -c "create user pgx_md5 SUPERUSER PASSWORD 'secret'" psql -U postgres -c "create user pgx_pw SUPERUSER PASSWORD 'secret'"