From ed1cf7ffef85d92afb8aebd6f9efad2375a1ce4d Mon Sep 17 00:00:00 2001 From: Todd Gruben Date: Mon, 27 Sep 2021 07:05:08 -0500 Subject: [PATCH] cleanup and applied review suggestions --- go.mod | 2 +- pg/message/io.go | 2 - pg/message/message.go | 6 +- pg/protocol.go | 156 +++++++++++++++++++----------------------- planner.go | 3 +- 5 files changed, 78 insertions(+), 91 deletions(-) diff --git a/go.mod b/go.mod index e590bdd1b..d64a32730 100644 --- a/go.mod +++ b/go.mod @@ -54,7 +54,7 @@ require ( golang.org/x/sys v0.0.0-20210510120138-977fb7262007 // indirect golang.org/x/text v0.3.5 // indirect google.golang.org/grpc v1.28.0 - gopkg.in/yaml.v2 v2.3.0 // indirect + gopkg.in/yaml.v2 v2.3.0 modernc.org/mathutil v1.0.0 modernc.org/strutil v1.0.0 sigs.k8s.io/yaml v1.2.0 // indirect diff --git a/pg/message/io.go b/pg/message/io.go index ec2321675..e09150b9d 100644 --- a/pg/message/io.go +++ b/pg/message/io.go @@ -18,7 +18,6 @@ import ( "bufio" "encoding/binary" "errors" - "fmt" "io" ) @@ -93,7 +92,6 @@ type WireWriter struct { // WriteMessage writes a message onto the wire. func (w *WireWriter) WriteMessage(message Message) error { - fmt.Printf("SendToClient %c (%d)\n", message.Type, len(message.Data)) if uint(len(message.Data))+4 >= 1<<31 { return ErrMessageTooBig } diff --git a/pg/message/message.go b/pg/message/message.go index 757053297..9fb6871ee 100644 --- a/pg/message/message.go +++ b/pg/message/message.go @@ -30,7 +30,7 @@ const ( // TypeReadyForQuery is a message used to indicate that the server is ready for another query. TypeReadyForQuery Type = 'Z' - // TypeCommandComplete is a message used to indicate that a query has completed. Backend + // TypeCommandComplete is a Backend message used to indicate that a query has completed. TypeCommandComplete Type = 'C' // TypeClos is a message used to indicate that a query has completed. Frontend @@ -60,8 +60,8 @@ const ( TypeBind Type = 'B' TypeBindComplete Type = '2' - TypeExecute Type = 'E' // Frontend TODO(TWG) - TypeError Type = 'E' // Backend SOMETHING NOT RIGHT HERE + TypeExecute Type = 'E' // Frontend + TypeError Type = 'E' // Backend TypeSync Type = 'S' // Frontend TypeParameterStatus Type = 'S' // Backend TypeDescribe Type = 'D' // Frontend diff --git a/pg/protocol.go b/pg/protocol.go index fdfce083e..0fa52a5e9 100644 --- a/pg/protocol.go +++ b/pg/protocol.go @@ -272,12 +272,11 @@ const ( ) type Portal struct { - Name string - Writer *message.WireWriter - commands []message.Message - Encoder *message.Encoder - mapper *sql.Mapper - //results []Result + Name string + Writer *message.WireWriter + commands []message.Message + Encoder *message.Encoder + mapper *sql.Mapper sql string pgspecial PgType pid int32 @@ -332,76 +331,75 @@ func (p *Portal) Parse(data []byte) { query, err := p.mapper.MapSQL(queryStr) if err != nil { return + } + + if strings.Contains(strings.ToLower(query.SQL), "select 1") { + p.pgspecial = pgSelect1 + p.Name = "SELECT" } else { - - if strings.Contains(strings.ToLower(query.SQL), "select 1") { - p.pgspecial = pgSelect1 - p.Name = "SELECT" - } else { - switch query.SQLType { - case sql.SQLTypeSet: - p.Name = "SET" - set := query.Statement.(*sqlparser.Set) - p.pgspecial = 0 - for _, item := range set.Exprs { - if item.Name.String() == "application_name" { - switch item.Expr.(type) { - case *sqlparser.SQLVal: - p.pgspecial = pgSetApplication - } + switch query.SQLType { + case sql.SQLTypeSet: + p.Name = "SET" + set := query.Statement.(*sqlparser.Set) + p.pgspecial = 0 + for _, item := range set.Exprs { + if item.Name.String() == "application_name" { + switch item.Expr.(type) { + case *sqlparser.SQLVal: + p.pgspecial = pgSetApplication } } - case sql.SQLTypeSelect: - p.Name = "SELECT" - p.pgspecial = pgPassOn - stmt := query.Statement.(*sqlparser.Select) - for _, item := range stmt.SelectExprs { - switch expr := item.(type) { - case *sqlparser.AliasedExpr: - switch colExpr := expr.Expr.(type) { - case *sqlparser.FuncExpr: - funcName := strings.ToLower(colExpr.Name.String()) - switch funcName { - case "pg_backend_pid": - //SELECT pg_backend_pid() - p.pgspecial = pgBackendPid - case "pg_terminate_backend": - //select pg_terminate_backend(100) - p.pgspecial = pgTerminate - case "version": - //SELECT VERSION() AS version - p.pgspecial = pgVersion - } - //need to return the pid from the cancelation object - //add row description object - //add data row for item - } - } - - } - for _, item := range stmt.From { - switch from := item.(type) { - case *sqlparser.AliasedTableExpr: - tableName := from.Expr.(sqlparser.TableName).ToViewName().Name.String() - switch tableName { - case "pg_type": - p.pgspecial = pgCountType - case "pg_stat_activity": - p.pgspecial = pgQueryTime - case "tables": - p.pgspecial = pgSchema - } - } - } - p.sql = queryStr - case sql.SQLTypeBegin: - // Ignore BEGIN - p.pgspecial = pgBegin - case sql.SQLTypeShow: - p.Name = "SHOW" - p.pgspecial = pgPassOn - p.sql = queryStr } + case sql.SQLTypeSelect: + p.Name = "SELECT" + p.pgspecial = pgPassOn + stmt := query.Statement.(*sqlparser.Select) + for _, item := range stmt.SelectExprs { + switch expr := item.(type) { + case *sqlparser.AliasedExpr: + switch colExpr := expr.Expr.(type) { + case *sqlparser.FuncExpr: + funcName := strings.ToLower(colExpr.Name.String()) + switch funcName { + case "pg_backend_pid": + //SELECT pg_backend_pid() + p.pgspecial = pgBackendPid + case "pg_terminate_backend": + //select pg_terminate_backend(100) + p.pgspecial = pgTerminate + case "version": + //SELECT VERSION() AS version + p.pgspecial = pgVersion + } + //need to return the pid from the cancelation object + //add row description object + //add data row for item + } + } + + } + for _, item := range stmt.From { + switch from := item.(type) { + case *sqlparser.AliasedTableExpr: + tableName := from.Expr.(sqlparser.TableName).ToViewName().Name.String() + switch tableName { + case "pg_type": + p.pgspecial = pgCountType + case "pg_stat_activity": + p.pgspecial = pgQueryTime + case "tables": + p.pgspecial = pgSchema + } + } + } + p.sql = queryStr + case sql.SQLTypeBegin: + // Ignore BEGIN + p.pgspecial = pgBegin + case sql.SQLTypeShow: + p.Name = "SHOW" + p.pgspecial = pgPassOn + p.sql = queryStr } } } else { @@ -410,18 +408,9 @@ func (p *Portal) Parse(data []byte) { p.Add(message.ParseOK) } func (p *Portal) Describe() { - /* - if p.pgspecial != pgNotPg { - //do custom handling for setup - } - if len(p.results) == 0 { - p.Add(&message.NoData) - p.Add(&message.EmptyQueryResponse) - return - } - */ - + // Placeholder should we need to handle the Decribe request } + func (p *Portal) Execute() (shouldTerminate bool, queryReady bool, err error) { queryReady = true switch p.pgspecial { @@ -824,7 +813,6 @@ func (s *Server) handleStandard(ctx context.Context, proto Protocol, conn net.Co // Read the next packet. msg, err := r.ReadMessage() - msg.Dump("start-") // TODO(twg) remove if err != nil { if err == errPreempted { // The server is shutting down. diff --git a/planner.go b/planner.go index 985a92001..196bfa817 100644 --- a/planner.go +++ b/planner.go @@ -645,10 +645,11 @@ func (rs *StmtRows) Columns() []*StmtColumn { return rs.node.Columns() } + /* func (rs *StmtRows) Row() int64 { return rs.node.Row()[0].(int64) } - + */ func (rs *StmtRows) Next() bool { if rs.err != nil { return false