diff --git a/README.md b/README.md index 3491780e..565b3efd 100644 --- a/README.md +++ b/README.md @@ -5,4 +5,8 @@ Package pgproto3 is a encoder and decoder of the PostgreSQL wire protocol version 3. +pgproto3 can be used as a foundation for PostgreSQL drivers, proxies, mock servers, load balancers and more. + +See example/pgfortune for a playful example of a fake PostgreSQL server. + Extracted from original implementation in https://github.com/jackc/pgx. diff --git a/example/pgfortune/README.md b/example/pgfortune/README.md new file mode 100644 index 00000000..c181c38a --- /dev/null +++ b/example/pgfortune/README.md @@ -0,0 +1,53 @@ +# pgfortune + +pgfortune is a mock PostgreSQL server that responds to every query with a fortune. + +## Installation + +Install `fortune` and `cowsay`. They should be available in any Unix package manager (apt, yum, brew, etc.) + +``` +go get -u github.com/jackc/pgproto3/example/pgfortune +``` + +## Usage + +``` +$ pgfortune +``` + +By default pgfortune listens on 127.0.0.1:15432 and responds to queries with `fortune | cowsay -f elephant`. These are +configurable with the `listen` and `response-command` arguments respectively. + +While `pgfortune` is running connect to it with `psql`. + +``` +$ psql -h 127.0.0.1 -p 15432 +Timing is on. +Null display is "∅". +Line style is unicode. +psql (11.5, server 0.0.0) +Type "help" for help. + +jack@127.0.0.1:15432 jack=# select foo; + fortune +───────────────────────────────────────────── + _________________________________________ ↵ + / Ships are safe in harbor, but they were \↵ + \ never meant to stay there. /↵ + ----------------------------------------- ↵ + \ /\ ___ /\ ↵ + \ // \/ \/ \\ ↵ + (( O O )) ↵ + \\ / \ // ↵ + \/ | | \/ ↵ + | | | | ↵ + | | | | ↵ + | o | ↵ + | | | | ↵ + |m| |m| ↵ + +(1 row) + +Time: 28.161 ms +``` diff --git a/example/pgfortune/main.go b/example/pgfortune/main.go new file mode 100644 index 00000000..45970eb3 --- /dev/null +++ b/example/pgfortune/main.go @@ -0,0 +1,48 @@ +package main + +import ( + "flag" + "fmt" + "log" + "net" + "os" + "os/exec" +) + +var options struct { + listenAddress string + responseCommand string +} + +func main() { + flag.Usage = func() { + fmt.Fprintf(os.Stderr, "usage: %s [options]\n", os.Args[0]) + flag.PrintDefaults() + } + + flag.StringVar(&options.listenAddress, "listen", "127.0.0.1:15432", "Listen address") + flag.StringVar(&options.responseCommand, "response-command", "fortune | cowsay -f elephant", "Command to execute to generate query response") + flag.Parse() + + ln, err := net.Listen("tcp", options.listenAddress) + if err != nil { + log.Fatal(err) + } + + for { + conn, err := ln.Accept() + if err != nil { + log.Fatal(err) + } + + b := NewPgFortuneBackend(conn, func() ([]byte, error) { + return exec.Command("sh", "-c", options.responseCommand).CombinedOutput() + }) + go func() { + err := b.Run() + if err != nil { + log.Println(err) + } + }() + } +} diff --git a/example/pgfortune/server.go b/example/pgfortune/server.go new file mode 100644 index 00000000..777192a6 --- /dev/null +++ b/example/pgfortune/server.go @@ -0,0 +1,104 @@ +package main + +import ( + "fmt" + "net" + + "github.com/jackc/pgproto3/v2" +) + +type PgFortuneBackend struct { + backend *pgproto3.Backend + conn net.Conn + responder func() ([]byte, error) +} + +func NewPgFortuneBackend(conn net.Conn, responder func() ([]byte, error)) *PgFortuneBackend { + backend := pgproto3.NewBackend(pgproto3.NewChunkReader(conn), conn) + + connHandler := &PgFortuneBackend{ + backend: backend, + conn: conn, + responder: responder, + } + + return connHandler +} + +func (p *PgFortuneBackend) Run() error { + defer p.Close() + + err := p.handleStartup() + if err != nil { + return err + } + + for { + msg, err := p.backend.Receive() + if err != nil { + return fmt.Errorf("error receiving message: %w", err) + } + + switch msg.(type) { + case *pgproto3.Query: + response, err := p.responder() + if err != nil { + return fmt.Errorf("error generating query response: %w", err) + } + + buf := (&pgproto3.RowDescription{Fields: []pgproto3.FieldDescription{ + { + Name: []byte("fortune"), + TableOID: 0, + TableAttributeNumber: 0, + DataTypeOID: 25, + DataTypeSize: -1, + TypeModifier: -1, + Format: 0, + }, + }}).Encode(nil) + buf = (&pgproto3.DataRow{Values: [][]byte{response}}).Encode(buf) + buf = (&pgproto3.CommandComplete{CommandTag: []byte("SELECT 1")}).Encode(buf) + buf = (&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(buf) + _, err = p.conn.Write(buf) + if err != nil { + return fmt.Errorf("error writing query response: %w", err) + } + case *pgproto3.Terminate: + return nil + default: + return fmt.Errorf("received message other than Query from client: %#v", msg) + } + } +} + +func (p *PgFortuneBackend) handleStartup() error { + startupMessage, err := p.backend.ReceiveStartupMessage() + if err != nil { + return fmt.Errorf("error receiving startup message: %w", err) + } + + switch startupMessage.(type) { + case *pgproto3.StartupMessage: + buf := (&pgproto3.AuthenticationOk{}).Encode(nil) + buf = (&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(buf) + _, err = p.conn.Write(buf) + if err != nil { + return fmt.Errorf("error sending ready for query: %w", err) + } + case *pgproto3.SSLRequest: + _, err = p.conn.Write([]byte("N")) + if err != nil { + return fmt.Errorf("error sending deny SSL request: %w", err) + } + return p.handleStartup() + default: + return fmt.Errorf("unknown startup message: %#v", startupMessage) + } + + return nil +} + +func (p *PgFortuneBackend) Close() error { + return p.conn.Close() +}