featurebase/pg/protocol.go

// Copyright 2022 Molecula Corp. (DBA FeatureBase).
// SPDX-License-Identifier: Apache-2.0
package pg

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/featurebasedb/featurebase/v3/pg/message"
	"github.com/featurebasedb/featurebase/v3/sql"
	"github.com/pkg/errors"
	"vitess.io/vitess/go/vt/sqlparser"
)

// Protocol is a Postgres protocol version.
type Protocol uint32

const (
	// ProtocolPostgres30 is version 3.0 of the Postgres wire protocol.
	ProtocolPostgres30 Protocol = (3 << 16)

	// ProtocolCancel is the protocol used for query cancellation.
	ProtocolCancel Protocol = (1234 << 16) | 5678

	// ProtocolSSL is the protocol used for SSL upgrades.
	ProtocolSSL Protocol = (1234 << 16) | 5679

	// ProtocolSupported is the main protocol version supported by this package.
	ProtocolSupported Protocol = ProtocolPostgres30

	// PgServerVersion is the latest version of postgres that we claim to support.
	PgServerVersion = "13.0.0"
)

// Major returns the major revision of the protocol.
func (p Protocol) Major() uint16 {
	return uint16(p >> 16)
}

// Minor returns the minor revision of the protocol.
func (p Protocol) Minor() uint16 {
	return uint16(p)
}

func (p Protocol) String() string {
	switch p {
	case ProtocolCancel:
		return "cancel"
	case ProtocolSSL:
		return "SSL"
	}

	return fmt.Sprintf("v%d.%d", p.Major(), p.Minor())
}

// handle reads the startup packet and dispatches an appropriate protocol handler for the connection.
func (s *Server) handle(ctx context.Context, conn net.Conn) (err error) {
	var hasTLS bool

	defer func() {
		cerr := conn.Close()
		if cerr != nil && err == nil {
			if hasTLS {
				if nerr, ok := cerr.(net.Error); ok && nerr.Timeout() {
					// TLS does this sometimes.
					return
				}
			}
			err = errors.Wrap(cerr, "closing connection")
		}
	}()

	if tcpconn, ok := conn.(*net.TCPConn); ok {
		// Postgres does not have any real mechanism for confirming that a connection is still alive.
		// Without this, a connection that breaks while idle would live indefinitely.
		// With a TCP keepalive, this should return an error after approximately 2 hours (depending on OS configuration).
		err := tcpconn.SetKeepAlive(true)
		if err != nil {
			return errors.Wrap(err, "enabling TCP keepalive")
		}
	}

	var startupDeadline time.Time
	if s.StartupTimeout > 0 {
		// Set deadline for processing the startup.
		startupDeadline = time.Now().Add(s.StartupTimeout)
		err = conn.SetDeadline(startupDeadline)
		if err != nil {
			return errors.Wrap(err, "setting deadline on protocol startup")
		}
	}

startup:
	// Read startup packet.
	var buf [4]byte
	_, err = io.ReadFull(conn, buf[:])
	if err != nil {
		return errors.Wrap(err, "reading startup message length")
	}
	size := binary.BigEndian.Uint32(buf[:])
	if size < 4 {
		return errors.Errorf("invalid startup packet length: %d bytes", size)
	}
	maxLen := s.MaxStartupSize
	if maxLen == 0 {
		maxLen = 1024 * 1024
	}
	if size > maxLen {
		return errors.Errorf("oversized startup frame of %d bytes (max: %d bytes)", size, maxLen)
	}
	data := make([]byte, size-4)
	_, err = io.ReadFull(conn, data)
	if err != nil {
		return errors.Wrap(err, "reading startup packet")
	}

	// Extract protocol ID.
	if len(data) < 4 {
		return errors.Errorf("startup packet is too small for protocol ID: %d bytes", len(data))
	}
	proto := Protocol(binary.BigEndian.Uint32(data))
	data = data[4:]

	if proto == ProtocolSSL {
		if s.TLSConfig != nil {
			// Upgrade the connection to TLS and renegotiate on the tunneled connection.
			_, err = conn.Write([]byte{'S'})
			if err != nil {
				return errors.Wrap(err, "sending SSL support confirmation")
			}
			conn = tls.Server(conn, s.TLSConfig)
			if s.StartupTimeout > 0 {
				err := conn.SetDeadline(startupDeadline)
				if err != nil {
					return errors.Wrap(err, "transferring startup deadline to TLS connection")
				}
			}
			hasTLS = true
			goto startup
		}

		// Inform the client that SSL is not available and try again.
		s.Logger.Debugf("client at %s requested a secure postgres connection but TLS is not configured", conn.RemoteAddr())
		_, err = conn.Write([]byte{'N'})
		if err != nil {
			return errors.Wrap(err, "sending SSL unsupported notification")
		}
		goto startup
	}

	if s.TLSConfig != nil && !hasTLS {
		// Reject the unsecured connection.
		return errors.Errorf("client at %s attempted to initiate an unsecured postgres conenction", conn.RemoteAddr())
	}

	switch proto {
	case ProtocolCancel:
		// Handle cancellation.
		return s.handleCancel(ctx, conn, data)
	default:
		// Handle regular postgres.
		return s.handleStandard(ctx, proto, conn, data)
	}
}

