github.com/jackc/pgx/v5@v5.5.5/pgproto3/example/pgfortune/server.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  
     7  	"github.com/jackc/pgx/v5/pgproto3"
     8  )
     9  
    10  type PgFortuneBackend struct {
    11  	backend   *pgproto3.Backend
    12  	conn      net.Conn
    13  	responder func() ([]byte, error)
    14  }
    15  
    16  func NewPgFortuneBackend(conn net.Conn, responder func() ([]byte, error)) *PgFortuneBackend {
    17  	backend := pgproto3.NewBackend(conn, conn)
    18  
    19  	connHandler := &PgFortuneBackend{
    20  		backend:   backend,
    21  		conn:      conn,
    22  		responder: responder,
    23  	}
    24  
    25  	return connHandler
    26  }
    27  
    28  func (p *PgFortuneBackend) Run() error {
    29  	defer p.Close()
    30  
    31  	err := p.handleStartup()
    32  	if err != nil {
    33  		return err
    34  	}
    35  
    36  	for {
    37  		msg, err := p.backend.Receive()
    38  		if err != nil {
    39  			return fmt.Errorf("error receiving message: %w", err)
    40  		}
    41  
    42  		switch msg.(type) {
    43  		case *pgproto3.Query:
    44  			response, err := p.responder()
    45  			if err != nil {
    46  				return fmt.Errorf("error generating query response: %w", err)
    47  			}
    48  
    49  			buf := mustEncode((&pgproto3.RowDescription{Fields: []pgproto3.FieldDescription{
    50  				{
    51  					Name:                 []byte("fortune"),
    52  					TableOID:             0,
    53  					TableAttributeNumber: 0,
    54  					DataTypeOID:          25,
    55  					DataTypeSize:         -1,
    56  					TypeModifier:         -1,
    57  					Format:               0,
    58  				},
    59  			}}).Encode(nil))
    60  			buf = mustEncode((&pgproto3.DataRow{Values: [][]byte{response}}).Encode(buf))
    61  			buf = mustEncode((&pgproto3.CommandComplete{CommandTag: []byte("SELECT 1")}).Encode(buf))
    62  			buf = mustEncode((&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(buf))
    63  			_, err = p.conn.Write(buf)
    64  			if err != nil {
    65  				return fmt.Errorf("error writing query response: %w", err)
    66  			}
    67  		case *pgproto3.Terminate:
    68  			return nil
    69  		default:
    70  			return fmt.Errorf("received message other than Query from client: %#v", msg)
    71  		}
    72  	}
    73  }
    74  
    75  func (p *PgFortuneBackend) handleStartup() error {
    76  	startupMessage, err := p.backend.ReceiveStartupMessage()
    77  	if err != nil {
    78  		return fmt.Errorf("error receiving startup message: %w", err)
    79  	}
    80  
    81  	switch startupMessage.(type) {
    82  	case *pgproto3.StartupMessage:
    83  		buf := mustEncode((&pgproto3.AuthenticationOk{}).Encode(nil))
    84  		buf = mustEncode((&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(buf))
    85  		_, err = p.conn.Write(buf)
    86  		if err != nil {
    87  			return fmt.Errorf("error sending ready for query: %w", err)
    88  		}
    89  	case *pgproto3.SSLRequest:
    90  		_, err = p.conn.Write([]byte("N"))
    91  		if err != nil {
    92  			return fmt.Errorf("error sending deny SSL request: %w", err)
    93  		}
    94  		return p.handleStartup()
    95  	default:
    96  		return fmt.Errorf("unknown startup message: %#v", startupMessage)
    97  	}
    98  
    99  	return nil
   100  }
   101  
   102  func (p *PgFortuneBackend) Close() error {
   103  	return p.conn.Close()
   104  }
   105  
   106  func mustEncode(buf []byte, err error) []byte {
   107  	if err != nil {
   108  		panic(err)
   109  	}
   110  	return buf
   111  }