pizzasql/pkg/pgserver/connection_test.go

package pgserver

import (
	"bufio"
	"bytes"
	"net"
	"testing"

	"github.com/danfragoso/pizzasql-next/pkg/executor"
	"github.com/danfragoso/pizzasql-next/pkg/lexer"
	"github.com/danfragoso/pizzasql-next/pkg/parser"
)

func TestBindQuery(t *testing.T) {
	query, err := bindQuery(
		"SELECT $1, $2, $3, $4, '$5'",
		[]boundParameter{
			{value: []byte("O'Reilly")},
			{value: []byte("42")},
			{null: true},
			{value: []byte("true")},
		},
	)
	if err != nil {
		t.Fatal(err)
	}
	want := "SELECT 'O''Reilly', 42, NULL, TRUE, '$5'"
	if query != want {
		t.Fatalf("bound query = %q, want %q", query, want)
	}
}

func TestBindQueryRequiresEveryParameter(t *testing.T) {
	if _, err := bindQuery("SELECT $2", []boundParameter{{value: []byte("one")}}); err == nil {
		t.Fatal("expected missing parameter error")
	}
}

func parseStmt(t *testing.T, sql string) parser.Statement {
	t.Helper()
	l := lexer.New(sql)
	p := parser.New(l)
	stmt, err := p.Parse()
	if err != nil {
		t.Fatalf("parse %q: %v", sql, err)
	}
	return stmt
}

func TestGetCommandTagSavepointKeepsTransaction(t *testing.T) {
	c := &Connection{txStatus: TxStatusIdle}
	tag := c.getCommandTag(parseStmt(t, "SAVEPOINT sp1"), executor.NewResult("SAVEPOINT"))
	if tag != "SAVEPOINT" {
		t.Fatalf("tag = %q, want %q", tag, "SAVEPOINT")
	}
	if c.txStatus != TxStatusInBlock {
		t.Fatalf("txStatus = %c, want %c", c.txStatus, TxStatusInBlock)
	}
}

func TestGetCommandTagReleaseKeepsTransaction(t *testing.T) {
	c := &Connection{txStatus: TxStatusInBlock}
	tag := c.getCommandTag(parseStmt(t, "RELEASE SAVEPOINT sp1"), executor.NewResult("RELEASE"))
	if tag != "RELEASE" {
		t.Fatalf("tag = %q, want %q", tag, "RELEASE")
	}
	if c.txStatus != TxStatusInBlock {
		t.Fatalf("txStatus = %c, want %c", c.txStatus, TxStatusInBlock)
	}
}

func TestGetCommandTagRollbackToSavepointKeepsTransaction(t *testing.T) {
	c := &Connection{txStatus: TxStatusFailed}
	tag := c.getCommandTag(parseStmt(t, "ROLLBACK TO SAVEPOINT sp1"), executor.NewResult("ROLLBACK"))
	if tag != "ROLLBACK" {
		t.Fatalf("tag = %q, want %q", tag, "ROLLBACK")
	}
	if c.txStatus != TxStatusInBlock {
		t.Fatalf("txStatus = %c, want %c", c.txStatus, TxStatusInBlock)
	}
}

func TestGetCommandTagFullRollbackEndsTransaction(t *testing.T) {
	c := &Connection{txStatus: TxStatusFailed}
	tag := c.getCommandTag(parseStmt(t, "ROLLBACK"), executor.NewResult("ROLLBACK"))
	if tag != "ROLLBACK" {
		t.Fatalf("tag = %q, want %q", tag, "ROLLBACK")
	}
	if c.txStatus != TxStatusIdle {
		t.Fatalf("txStatus = %c, want %c", c.txStatus, TxStatusIdle)
	}
}

func newTestConnection(t *testing.T) (*Connection, net.Conn) {
	t.Helper()
	server, client := net.Pipe()
	t.Cleanup(func() {
		server.Close()
		client.Close()
	})
	c := &Connection{
		conn:       server,
		reader:     bufio.NewReader(server),
		writer:     bufio.NewWriter(server),
		params:     map[string]string{"user": "tester"},
		statements: make(map[string]*preparedStatement),
		portals:    make(map[string]*portal),
		txStatus:   TxStatusIdle,
		quiet:      true,
	}
	return c, client
}