// parseParams parses a parameter list from a startup packet.
func parseParams(data []byte) (map[string]string, error) {
	params := make(map[string]string)
	for {
		idx := bytes.IndexByte(data, 0)
		switch idx {
		case 0:
			return params, nil
		case -1:
			return nil, errors.New("malformed startup parameter list")
		}

		key := string(data[:idx])
		data = data[idx+1:]

		idx = bytes.IndexByte(data, 0)
		if idx == -1 {
			return nil, errors.New("malformed startup parameter list")
		}
		val := string(data[:idx])
		data = data[idx+1:]

		params[key] = val
	}
}

// handleCancel handles cancel request connections.
func (s *Server) handleCancel(ctx context.Context, conn net.Conn, data []byte) error {
	if len(data) != 8 {
		return errors.New("malformed cancellation packet")
	}

	if s.CancellationManager == nil {
		return errors.New("cancellation is not configured")
	}

	pid := int32(binary.BigEndian.Uint32(data[:4]))
	key := int32(binary.BigEndian.Uint32(data[4:]))

	err := s.CancellationManager.Cancel(CancellationToken{PID: pid, Key: key})
	switch err {
	case nil:
	case ErrCancelledMissingConnection:
		// This is usually not a real error (race condition in the protocol).
		// This can happen if a client cancels a request and shuts down.
		s.Logger.Debugf("client at %v sent a mismatched cancellation token (is a load balancer misconfigured?)", conn.RemoteAddr())
	default:
		return err
	}

	return nil
}
func (s *Server) SendParameterStatus(w *message.WireWriter, param, value string, encoder *message.Encoder) error {
	msg, err := encoder.ParameterStatus(param, value)
	if err != nil {
		return err
	}
	err = w.WriteMessage(msg)
	if err != nil {
		return errors.Wrap(err, "sending parameter status")
	}
	return nil
}

type Result struct {
}
type PgType byte

// Constants used to indicated query interception
// Only pgPassOn is allowed to be processed in Featurebase query handling
const (
	pgPassOn         PgType = 'x'
	pgBackendPid     PgType = 'a'
	pgVersion        PgType = 'b'
	pgCountType      PgType = 'c'
	pgQueryTime      PgType = 'd'
	pgTerminate      PgType = 'e'
	pgEmpty          PgType = 'f'
	pgSetApplication PgType = 'g'
	pgSelect1        PgType = 'h'
	pgSchema         PgType = 'i'
	pgBegin          PgType = 'j'
	pgTypeLen        PgType = 'k'
)

type Portal struct {
	Name         string
	Writer       *message.WireWriter
	commands     []message.Message
	Encoder      *message.Encoder
	mapper       *sql.Mapper
	sql          string
	pgspecial    PgType
	pid          int32
	queryStart   time.Time
	server       *Server
	cancelNotify <-chan struct{}
}

func (p *Portal) Reset() {
	p.Name = ""
	p.sql = ""
	p.pgspecial = pgPassOn
	p.commands = p.commands[:0]
}

