Add pgfortune example

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
+```

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)
+			}
+		}()
+	}
+}

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()
+}