func runQuery(t *testing.T, c *Connection, client net.Conn, msg *Message) []*Message {
	t.Helper()
	done := make(chan error, 1)
	go func() { done <- c.handleQuery(msg) }()

	reader := bufio.NewReader(client)
	var msgs []*Message
	for {
		m, err := ReadMessage(reader)
		if err != nil {
			t.Fatalf("read response: %v", err)
		}
		msgs = append(msgs, m)
		if m.Type == MsgReadyForQuery {
			break
		}
	}
	if err := <-done; err != nil {
		t.Fatalf("handleQuery: %v", err)
	}
	return msgs
}

func errorCode(msg *Message) string {
	if msg == nil || msg.Type != MsgErrorResponse {
		return ""
	}
	data := msg.Data
	for i := 0; i < len(data); {
		if data[i] == 0 {
			break
		}
		field := data[i]
		i++
		end := bytes.IndexByte(data[i:], 0)
		if end < 0 {
			break
		}
		value := string(data[i : i+end])
		if field == ErrorFieldCode {
			return value
		}
		i += end + 1
	}
	return ""
}

func TestHandleQueryEmptyPayloadReturnsProtocolError(t *testing.T) {
	c, client := newTestConnection(t)
	msgs := runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte{}})
	if len(msgs) != 2 {
		t.Fatalf("got %d messages, want 2", len(msgs))
	}
	if code := errorCode(msgs[0]); code != ErrCodeProtocolViolation {
		t.Fatalf("error code = %q, want %q", code, ErrCodeProtocolViolation)
	}
}

func TestHandleQueryMissingNullTerminatorReturnsProtocolError(t *testing.T) {
	c, client := newTestConnection(t)
	msgs := runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte("SELECT 1")})
	if len(msgs) != 2 {
		t.Fatalf("got %d messages, want 2", len(msgs))
	}
	if code := errorCode(msgs[0]); code != ErrCodeProtocolViolation {
		t.Fatalf("error code = %q, want %q", code, ErrCodeProtocolViolation)
	}
}

func TestHandleQueryEmptyQueryReturnsEmptyQueryResponse(t *testing.T) {
	for _, data := range [][]byte{{0}, []byte(";\x00"), []byte("   \x00")} {
		c, client := newTestConnection(t)
		msgs := runQuery(t, c, client, &Message{Type: MsgQuery, Data: data})
		if len(msgs) != 2 {
			t.Fatalf("payload %q: got %d messages, want 2", data, len(msgs))
		}
		if msgs[0].Type != MsgEmptyQueryResponse {
			t.Fatalf("payload %q: first message type = %c, want %c", data, msgs[0].Type, MsgEmptyQueryResponse)
		}
	}
}

func TestHandleQueryFailedTransactionInterceptsSpecialQueries(t *testing.T) {
	queries := []string{
		"SELECT version()",
		"SELECT current_user",
		"SHOW server_version",
		"SELECT * FROM information_schema.tables",
	}
	for _, q := range queries {
		c, client := newTestConnection(t)
		c.txStatus = TxStatusFailed
		msgs := runQuery(t, c, client, &Message{Type: MsgQuery, Data: []byte(q + "\x00")})
		if len(msgs) != 2 {
			t.Fatalf("query %q: got %d messages, want 2", q, len(msgs))
		}
		if code := errorCode(msgs[0]); code != ErrCodeTransactionAborted {
			t.Fatalf("query %q: error code = %q, want %q", q, code, ErrCodeTransactionAborted)
		}
	}
}

func TestIsRollbackStatement(t *testing.T) {
	for _, sql := range []string{"ROLLBACK", "ROLLBACK TO SAVEPOINT sp1", "rollback;", "  rollback  ;"} {
		if !isRollbackStatement(sql) {
			t.Errorf("isRollbackStatement(%q) = false, want true", sql)
		}
	}
	for _, sql := range []string{"SELECT version()", "SHOW server_version", "SAVEPOINT sp1"} {
		if isRollbackStatement(sql) {
			t.Errorf("isRollbackStatement(%q) = true, want false", sql)
		}
	}
	for _, sql := range []string{"ROLLBACKBOGUS", "ROLLBACK; SELECT 1"} {
		if isRollbackStatement(sql) {
			t.Errorf("isRollbackStatement(%q) = true, want false", sql)
		}
	}
}