func (p *Portal) Bind() {
	p.Add(message.BindComplete)
}

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"))
	foundPQL := lookPQL.FindStringSubmatch(queryStr)
	if len(foundPQL) > 0 {

		p.sql = foundPQL[0]
		p.Name = "PQL"
		p.pgspecial = pgPassOn
		p.Add(message.ParseOK)
		return
	}
	if strings.Contains(queryStr, "EXTRACT") {
		// had to add this hack because the vitis parser doesn't handle...
		/*
		  	SELECT pid as id,
		   	 query as stmt,
		     EXTRACT(seconds from query_start - NOW()) as elapsed_time
		     FROM pg_stat_activity
		     WHERE usename='docker'`
		*/
		p.pgspecial = pgQueryTime
		p.Name = "SELECT"
		p.sql = queryStr
		p.Add(message.ParseOK)
		return
	}

	if len(queryStr) > 2 {
		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(queryStr, POSTGRESLENSQL) {
			p.pgspecial = pgTypeLen
			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
						}
					}
				}
			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 {
		p.pgspecial = pgEmpty
	}
	p.Add(message.ParseOK)
}
func (p *Portal) Describe() {
	// Placeholder should we need to handle the Decribe request
}

func (p *Portal) Execute() (shouldTerminate bool, queryReady bool, err error) {
	queryReady = true
	switch p.pgspecial {
	case pgBackendPid:
		rowDescription, e := p.Encoder.EncodeColumn("pg_backend_pid", int32(23), 4)
		if e != nil {
			err = e
			return
		}
		p.Add(rowDescription)
		pid := fmt.Sprintf("%v", p.pid)
		dataRow, _ := p.Encoder.TextRow(pid)
		p.Add(dataRow)
		//needs data row with cancel token
	case pgVersion:
		rowDescription, e := p.Encoder.EncodeColumn("version", int32(25), -1)
		if e != nil {
			err = e
			return
		}
		p.Add(rowDescription)
		mesg := fmt.Sprintf("PostgresSQL 13.0 (molecula.%v)", p.server.QueryHandler.Version())
		dataRow, _ := p.Encoder.TextRow(mesg)
		p.Add(dataRow)
	case pgSelect1:
		rowDescription, e := p.Encoder.EncodeColumn("?column?", int32(23), 4)
		if e != nil {
			err = e
			return
		}
		p.Add(rowDescription)
		dataRow, _ := p.Encoder.TextRow("1")
		p.Add(dataRow)
	case pgCountType:
		//need to block
		<-p.server.lookerChannel
		rowDescription, e := p.Encoder.EncodeColumn("count", int32(20), 8)
		if err != nil {
			err = e
			return
		}
		p.Add(rowDescription)
		errorResponse, _ := p.Encoder.Error(
			message.NoticeField{
				Type: message.NoticeFieldSeverity,
				Data: "FATAL",
			},
			message.NoticeField{
				Type: message.NoticeFieldMessage,
				Data: "terminating connection due to administrator command",
			},
			message.NoticeField{
				Type: message.NoticeFieldCode,
				Data: "57P01",
			},
		)
		p.Add(errorResponse)
		e = p.Sync() //send and Reset

		if e != nil {
			err = e
			return
		}
		return true, queryReady, nil
	case pgQueryTime:
		// need to return something so that the id can be queried
		//need to return  SELECT pid as id, query as stmt, EXTRACT(seconds from query_start - NOW()) as elapsed_time          FROM pg_stat_activity
		//seems like we need a map of pids to querys
		e := p.server.dumpPortalsTo(p)
		if e != nil {
			err = e
			return
		}

	case pgTerminate:
		//note just have 1 lock that blocks all who try to count the activities
		//TODO (twg) lock this
		close(p.server.lookerChannel)                //release all the other blockers and allow them to terminate
		p.server.lookerChannel = make(chan struct{}) //create a new one just in case
		// i think it needs to return boolean true
		rowDescription, e := p.Encoder.EncodeColumn("pg_terminate_backend", int32(16), 1)
		if e != nil {
			err = e
			return
		}
		p.Add(rowDescription)
		dataRow, _ := p.Encoder.TextRow("t")
		p.Add(dataRow)
		commandComplete, e := p.Encoder.CommandComplete("SELECT 1")
		if e != nil {
			err = e
			return
		}
		p.Add(commandComplete)
		e = p.Sync()
		if e != nil {
			err = e
			return
		}
		return false, queryReady, nil
	case pgEmpty:
		p.Add(message.NoData)
		p.Add(message.EmptyQueryResponse)
		e := p.Sync()
		if e != nil {
			err = e
			return
		}
		queryReady = true
		return false, queryReady, nil
	case pgSetApplication:
		//needs to add/send status
		msg, e := p.Encoder.ParameterStatus("application_name", "PostgreSQL JDBC Driver")
		if e != nil {
			err = e
			return
		}
		p.Add(msg)
	case pgSchema:
		parts := []message.SimpleColumn{
			{
				Name:    "table_schema",
				Typeid:  int32(19),
				Typelen: 64,
			},
			{
				Name:    "table_name",
				Typeid:  int32(19),
				Typelen: 64,
			},
		}
		rowDescription, e := p.Encoder.EncodeColumns(parts...)
		if e != nil {
			err = e
			return
		}
		p.Add(rowDescription)
		e = p.HandleSchema()
		if e != nil {
			err = e
			return
		}

	case pgPassOn:
		query := SimpleQuery(p.sql)
		e := p.server.handleQuery(p, query, p.cancelNotify)
		if e != nil {
			err = e
			return
		}
		return

	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
	if len(p.Name) > 0 { //only send command complete for those that have names
		message, _ := p.Encoder.CommandComplete(p.Name)
		p.Add(message)
	}
	return
}

// handleStandard handles a connection in the standard postgres wire protocol.
func (p *Portal) Sync() error {
	for _, m := range p.commands {
		err := p.Writer.WriteMessage(m)
		if err != nil {
			return err
		}
	}
	p.Writer.Flush()
	p.Reset()
	return nil
}
func (p *Portal) Add(m message.Message) {
	cp := message.Message{Type: m.Type, Data: make([]byte, len(m.Data))}
	copy(cp.Data, m.Data)
	p.commands = append(p.commands, cp)
}

func (p *Portal) DumpComands() {

	for _, m := range p.commands {
		m.Dump("DUMPING:")
	}
}
func (p *Portal) HandleSchema() error {
	return p.server.QueryHandler.HandleSchema(context.Background(), p)
}

func (p *Portal) WriteMessage(m message.Message) error {
	p.Add(m)
	return nil
}
func (p *Portal) Flush() error {
	return nil
}

// handleStandard handles a connection in the standard postgres wire protocol.
// The client is responsible for closing the connection when this finishes.
func (s *Server) handleStandard(ctx context.Context, proto Protocol, conn net.Conn, data []byte) error {
	// Wait for helper goroutines to finish.
	var wg sync.WaitGroup
	defer wg.Wait()

	// Set up context.
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	// Check the major version.
	if proto.Major() != ProtocolSupported.Major() {
		return errors.Errorf("unsupported protocol %v", proto)
	}

	// Parse the parameters bundled in the startup packet.
	params, err := parseParams(data)
	if err != nil {
		return errors.Wrap(err, "parsing parameters")
	}
	if user, ok := params["user"]; ok {
		// Log the connection.
		s.Logger.Debugf("new postgres connection from user %q at %v", user, conn.RemoteAddr())
	} else {
		// We do not use this much yet, but the wire protocol says that it is required.
		return errors.New("missing username")
	}

	// Set up message input and output.
	// Set up a reader that will preempt the connection when the context is canceled.
	ir := idleReader{
		conn:    conn,
		timeout: s.ReadTimeout,
	}
	wg.Add(1)
	go func() {
		defer wg.Done()

		<-ctx.Done()

		ir.preempt() //nolint:errcheck
	}()

	// Clear the startup deadline.
	err = conn.SetDeadline(time.Time{})
	if err != nil {
		return err
	}

	// Set up a message reader with buffering.
	rbuf := bufio.NewReader(&ir)
	r := message.NewWireReader(rbuf)

	// Set up a writer on the connection.
	var ww io.Writer = conn
	if s.WriteTimeout != 0 {
		// Apply the write timeout.
		ww = &timeoutWriter{
			conn:    conn,
			timeout: s.WriteTimeout,
		}
	}

	// Set up a message writer with buffering.
	w := message.NewWireWriter(bufio.NewWriter(ww))

	var encoder message.Encoder
	if proto.Minor() > ProtocolSupported.Minor() {
		// Negotiate the version down.
		s.Logger.Debugf("client requested unsupported protocol version %v; attempting to downgrade to %v", proto, ProtocolSupported)
		msg, err := encoder.NegotiateProtocolVersion(int32(ProtocolSupported.Minor()))
		if err != nil {
			return errors.Wrap(err, "negotiating version")
		}
		err = w.WriteMessage(msg)
		if err != nil {
			return errors.Wrap(err, "negotiating version")
		}
	}

	// TODO: real auth
	err = w.WriteMessage(message.AuthenticationOK)
	if err != nil {
		return errors.Wrap(err, "sending authentication confirmation")
	}
	err = s.SendParameterStatus(w, "application_name", "", &encoder)
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "client_encoding", "UTF8", &encoder)
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "DateStyle", "ISO, MDY", &encoder)
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "integer_datetimes", "on", &encoder)
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "IntervalStyle", "postgres", &encoder)
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "is_superuser", "on", &encoder)
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "server_encoding", "UTF8", &encoder)
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "server_version", PgServerVersion, &encoder)
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	/*
		-PARAMETER STATUS name='TimeZone', value='GMT'
	*/
	err = s.SendParameterStatus(w, "session_authorization", "docker", &encoder) //TODO(twg) figure out valid values here
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "standard_conforming_strings", "on", &encoder) //TODO(twg) figure out valid values here
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}
	err = s.SendParameterStatus(w, "TimeZone", "GMT", &encoder) //TODO(twg) figure out valid values here
	if err != nil {
		return errors.Wrap(err, "sending parameter status server version")
	}

	var cancelNotify <-chan struct{}
	var pid int32
	if s.CancellationManager != nil {
		notify, cancel, token, err := s.CancellationManager.Token()
		if err != nil {
			return errors.Wrap(err, "setting up cancellation")
		}
		defer cancel()

		msg, err := encoder.BackendKeyData(token.PID, token.Key)
		pid = token.PID
		if err != nil {
			return errors.Wrap(err, "encoding cancellation key data")
		}
		err = w.WriteMessage(msg)
		if err != nil {
			return errors.Wrap(err, "sending cancellation key data")
		}
		cancelNotify = notify
	}

	var queryReady bool
	portal := &Portal{
		Writer:       w,
		Encoder:      &encoder,
		commands:     make([]message.Message, 0),
		mapper:       sql.NewMapper(),
		pid:          pid,
		server:       s,
		cancelNotify: cancelNotify,
	}
	s.addPortal(portal)
	defer s.removePortal(portal)
	//mapper.Logger = logger
	for {
		if !queryReady {
			// Indicate that we are ready for a query.
			portal.sql = ""
			msg, err := encoder.ReadyForQuery(message.TransactionStatusIdle)
			if err != nil {
				return errors.Wrap(err, "sending query ready status")
			}
			err = w.WriteMessage(msg)
			if err != nil {
				return errors.Wrap(err, "sending query ready status")
			}

			// Flush the write buffer so that the client can respond.
			err = w.Flush()
			if err != nil {
				return errors.Wrap(err, "flushing status")
			}

			if rbuf.Buffered() == 0 {
				// Put the connection into idle mode.
				err = ir.setIdle()
				if err != nil {
					return errors.Wrap(err, "setting idle mode")
				}
			} else {
				// If the client follows the spec, then it should not have sent anything more.
				// However, it seems that no clients completely follow the spec, so we shouldn't rely on anything that isn't entirely straightforward.
				s.Logger.Debugf("postgres client sent additional data without waiting for completion")
			}
			queryReady = true
		}

		// Read the next packet.
		msg, err := r.ReadMessage()
		if err != nil {
			if err == errPreempted {
				// The server is shutting down.
				return errors.Wrap(s.handleShutdown(
					conn, w, &encoder,

					message.NoticeField{
						Type: message.NoticeFieldSeverity,
						Data: "ERROR",
					},
					message.NoticeField{
						Type: message.NoticeFieldMessage,
						Data: "server shutting down",
					},
					message.NoticeField{
						Type: message.NoticeFieldHint,
						Data: "This is normal. This message is sent when a server is shutting down and terminating its connections.",
					},
				), "processing connection shutdown")
			}

			return err
		}

		switch msg.Type {
		case message.TypeTermination:
			// We are done.
			return w.Flush()

		case message.TypeParse:
			portal.Parse(msg.Data)
		case message.TypeBind:
			portal.Bind()
		case message.TypeExecute:
			term, qr, err := portal.Execute()
			if err != nil {
				return err
			}
			if term {
				return w.Flush()
			}
			queryReady = qr
		case message.TypeSync:
			err := portal.Sync()
			if err != nil {
				return err
			}
			queryReady = false
		case message.TypeSimpleQuery:
			queryReady = false
			query := SimpleQuery(strings.TrimSuffix(string(msg.Data), "\x00"))

			// Execute the query.
			err := s.handleQuery(w, query, cancelNotify)
			if err != nil {
				return err
			}
		case message.TypeDescribe:
			portal.Describe()
		case message.TypeClose:
			return w.Flush()
		default:
			// The message is not supported yet.
			// Send an error.
			s.Logger.Errorf("unrecognized postgres packet %v", msg)
			msg, err = encoder.Error(
				message.NoticeField{
					Type: message.NoticeFieldSeverity,
					Data: "ERROR",
				},
				message.NoticeField{
					Type: message.NoticeFieldMessage,
					Data: fmt.Sprintf("unrecognized message type %q", msg.Type),
				},
				message.NoticeField{
					Type: message.NoticeFieldDetail,
					Data: "message body:" + hex.Dump(msg.Data),
				},
			)
			if err != nil {
				return errors.Wrap(err, "sending unrecognized message error")
			}
			err = w.WriteMessage(msg)
			if err != nil {
				return errors.Wrap(err, "sending unrecognized message error")
			}
			err = w.Flush()
			if err != nil {
				return errors.Wrap(err, "sending unrecognized message error")
			}
		}
	}
}
func (s *Server) addPortal(p *Portal) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.portals = append(s.portals, p)
}
func (s *Server) removePortal(p *Portal) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i, portal := range s.portals {
		if portal.pid == p.pid {
			//remove i
			s.portals = append(s.portals[:i], s.portals[i+1:]...)
			return
		}

	}

}
func (s *Server) dumpPortalsTo(p *Portal) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	//need to add the descrition for the 3 fields
	//	<-:-ROW DESCRIPTION: num_fields=3
	//---[Field 01]---  name='id'  type=23  type_len=4  type_mod=4294967295  relid=12250  attnum=3  format=0
	//---[Field 02]---  name='stmt'  type=25  type_len=65535  type_mod=4294967295  relid=12250  attnum=20  format=0  -
	//--[Field 03]---  name='elapsed_time'  type=701  type_len=8  type_mod=4294967295  relid=0  attnum=0  format=0
	parts := []message.SimpleColumn{
		{
			Name:    "id",
			Typeid:  int32(23),
			Typelen: 4,
		},
		{
			Name:    "stmt",
			Typeid:  int32(25),
			Typelen: -1,
		},
		{
			Name:    "elapsed_time",
			Typeid:  int32(701),
			Typelen: 8,
		},
	}
	rowDescription, err := p.Encoder.EncodeColumns(parts...)
	if err != nil {
		return err
	}
	p.Add(rowDescription)

	for _, portal := range s.portals {
		dataRow, err := p.Encoder.TextRow(
			fmt.Sprintf("%v", portal.pid),
			portal.sql,
			fmt.Sprintf("%v", time.Since(portal.queryStart).Seconds()))
		if err != nil {
			return err
		}
		//if sql == "" need to put in a null record
		//need to add the dararow
		//also need to figure out null types
		p.Add(dataRow)
	}
	return nil
}

