Merge pull request #1718 from tgruben/sql2-type

Better type support in result set for looker (postgres) sql2 interface

pg/protocol.go

@@ -269,6 +269,7 @@ const (
 	pgSelect1        PgType = 'h'
 	pgSchema         PgType = 'i'
 	pgBegin          PgType = 'j'
+	pgTypeLen        PgType = 'k'
 )
 
 type Portal struct {
@@ -298,6 +299,8 @@ func (p *Portal) Bind() {
 
 var lookPQL = regexp.MustCompile(`\[.*\].*\)\z`)
 
+const POSTGRESLENSQL = `SELECT t.typlen FROM pg_catalog.pg_type t, pg_catalog.pg_namespace n WHERE t.typnamespace=n.oid AND t.typname='name' AND n.nspname='pg_catalog'`
+
 func (p *Portal) Parse(data []byte) {
 	p.queryStart = time.Now()
 	queryStr := string(bytes.Trim(data, "\x00"))
@@ -327,7 +330,6 @@ func (p *Portal) Parse(data []byte) {
 	}
 
 	if len(queryStr) > 2 {
-
 		query, err := p.mapper.MapSQL(queryStr)
 		if err != nil {
 			return
@@ -336,6 +338,9 @@ func (p *Portal) Parse(data []byte) {
 		if strings.Contains(strings.ToLower(query.SQL), "select 1") {
 			p.pgspecial = pgSelect1
 			p.Name = "SELECT"
+		} else if strings.Contains(queryStr, POSTGRESLENSQL) {
+			p.pgspecial = pgTypeLen
+			p.Name = "SELECT"
 		} else {
 			switch query.SQLType {
 			case sql.SQLTypeSet:
@@ -565,6 +570,16 @@ func (p *Portal) Execute() (shouldTerminate bool, queryReady bool, err error) {
 
 	case pgBegin:
 		p.Name = "BEGIN"
+	case pgTypeLen:
+		rowDescription, e := p.Encoder.EncodeColumn("typelen", int32(21), 2)
+		if e != nil {
+			err = e
+			return
+		}
+		p.Add(rowDescription)
+		mesg := "64"
+		dataRow, _ := p.Encoder.TextRow(mesg)
+		p.Add(dataRow)
 	}
 
 	//maybe add in the number of items in select clause

pg/type.go

@@ -20,18 +20,17 @@ import "github.com/molecula/featurebase/v2/pg/message"
 type Type struct {
 	// I am not entirely sure what should be in here long term.
 	// For now, I am just going to leave it like this.
-	id int32
+	Id      int32
+	Typelen int16
 }
 
 // TypeCharoid is a postgres type for text.
-var TypeCharoid = Type{id: 18}
-
-// TypeData is a type containing raw postgres wire type information.
-type TypeData struct {
-	TypeID       int32
-	TypeLen      int16
-	TypeModifier int32
-}
+// found in postgres source src/include/catalog/pg_type.h
+var TypeCharoid = Type{Id: 18, Typelen: -1}
+var TypeNAMEOID = Type{Id: 19, Typelen: 64}
+var TypeINT4OID = Type{Id: 23, Typelen: 4}
+var TypeTEXTOID = Type{Id: 25, Typelen: -1}
+var TypeFLOAT8OID = Type{Id: 701, Typelen: 8}
 
 // TypeEngine is a system for managing types.
 // This is necessary for compound types like arrays which need ID generation.
@@ -45,9 +44,31 @@ type PrimitiveTypeEngine struct{}
 
 // TranslateType translates a type to a column description.
 func (pte PrimitiveTypeEngine) TranslateType(t Type) (message.ColumnDescription, error) {
+	var TypeLen int16
+	var TypeID int32
+	switch t {
+	case TypeCharoid:
+		TypeID = TypeCharoid.Id
+		TypeLen = TypeCharoid.Typelen
+	case TypeNAMEOID:
+		TypeID = TypeNAMEOID.Id
+		TypeLen = TypeNAMEOID.Typelen
+	case TypeINT4OID:
+		TypeID = TypeINT4OID.Id
+		TypeLen = TypeINT4OID.Typelen
+	case TypeTEXTOID:
+		TypeID = TypeTEXTOID.Id
+		TypeLen = TypeTEXTOID.Typelen
+	case TypeFLOAT8OID:
+		TypeID = TypeFLOAT8OID.Id
+		TypeLen = TypeFLOAT8OID.Typelen
+	default: // treat like TypeCharoid:
+		TypeID = TypeCharoid.Id
+		TypeLen = TypeCharoid.Typelen
+	}
 	return message.ColumnDescription{
-		TypeID:       t.id, // just charoid for now; as far as I can tell most implementations do not really use this
-		TypeLen:      -1,   // vdsm had 4. . . but the spec says this should be negative
+		TypeID:       TypeID,
+		TypeLen:      TypeLen,
 		TypeModifier: -1,
 		Mode:         0, // send as text
 	}, nil

server/pg.go

@@ -28,6 +28,7 @@ import (
 	pilosa "github.com/molecula/featurebase/v2"
 	"github.com/molecula/featurebase/v2/logger"
 	"github.com/molecula/featurebase/v2/pg"
+	"github.com/molecula/featurebase/v2/sql2"
 
 	//"github.com/molecula/featurebase/v2/pg"
 	"github.com/molecula/featurebase/v2/pql"
@@ -336,6 +337,15 @@ func pgWriteGroupCount(w pg.QueryResultWriter, counts *pilosa.GroupCounts) error
 	return nil
 }
 
+// TODO(twg) move this to a better area
+func getPgType(sql2type string) pg.Type {
+	ret := pg.TypeCharoid
+	switch sql2type {
+	case sql2.DataTypeInt:
+		ret = pg.TypeINT4OID
+	}
+	return ret
+}
 func pgWriteStmtRows(w pg.QueryResultWriter, rows *pilosa.StmtRows) error {
 	//TODO(twg) writeHeader
 	//TODO(twg) writeColumns
@@ -348,9 +358,10 @@ func pgWriteStmtRows(w pg.QueryResultWriter, rows *pilosa.StmtRows) error {
 			//TODO (twg) types:=rows.Types()
 			headers := make([]pg.ColumnInfo, len(columns))
 			for i, column := range columns {
+				pgType := getPgType(column.Type)
 				headers[i] = pg.ColumnInfo{
 					Name: column.Name,
-					Type: pg.TypeCharoid, //TODO(twg) types[i]
+					Type: pgType,
 				}
 			}
 			err := w.WriteHeader(headers...)
@@ -407,6 +418,17 @@ func pgWriteStmtRows(w pg.QueryResultWriter, rows *pilosa.StmtRows) error {
 	}
 	return nil
 }
+func getPgTypeFromColumnInfo(sql2type string) pg.Type {
+	switch sql2type {
+	case sql2.DataTypeInt:
+		return pg.TypeINT4OID
+	default:
+		return pg.TypeCharoid
+
+	}
+}
+
+var _ = getPgTypeFromColumnInfo //make linter happy for this function will be needed in future
 
 func pgWriteRowser(w pg.QueryResultWriter, result pb.ToRowser) error {
 	var data []string
@@ -416,7 +438,8 @@ func pgWriteRowser(w pg.QueryResultWriter, result pb.ToRowser) error {
 			for i, h := range row.Headers {
 				headers[i] = pg.ColumnInfo{
 					Name: h.Name,
-					Type: pg.TypeCharoid,
+					Type: pg.TypeCharoid, // TODO(twg) this needs to be updated with type
+					// information so it works from psql client
 				}
 			}
 			err := w.WriteHeader(headers...)
@@ -481,7 +504,7 @@ func pgWriteResult(w pg.QueryResultWriter, result interface{}) error {
 		return pgWriteGroupCount(w, result)
 	case pb.ToRowser: // we should avoid protobuf where we can...
 		return pgWriteRowser(w, result)
-	case *pilosa.StmtRows: // we should avoid protobuf where we can...
+	case *pilosa.StmtRows:
 		return pgWriteStmtRows(w, result)
 	case uint64:
 		err := w.WriteHeader(pg.ColumnInfo{