// handleQuery processes a single query on a connection.
func (s *Server) handleQuery(w message.Writer, query Query, cancelNotify <-chan struct{}) error {
	// Configure cancellation.
	// This is not the connection context, since we want the request to finish safely before connection shutdown.
	ctx := context.Background()
	if cancelNotify != nil {
		defer func() {
			// Flush any cancel notifications.
			// This works on a best-effort basis.
			// It is still entirely possible that the cancel notification may be delivered to the next request.
			// Regardless of what we do, we either get false positives or false negatives.
			// This code chooses false positives.
			for len(cancelNotify) > 0 {
				<-cancelNotify
			}
		}()

		var wg sync.WaitGroup
		defer wg.Add(1)

		var cancel context.CancelFunc
		ctx, cancel = context.WithCancel(ctx)
		defer cancel()

		wg.Add(1)
		go func() {
			defer wg.Done()

			select {
			case <-ctx.Done():
			case <-cancelNotify:
				cancel()
			}
		}()
	}

	// Set up a result writer.
	// SELECT is used as a default tag, which seems to be handled decently by clients.
	// The encoder is intentionally not re-used because its buffer may be huge.
	qwriter := &queryResultWriter{
		w:   w,
		te:  s.TypeEngine,
		tag: "SELECT",
	}
	// Dispatch the query handler.
	qerr := s.QueryHandler.HandleQuery(ctx, qwriter, query)
	if qerr != nil {
		// There was an error in processing the query.
		// Send the error back to the client and keep going.
		s.Logger.Debugf("failed to execute query %q: %v", query, qerr)
		msg, err := qwriter.enc.GoError(qerr)
		if err != nil {
			return errors.Wrap(err, "failed to send query error to client")
		}
		err = w.WriteMessage(msg)
		if err != nil {
			return errors.Wrap(err, "failed to send query error to client")
		}
	} else {
		if !qwriter.wroteHeaders {
			// The handler did not write headers.
			// Write back an empty set of headers.
			err := qwriter.WriteHeader()
			if err != nil {
				return errors.Wrap(err, "sending empty column headers")
			}
		}
		// The query completed normally.
		// Notify the client of completion.
		msg, err := qwriter.enc.CommandComplete(qwriter.tag)
		if err != nil {
			return errors.Wrap(err, "sending command completion notification")
		}
		err = w.WriteMessage(msg)
		if err != nil {
			return errors.Wrap(err, "sending command completion notification")
		}
	}

	// The data will be flushed after we write back the "ready for query" state.
	return nil
}

func (s *Server) handleShutdown(conn net.Conn, w message.Writer, encoder *message.Encoder, notice ...message.NoticeField) error {
	var wg sync.WaitGroup
	defer wg.Wait()

	// Try to send a message to the client before closing the connection.
	msg, err := encoder.Error(notice...)
	if err != nil {
		return errors.Wrap(err, "generating shutdown notification")
	}

	if s.WriteTimeout == 0 {
		// The client is likely to not listen for incoming messages.
		// Force a write timeout to ensure that this terminates.
		err := conn.SetWriteDeadline(time.Now().Add(time.Second))
		if err != nil {
			return errors.Wrap(err, "setting shutdown write deadline")
		}
	}

	// The client may be waiting on a write, so we need to drain the incoming data stream.
	err = conn.SetReadDeadline(time.Time{})
	if err != nil {
		return errors.Wrap(err, "clearing read deadline for shutdown")
	}
	defer conn.SetReadDeadline(time.Now()) //nolint:errcheck
	wg.Add(1)
	go func() {
		defer wg.Done()

		io.Copy(io.Discard, conn) //nolint:errcheck
	}()

	// Attempt to send the shutdown notification.
	// This will fail under many scenarios, as the client is not necessarily reading.
	err = w.WriteMessage(msg)
	if err != nil {
		return nil
	}
	w.Flush()

	return nil
}