scriptoriumadmin/featurebase
1
0
Fork 0
mirror of https://github.com/featurebasedb/featurebase.git synced 2026-09-11 15:21:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #679 from niaow/pg

Browse source 
Add a postgres endpoint to pilosa
This commit is contained in:
Nia 2020-08-20 14:24:56 -04:00 • committed by GitHub
commit 272f6708a3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
18 changed files with 2476 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
ctl/server.go
View file

@ -88,4 +88,7 @@ func BuildServerFlags(cmd *cobra.Command, srv *server.Command) {
// Transactional storage engine
flags.StringVarP(&srv.Config.Txsrc, "tx", "", "", "transaction/storage to use: one of roaring, rbf, badger, rbf_roaring, roaring_rbf, badger_roaring, roaring_badger, badger_rbf, or rbf_badger (default roaring)")
// Postgres endpoint
flags.StringVar(&srv.Config.Postgres.Addr, "postgres.addr", "", "address to which to bind a postgres endpoint")
}

1
go.mod
View file

@ -20,6 +20,7 @@ require (
github.com/gorilla/handlers v1.3.0
github.com/gorilla/mux v1.7.0
github.com/hashicorp/memberlist v0.1.3
github.com/lib/pq v1.8.0
github.com/opentracing/opentracing-go v1.1.0
github.com/pelletier/go-toml v1.2.0
github.com/pkg/errors v0.8.1

2
go.sum
View file

@ -117,6 +117,8 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/lib/pq v1.8.0 h1:9xohqzkUwzR4Ga4ivdTcawVS89YSDVxXMa3xJX3cGzg=
github.com/lib/pq v1.8.0/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
github.com/magiconair/properties v1.8.0 h1:LLgXmsheXeRoUOBOjtwPQCWIYqM/LU1ayDtDePerRcY=
github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU=

249
pg/io.go Normal file
View file

@ -0,0 +1,249 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pg
import (
"net"
"sync"
"sync/atomic"
"time"
"unsafe"
"github.com/pkg/errors"
)
// timeoutWriter wraps a connection and implements io.Writer with a timeout for each write.
type timeoutWriter struct {
conn net.Conn
timeout time.Duration
}
func (w *timeoutWriter) Write(data []byte) (int, error) {
err := w.conn.SetWriteDeadline(time.Now().Add(w.timeout))
if err != nil {
return 0, err
}
return w.conn.Write(data)
}
// errPreempted is an error used to indicate preemption of an idle connection.
var errPreempted = errors.New("preempted during idle")
// idleState is an atomic state value used to track a preemptible connection.
type idleState uint32
const (
idleStateActive idleState = iota
idleStateIdle
idleStatePreempted
idleStatePendingPreemption
)
func (s *idleState) load() idleState {
return idleState(atomic.LoadUint32((*uint32)(unsafe.Pointer(s))))
}
func (s *idleState) cas(old, new idleState) bool {
return atomic.CompareAndSwapUint32((*uint32)(unsafe.Pointer(s)), uint32(old), uint32(new))
}
// idleReader is an io.Reader implementation on a preemptible network connection.
// The connection has 2 modes: "idle" and "active".
// While in idle mode, the connection has no timeout but can be preempted.
// While in active mode, the connection may have a read timeout but cannot be immediately preempted.
// When a read completes in idle mode, the connection returns to active mode.
// If the connection is preempted in active mode, the preemption will be deferred until the connection returns to idle mode.
// This also provides a read timeout.
type idleReader struct {
conn net.Conn
timeout time.Duration
state idleState
preemptMu sync.Mutex
}
// setIdle pushes the reader into idle mode.
// If a preemption is pending, it will be delivered on the next call to Read.
func (r *idleReader) setIdle() error {
// Clear the read deadline.
err := r.conn.SetReadDeadline(time.Time{})
if err != nil {
return errors.Wrap(err, "failed to clear deadline")
}
for {
// Transition to idle mode.
state := r.state.load()
var target idleState
switch state {
case idleStateActive:
// active -> idle
target = idleStateIdle
case idleStatePendingPreemption:
// pending preemption -> preempted
// Switching to idle mode activates the preemption.
target = idleStatePreempted
default:
panic("inconsistent state")
}
if r.state.cas(state, target) {
return nil
}
}
}
// preempt the reader.
// If the reader is not currently idle, the preemption will be delivered next time the connection enters idle mode.
// This does not wait until the preemption error is delivered.
func (r *idleReader) preempt() error {
r.preemptMu.Lock()
defer r.preemptMu.Unlock()
for {
state := r.state.load()
var target idleState
switch state {
case idleStateActive:
// active -> pending preemption
target = idleStatePendingPreemption
case idleStateIdle:
// idle -> preempted
target = idleStatePreempted
case idleStatePendingPreemption, idleStatePreempted:
// A preemption has already been delivered.
return nil
default:
panic("inconsistent state")
}
ok := r.state.cas(state, target)
if ok && target == idleStatePreempted {
// We have entered preemption mode.
// Preempt the current read on the connection.
return r.conn.SetReadDeadline(time.Now())
}
}
}
// Read from the connection.
func (r *idleReader) Read(data []byte) (int, error) {
var needsDeadlineReset bool
state := r.state.load()
switch state {
case idleStateActive, idleStatePendingPreemption:
// Connection is active.
// There is no need to worry about preemption.
if r.timeout != 0 {
// Apply a read timeout.
err := r.conn.SetReadDeadline(time.Now().Add(r.timeout))
if err != nil {
return 0, err
}
}
return r.conn.Read(data)
case idleStateIdle:
// Read, and handle preemption.
n, err := r.conn.Read(data)
if err != nil {
// Check if the error was caused by preemption.
state = r.state.load()
switch {
case state == idleStatePreempted && n != 0:
// Some data was read before the preemption was delivered.
// Re-activate and discard the error.
// Synchronize against the preempter.
// This is necessary to ensure that the cancellation deadline is cleared.
r.preemptMu.Lock()
defer r.preemptMu.Unlock()
// Re-activate the connection.
// No CAS loop is necessary since we are synchronized against preempters.
r.state = idleStatePendingPreemption
// The deadline may need to reset since the preempter may have changed it.
needsDeadlineReset = true
case state == idleStatePreempted:
// The read was preempted.
return 0, errPreempted
case state != idleStateIdle:
// No other states make sense here.
panic("inconsistent state")
default:
// No preemption was involved.
// It is just a regular network error.
return n, err
}
} else {
// The read went through.
// Exit from idle mode.
// Ideally, transition to active mode.
// However, a preemption may trigger while this is running.
for state == idleStateIdle {
if r.state.cas(idleStateIdle, idleStateActive) {
state = idleStateActive
break
}
state = r.state.load()
}
switch state {
case idleStateActive:
// The connection was reactivated normally.
case idleStatePreempted:
// The connection was preempted after the read completed.
// Defer the preemption and complete successfully.
// Synchronize against the preempter.
// This is necessary to ensure that the cancellation deadline is cleared.
r.preemptMu.Lock()
defer r.preemptMu.Unlock()
// Re-activate the connection.
// No CAS loop is necessary since we are synchronized against preempters.
r.state = idleStatePendingPreemption
// The deadline may need to reset since the preempter may have changed it.
needsDeadlineReset = true
default:
panic("inconsistent state")
}
}
if needsDeadlineReset && r.timeout == 0 {
// Clear the deadline.
err := r.conn.SetReadDeadline(time.Time{})
if err != nil {
return n, err
}
}
return n, nil
case idleStatePreempted:
// The connection is preempted.
return 0, errPreempted
default:
panic("inconsistent state")
}
}

124
pg/message/io.go Normal file
View file

@ -0,0 +1,124 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package message
import (
"bufio"
"encoding/binary"
"errors"
"io"
)
// Reader reads messages.
type Reader interface {
ReadMessage() (Message, error)
}
// Writer writes messages.
type Writer interface {
WriteMessage(Message) error
Flush() error
}
// WireReader reads messages in Postgres wire protocol format.
type WireReader struct {
buf []byte
r *bufio.Reader
scratch [4]byte
}
// ReadMessage reads a single message off of the wire.
// The returned message is only valid until the next read call, as the data buffer may be re-used.
func (r *WireReader) ReadMessage() (Message, error) {
t, err := r.r.ReadByte()
if err != nil {
return Message{}, err
}
_, err = r.r.Read(r.scratch[:])
if err != nil {
return Message{}, err
}
len := binary.BigEndian.Uint32(r.scratch[:4])
if len < 4 {
return Message{}, errors.New("invalid message length")
}
len -= 4
if cap(r.buf) < int(len) {
r.buf = make([]byte, len)
} else {
r.buf = r.buf[:len]
}
_, err = io.ReadFull(r.r, r.buf)
if err != nil {
return Message{}, err
}
return Message{
Type: Type(t),
Data: r.buf,
}, nil
}
var _ Reader = (*WireReader)(nil)
// NewWireReader returns a message reader that reads postgres wire protocol format.
func NewWireReader(r *bufio.Reader) *WireReader {
return &WireReader{r: r}
}
// ErrMessageTooBig is an error indicating that a message is too big to be sent or received.
var ErrMessageTooBig = errors.New("message is too big")
// WireWriter writes messages in Postgres wire protocol.
type WireWriter struct {
w *bufio.Writer
scratch [4]byte
}
// WriteMessage writes a message onto the wire.
func (w *WireWriter) WriteMessage(message Message) error {
if uint(len(message.Data))+4 >= 1<<31 {
return ErrMessageTooBig
}
err := w.w.WriteByte(byte(message.Type))
if err != nil {
return err
}
binary.BigEndian.PutUint32(w.scratch[:], uint32(len(message.Data))+4)
_, err = w.w.Write(w.scratch[:])
if err != nil {
return err
}
_, err = w.w.Write(message.Data)
return err
}
// Flush writes any buffered data to the underlying stream.
func (w *WireWriter) Flush() error {
return w.w.Flush()
}
var _ Writer = (*WireWriter)(nil)
// NewWireWriter returns a message writer that writes in postgres wire protocol format.
func NewWireWriter(w *bufio.Writer) *WireWriter {
return &WireWriter{w: w}
}

344
pg/message/message.go Normal file
View file

@ -0,0 +1,344 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package message
import (
"bytes"
"encoding/binary"
)
// Type is a byte indicating the type of a Postgres message.
type Type byte
const (
// TypeAuthentication is a message used to transfer authentication info.
TypeAuthentication Type = 'R'
// TypeReadyForQuery is a message used to indicate that the server is ready for another query.
TypeReadyForQuery Type = 'Z'
// TypeCommandComplete is a message used to indicate that a query has completed.
TypeCommandComplete Type = 'C'
// TypeError is an error message.
TypeError Type = 'E'
// TypeRowDescription is a message indicating the column types of the result rows from a query.
TypeRowDescription Type = 'T'
// TypeDataRow is a message with the contents of a single row.
TypeDataRow Type = 'D'
// TypeTermination is a message indicating a request to terminate a connection.
TypeTermination Type = 'X'
// TypeNegotiateProtocolVersion is a message used when a client attempts to connect with a newer minor version than the server supports.
TypeNegotiateProtocolVersion Type = 'v'
// TypeSimpleQuery is a simple query request.
TypeSimpleQuery Type = 'Q'
)
// AuthenticationOK is a message indicating that authentication has completed.
var AuthenticationOK = Message{
Type: TypeAuthentication,
Data: []byte{0, 0, 0, 0},
}
// Message is a Postgres message value.
type Message struct {
Type Type
Data []byte
}
// TransactionStatus is the current transaction state.
type TransactionStatus byte
const (
// TransactionStatusIdle indicates that there is no active transaction.
TransactionStatusIdle TransactionStatus = 'I'
// TransactionStatusActive indicates that the connection currently has an active transaction.
TransactionStatusActive TransactionStatus = 'T'
// TransactionStatusFailed indicates that the connection currently has a failed transaction.
TransactionStatusFailed TransactionStatus = 'E'
)
// Encoder encodes messages.
type Encoder struct {
buf bytes.Buffer
scratch [4]byte
}
func (e *Encoder) i16(i int16) error {
binary.BigEndian.PutUint16(e.scratch[:2], uint16(i))
_, err := e.buf.Write(e.scratch[:2])
return err
}
func (e *Encoder) i32(i int32) error {
binary.BigEndian.PutUint32(e.scratch[:], uint32(i))
_, err := e.buf.Write(e.scratch[:])
return err
}
// ReadyForQuery encodes a "ready for query" message.
func (e *Encoder) ReadyForQuery(status TransactionStatus) (Message, error) {
e.buf.Reset()
err := e.buf.WriteByte(byte(status))
if err != nil {
return Message{}, err
}
return Message{
Type: TypeReadyForQuery,
Data: e.buf.Bytes(),
}, nil
}
// CommandComplete encodes a command completion message.
func (e *Encoder) CommandComplete(tag string) (Message, error) {
e.buf.Reset()
_, err := e.buf.WriteString(tag)
if err != nil {
return Message{}, err
}
err = e.buf.WriteByte(0)
if err != nil {
return Message{}, err
}
return Message{
Type: TypeCommandComplete,
Data: e.buf.Bytes(),
}, nil
}
// NoticeFieldType indicates the type of a notice/error field.
// https://www.postgresql.org/docs/9.3/protocol-error-fields.html
type NoticeFieldType byte
const (
// NoticeFieldSeverity indicates the severity of a notice/error.
NoticeFieldSeverity NoticeFieldType = 'S'
// NoticeFieldMessage is a short human-readable error/notice message.
NoticeFieldMessage NoticeFieldType = 'M'
// NoticeFieldDetail is an optional extended description of the error.
NoticeFieldDetail NoticeFieldType = 'D'
// NoticeFieldHint is a suggestion of how to address the issue.
NoticeFieldHint NoticeFieldType = 'H'
)
// NoticeField is a field in an error or notice.
type NoticeField struct {
Type NoticeFieldType
Data string
}
func (e *Encoder) messageOrNotice(fields ...NoticeField) error {
for _, f := range fields {
err := e.buf.WriteByte(byte(f.Type))
if err != nil {
return err
}
_, err = e.buf.WriteString(f.Data)
if err != nil {
return err
}
err = e.buf.WriteByte(0)
if err != nil {
return err
}
}
return e.buf.WriteByte(0)
}
// Error encodes a Postgres error message.
func (e *Encoder) Error(fields ...NoticeField) (Message, error) {
e.buf.Reset()
err := e.messageOrNotice(fields...)
if err != nil {
return Message{}, err
}
return Message{
Type: TypeError,
Data: e.buf.Bytes(),
}, nil
}
// GoError creates a simple Postgres error message from a Go error value.
func (e *Encoder) GoError(err error) (Message, error) {
return e.Error(
NoticeField{
Type: NoticeFieldSeverity,
Data: "ERROR",
},
NoticeField{
Type: NoticeFieldMessage,
Data: err.Error(),
},
)
}
// ColumnDescription is a description of a data column.
type ColumnDescription struct {
Name string
TableID int32 //either a table/col id or 0
FieldID int16 //either a table/col id or 0
TypeID int32 //field type
TypeLen int16 //size in bytes of field
TypeModifier int32 //type modifer?
Mode int16 //0=text 1=binary
}
// RowDescription describes the response rows from a query.
func (e *Encoder) RowDescription(cols ...ColumnDescription) (Message, error) {
if len(cols) >= 1<<15 {
return Message{}, ErrMessageTooBig
}
e.buf.Reset()
err := e.i16(int16(len(cols)))
if err != nil {
return Message{}, nil
}
for _, col := range cols {
_, err := e.buf.WriteString(col.Name)
if err != nil {
return Message{}, err
}
err = e.buf.WriteByte(0)
if err != nil {
return Message{}, err
}
err = e.i32(col.TableID)
if err != nil {
return Message{}, err
}
err = e.i16(col.FieldID)
if err != nil {
return Message{}, err
}
err = e.i32(col.TypeID)
if err != nil {
return Message{}, err
}
err = e.i16(col.TypeLen)
if err != nil {
return Message{}, err
}
err = e.i32(col.TypeModifier)
if err != nil {
return Message{}, err
}
err = e.i16(col.Mode)
if err != nil {
return Message{}, err
}
}
return Message{
Type: TypeRowDescription,
Data: e.buf.Bytes(),
}, nil
}
// TextRow encodes a data row in textual format.
func (e *Encoder) TextRow(row ...string) (Message, error) {
if len(row) >= 1<<15 {
return Message{}, ErrMessageTooBig
}
e.buf.Reset()
err := e.i16(int16(len(row)))
if err != nil {
return Message{}, err
}
for _, val := range row {
if uint(len(val)) >= 1<<31 {
return Message{}, ErrMessageTooBig
}
err = e.i32(int32(len(val)))
if err != nil {
return Message{}, err
}
_, err = e.buf.WriteString(val)
if err != nil {
return Message{}, err
}
}
return Message{
Type: TypeDataRow,
Data: e.buf.Bytes(),
}, nil
}
// NegotiateProtocolVersion encodes a protocol negotiation packet.
func (e *Encoder) NegotiateProtocolVersion(maxMinor int32, unrecognizedOptions ...string) (Message, error) {
if uint64(len(unrecognizedOptions)) >= 1<<31 {
return Message{}, ErrMessageTooBig
}
e.buf.Reset()
err := e.i32(maxMinor)
if err != nil {
return Message{}, err
}
err = e.i32(int32(len(unrecognizedOptions)))
if err != nil {
return Message{}, err
}
for _, opt := range unrecognizedOptions {
_, err = e.buf.WriteString(opt)
if err != nil {
return Message{}, err
}
err = e.buf.WriteByte(0)
if err != nil {
return Message{}, err
}
}
return Message{
Type: TypeNegotiateProtocolVersion,
Data: e.buf.Bytes(),
}, nil
}

31
pg/pgtest/handler.go Normal file
View file

@ -0,0 +1,31 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pgtest
import (
"context"
"github.com/pilosa/pilosa/v2/pg"
)
// HandlerFunc implements a postgres query handler with a function.
type HandlerFunc func(context.Context, pg.QueryResultWriter, pg.Query) error
// HandleQuery calls the user's query handler function.
func (h HandlerFunc) HandleQuery(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error {
return h(ctx, w, q)
}
var _ pg.QueryHandler = HandlerFunc(nil)

71
pg/pgtest/memnet.go Normal file
View file

@ -0,0 +1,71 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pgtest
import (
"errors"
"net"
"sync"
)
// errListenerClosed is an error returned when the listener is closed.
var errListenerClosed = errors.New("listener closed")
type inMemoryListener struct {
ch chan net.Conn
closed chan struct{}
once sync.Once
}
func (l *inMemoryListener) Accept() (net.Conn, error) {
select {
case <-l.closed:
return nil, errListenerClosed
default:
}
select {
case conn := <-l.ch:
return conn, nil
case <-l.closed:
return nil, errListenerClosed
}
}
func (l *inMemoryListener) Close() error {
l.once.Do(func() { close(l.closed) })
return nil
}
type memAddr struct{}
func (a memAddr) Network() string { return "memory" }
func (a memAddr) String() string { return "memory" }
func (l *inMemoryListener) Addr() net.Addr {
return memAddr{}
}
func (l *inMemoryListener) Dial() (net.Conn, error) {
serverConn, clientConn := net.Pipe()
select {
case l.ch <- serverConn:
return clientConn, nil
case <-l.closed:
serverConn.Close()
clientConn.Close()
return nil, errListenerClosed
}
}

91
pg/pgtest/server.go Normal file
View file

@ -0,0 +1,91 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pgtest
import (
"context"
"net"
"testing"
"github.com/pilosa/pilosa/v2/pg"
"github.com/pkg/errors"
"golang.org/x/sync/errgroup"
)
// ShutdownFunc is a function to use to shut down a test fixture.
// This function will send a shutdown signal and then wait for completion.
type ShutdownFunc func() error
// Finish invokes the shutdown function and fails the test if an error occurs.
func (f ShutdownFunc) Finish(tb testing.TB, name string) {
err := f()
if err != nil {
tb.Errorf("failed to shut down %s: %v", name, err)
}
}
// ServeTCP creates a TCP listener and serves postgres wire protocol on it.
func ServeTCP(addr string, server *pg.Server) (net.Addr, ShutdownFunc, error) {
listener, err := net.Listen("tcp", addr)
if err != nil {
return nil, nil, errors.Wrap(err, "listening on TCP")
}
laddr := listener.Addr()
ctx, cancel := context.WithCancel(context.Background())
var eg errgroup.Group
eg.Go(func() error { return server.Serve(ctx, listener) })
return laddr,
func() error {
cancel()
return eg.Wait()
},
nil
}
// ServeTLS sets up TLS on the server and invokes ServeTCP.
func ServeTLS(addr string, server *pg.Server) (net.Addr, ShutdownFunc, error) {
err := SetupTLS(server)
if err != nil {
return nil, nil, errors.Wrap(err, "server TLS setup failed")
}
return ServeTCP(addr, server)
}
// ConnectFunc is a function to connect to a server.
type ConnectFunc func() (net.Conn, error)
// ServeMem serves postgres on in-memory connections.
// TLS does not work here, as it relies on the OS to buffer and discard data.
func ServeMem(server *pg.Server) (ConnectFunc, ShutdownFunc, error) {
listener := &inMemoryListener{
ch: make(chan net.Conn),
closed: make(chan struct{}),
}
ctx, cancel := context.WithCancel(context.Background())
var eg errgroup.Group
eg.Go(func() error { return server.Serve(ctx, listener) })
return listener.Dial,
func() error {
cancel()
return eg.Wait()
},
nil
}

101
pg/pgtest/tls.go Normal file
View file

@ -0,0 +1,101 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pgtest
import (
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"math/big"
"time"
"github.com/pilosa/pilosa/v2/pg"
"github.com/pkg/errors"
)
// SetupTLS generates a TLS certificate and installs it into the server.
// TODO: have the client properly trust this (generate a CA to install instead of using self-signed).
func SetupTLS(server *pg.Server) error {
// Generate an ecdsa key for the cert.
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
return errors.Wrap(err, "generating TLS key")
}
// Generate a random 128-bit serial number.
serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
if err != nil {
return errors.Wrap(err, "generating serial number")
}
// Make the certificate valid starting now.
now := time.Now()
// Create a certificate template.
template := x509.Certificate{
SerialNumber: serialNumber,
Subject: pkix.Name{
Organization: []string{"Molecula"},
},
NotBefore: now,
NotAfter: now.Add(time.Hour),
KeyUsage: x509.KeyUsageDigitalSignature,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
BasicConstraintsValid: true,
}
// Generate a self-signed x509 cert from the template and the key.
certData, err := x509.CreateCertificate(rand.Reader, &template, &template, key.Public(), key)
if err != nil {
return errors.Wrap(err, "encoding certificate x509")
}
// Encode the cert to PEM so that the TLS package can load it.
certPEM := pem.EncodeToMemory(&pem.Block{
Type: "CERTIFICATE",
Bytes: certData,
})
// Encode the private key to x509.
keyData, err := x509.MarshalPKCS8PrivateKey(key)
if err != nil {
return errors.Wrap(err, "encoding key x509")
}
// Encode the key to PEM so that the TLS package can load it.
keyPEM := pem.EncodeToMemory(&pem.Block{
Type: "PRIVATE KEY",
Bytes: keyData,
})
// Load the certificate and key from their PEM encodings.
cert, err := tls.X509KeyPair(certPEM, keyPEM)
if err != nil {
return errors.Wrap(err, "loading TLS key pair")
}
// Install the certificate into the server.
if server.TLSConfig == nil {
server.TLSConfig = &tls.Config{}
}
server.TLSConfig.Certificates = append(server.TLSConfig.Certificates, cert)
return nil
}

471
pg/protocol.go Normal file
View file

@ -0,0 +1,471 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pg
import (
"bufio"
"bytes"
"context"
"crypto/tls"
"encoding/binary"
"encoding/hex"
"fmt"
"io"
"io/ioutil"
"net"
"strings"
"sync"
"time"
"github.com/pilosa/pilosa/v2/pg/message"
"github.com/pkg/errors"
)
// Protocol is a Postgres protocol version.
type Protocol uint32
const (
// ProtocolPostgres30 is version 3.0 of the Postgres wire protocol.
ProtocolPostgres30 Protocol = (3 << 16) | 0
// 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
)
// 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) {
defer func() {
cerr := conn.Close()
if cerr != nil && err == nil {
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:]
// Handle special protocols.
switch proto {
case ProtocolCancel:
// TODO: send an actual postgres error message.
return errors.New("cancellation protocol not yet supported")
case 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")
}
}
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
}
// 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
}
}
// 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")
}
var queryReady bool
for {
if !queryReady {
// Indicate that we are ready for a query.
// TODO: provide a valid transaction state.
msg, err := encoder.ReadyForQuery(message.TransactionStatusActive)
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")
}
}
// 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.TypeSimpleQuery:
// Execute a simple query.
queryReady = false
// Parse the query message (a null-terminated string).
query := SimpleQuery(strings.TrimSuffix(string(msg.Data), "\x00"))
// 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 used because its buffer may be huge.
qwriter := &queryResultWriter{
w: w,
te: s.TypeEngine,
tag: "SELECT",
}
// Dispatch the query handler.
// TODO: cancellation (requires crazy internode logic and fake process IDs)
// This is not the connection context, since we want the request to finish safely before connection shutdown.
qerr := s.QueryHandler.HandleQuery(context.Background(), 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 = encoder.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 {
// The query completed normally.
// Notify the client of completion.
msg, err = encoder.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.
default:
// The message is not supported yet.
// Send an error.
s.Logger.Printf("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) 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(ioutil.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
}

130
pg/query.go Normal file
View file

@ -0,0 +1,130 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pg
import (
"context"
"fmt"
"github.com/pilosa/pilosa/v2/pg/message"
"github.com/pkg/errors"
)
// Query is an interface to be implemented by queries.
type Query interface {
fmt.Stringer
}
// SimpleQuery is a query sent as only a string.
// It has no parameters.
type SimpleQuery string
func (q SimpleQuery) String() string {
return string(q)
}
// ColumnInfo contains metadata about a column.
type ColumnInfo struct {
Name string
Type Type
TableID int32
FieldID int16
}
// QueryResultWriter is used to write the results of a query back over the connection.
type QueryResultWriter interface {
// WriteHeader sets the column header information.
WriteHeader(...ColumnInfo) error
// WriteRowText sends a row of data in textual format.
WriteRowText(...string) error
// Tag assigns a tag to the query.
// This should be called before the query is completed.
Tag(tag string)
}
// QueryHandler handles a query.
type QueryHandler interface {
// HandleQuery executes a query and writes the results back.
HandleQuery(context.Context, QueryResultWriter, Query) error
}
// queryResultWriter implements QueryResultWrtiter over postgres wire protocol.
// The underlying message writer must be flushed by the caller once the query has finished.
type queryResultWriter struct {
w message.Writer
te TypeEngine
enc message.Encoder
width int
wroteHeaders bool
tag string
}
func (w *queryResultWriter) WriteHeader(info ...ColumnInfo) error {
if w.wroteHeaders {
return errors.New("double-write of query headers")
}
// Translate column information into a row description message.
desc := make([]message.ColumnDescription, len(info))
for i, c := range info {
t, err := w.te.TranslateType(c.Type)
if err != nil {
return errors.Wrap(err, "translating column type")
}
t.Name = c.Name
t.TableID = c.TableID
t.FieldID = c.FieldID
desc[i] = t
}
// Encode the row description.
msg, err := w.enc.RowDescription(desc...)
if err != nil {
return errors.Wrap(err, "encoding query header")
}
w.wroteHeaders = true
w.width = len(desc)
// Write the row description.
return w.w.WriteMessage(msg)
}
func (w *queryResultWriter) WriteRowText(text ...string) error {
// Check preconditions of the call.
switch {
case !w.wroteHeaders:
return errors.New("writing rows without headers")
case len(text) != w.width:
return errors.Errorf("expected %d columns but found %d", w.width, len(text))
}
// Encode the row data as text into a DataRow message.
msg, err := w.enc.TextRow(text...)
if err != nil {
return err
}
// Write the data row over the network.
return w.w.WriteMessage(msg)
}
func (w *queryResultWriter) Tag(tag string) {
w.tag = tag
}
var _ QueryResultWriter = (*queryResultWriter)(nil)

136
pg/server.go Normal file
View file

@ -0,0 +1,136 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pg
import (
"context"
"crypto/tls"
"net"
"sync"
"time"
"github.com/pilosa/pilosa/v2/logger"
)
// Server is a postgres wire protocol server.
type Server struct {
// QueryHandler will be used to serve query requests.
QueryHandler QueryHandler
// TypeEngine is the type engine to use to result columns in query requests.
TypeEngine TypeEngine
// TLSConfig is the TLS configuration to use to serve postgres TLS connections.
TLSConfig *tls.Config
// StartupTimeout is the timeout to use for connection startup.
// If a connection fails to set up a protocol before this completes, it will be terminated.
StartupTimeout time.Duration
// ReadTimeout is the timeout to apply for active reads (reads during the lifetime of a command).
// This timeout does not apply to an idling connection.
ReadTimeout time.Duration
// WriteTimeout is the timeout to apply to network writes.
WriteTimeout time.Duration
// MaxStartupSize is the maximum size of the startup packet (in bytes).
// This defaults to 2^31-1 bytes, which is the maximum size allowed by the protocol.
MaxStartupSize uint32
// ConnectionLimit is the maximum number of connections to allow at once.
ConnectionLimit uint16
// Logger is the logger to use for error conditions and state changes.
Logger logger.Logger
}
// ServeConn serves a single connection.
func (s *Server) ServeConn(ctx context.Context, conn net.Conn) error {
return s.handle(ctx, conn)
}
// Serve accepts postgres connections from a listener and processes them.
// If the context is cancelled, this will stop accepting requests and wait until all connections have terminated.
// No error will be returned if terminated by context cancellation.
// This will close the connection for the caller.
func (s *Server) Serve(ctx context.Context, l net.Listener) (err error) {
// Ignore errors triggered by a shutdown.
// Also propogate any error from terminating the listener.
var cerr error
defer func(ctx context.Context) {
if ctx.Err() == context.Canceled {
err = cerr
}
}(ctx)
// Wait for the listener to be closed and all connections to shut down.
var wg sync.WaitGroup
defer wg.Wait()
// Wrap the context to propogate a shutdown to the listeners and connection handlers.
ctx, cancel := context.WithCancel(ctx)
defer cancel()
// Start a goroutine to shut down the listener when the context is canceled.
wg.Add(1)
go func() {
defer wg.Done()
<-ctx.Done()
cerr = l.Close()
}()
// Set up a semaphore for the connection limit.
var limit chan struct{}
done := ctx.Done()
if s.ConnectionLimit != 0 {
limit = make(chan struct{}, s.ConnectionLimit)
}
for {
if limit != nil {
// Wait for connection limit.
if len(limit) == cap(limit) {
s.Logger.Printf("postgres connection limit reached")
}
select {
case limit <- struct{}{}:
case <-done:
return nil
}
}
// Accept a connection.
conn, err := l.Accept()
if err != nil {
return err
}
// Handle the connection in another goroutine.
wg.Add(1)
go func() {
defer wg.Done()
if limit != nil {
// Restore connection limit when done.
defer func() { <-limit }()
}
err := s.handle(ctx, conn)
if err != nil {
s.Logger.Printf("postgres connection terminated with error: %v", err)
}
}()
}
}

218
pg/server_test.go Normal file
View file

@ -0,0 +1,218 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pg_test
import (
"context"
"fmt"
"net"
"os/exec"
"strconv"
"testing"
"time"
"github.com/lib/pq"
"github.com/pilosa/pilosa/v2/logger"
"github.com/pilosa/pilosa/v2/pg"
"github.com/pilosa/pilosa/v2/pg/pgtest"
)
// TestStartupTimeout tests that an incoming connection that does nothing times out and gets closed.
func TestStartupTimeout(t *testing.T) {
t.Parallel()
connect, shutdown, err := pgtest.ServeMem(&pg.Server{
StartupTimeout: time.Millisecond,
Logger: logger.NewLogfLogger(t),
})
if err != nil {
t.Fatalf("starting in-memory postgres server: %v", err)
}
defer shutdown.Finish(t, "in-memory postgres server")
conn, err := connect()
if err != nil {
t.Fatalf("failed to acquire connection: %v", err)
}
defer conn.Close()
// The server isn't sending anything, so this should block until the connection dies.
conn.Read(make([]byte, 1024)) //nolint:errcheck
}
// TestStartupInvalidLength tests that sending an HTTP GET request does not cause the server to allocate 1.2 GiB of memory.
func TestStartupInvalidLength(t *testing.T) {
t.Parallel()
res := testing.Benchmark(func(b *testing.B) {
connect, shutdown, err := pgtest.ServeMem(&pg.Server{
MaxStartupSize: 1024,
Logger: logger.NewLogfLogger(t),
})
if err != nil {
t.Fatalf("starting in-memory postgres server: %v", err)
}
defer shutdown.Finish(t, "in-memory postgres server")
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
conn, err := connect()
if err != nil {
t.Fatalf("failed to acquire connection: %v", err)
}
_, err = conn.Write([]byte("GET "))
if err != nil {
t.Fatalf("failed to write invalid length: %v", err)
}
// The server isn't sending anything, so this should block until the connection dies.
conn.Read(make([]byte, 1024)) //nolint:errcheck
err = conn.Close()
if err != nil {
t.Fatalf("failed to close connection: %v", err)
}
}
})
bpo := res.AllocedBytesPerOp()
t.Logf("allocated %d bytes per op", bpo)
if bpo > 1024*1024 {
t.Errorf("allocated too much memory: %d bytes/connection", bpo)
}
}
// TestPQConnect tests connecting the Go SQL driver `pq` to this postgres server.
func TestPQConnect(t *testing.T) {
t.Parallel()
server := &pg.Server{
StartupTimeout: time.Second,
Logger: logger.NewLogfLogger(t),
}
addr, shutdown, err := pgtest.ServeTCP(":0", server)
if err != nil {
t.Fatalf("starting postgres server: %v", err)
}
defer shutdown.Finish(t, "postgres server")
tcpAddr := addr.(*net.TCPAddr)
connector, err := pq.NewConnector(fmt.Sprintf("user=molecula dbname=pilosa sslmode=disable host=%s port=%d", tcpAddr.IP, tcpAddr.Port))
if err != nil {
t.Fatalf("failed to create connector: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, err := connector.Connect(ctx)
if err != nil {
t.Fatalf("failed to connect to postgres: %v", err)
}
defer pgtest.ShutdownFunc(conn.Close).Finish(t, "postgres TLS conn")
}
// TestPQConnectSSL tests connecting the Go SQL driver `pq` to this postgres server, with SSL enabled.
func TestPQConnectSSL(t *testing.T) {
t.Parallel()
server := &pg.Server{
StartupTimeout: time.Second,
Logger: logger.NewLogfLogger(t),
}
addr, shutdown, err := pgtest.ServeTLS(":0", server)
if err != nil {
t.Fatalf("starting postgres server: %v", err)
}
defer shutdown.Finish(t, "postgres TLS server")
tcpAddr := addr.(*net.TCPAddr)
connector, err := pq.NewConnector(fmt.Sprintf("user=molecula dbname=pilosa sslmode=require host=%s port=%d", tcpAddr.IP, tcpAddr.Port))
if err != nil {
t.Fatalf("failed to create connector: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
conn, err := connector.Connect(ctx)
if err != nil {
t.Fatalf("failed to connect to postgres: %v", err)
}
defer pgtest.ShutdownFunc(conn.Close).Finish(t, "postgres TLS conn")
}
// TestPSQLQuery tests sending a query from the `psql` command line tool.
func TestPSQLQuery(t *testing.T) {
// Check if psql is present.
// Skip this test if it is not.
_, err := exec.LookPath("psql")
if err != nil {
if err, ok := err.(*exec.Error); ok {
if err.Err == exec.ErrNotFound {
t.Skip("psql is not available")
}
}
t.Fatalf("searching for psql: %v", err)
}
server := &pg.Server{
QueryHandler: pgtest.HandlerFunc(func(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error {
err := w.WriteHeader(pg.ColumnInfo{
Name: "field",
Type: pg.TypeCharoid,
})
if err != nil {
return err
}
err = w.WriteRowText("h")
if err != nil {
return err
}
err = w.WriteRowText("xyzzy")
if err != nil {
return err
}
return nil
}),
TypeEngine: pg.PrimitiveTypeEngine{},
StartupTimeout: time.Second,
Logger: logger.NewLogfLogger(t),
}
addr, shutdown, err := pgtest.ServeTCP(":0", server)
if err != nil {
t.Fatalf("starting postgres server: %v", err)
}
defer shutdown.Finish(t, "postgres server")
tcpAddr := addr.(*net.TCPAddr)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cmd := exec.CommandContext(ctx, "psql", "-h", tcpAddr.IP.String(), "-p", strconv.Itoa(tcpAddr.Port), "-c", "test query")
data, err := cmd.CombinedOutput()
if err != nil {
t.Fatalf("psql failed: %v", string(data))
}
}

54
pg/type.go Normal file
View file

@ -0,0 +1,54 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pg
import "github.com/pilosa/pilosa/v2/pg/message"
// Type represents a postgres type.
type Type struct {
// I am not entirely sure what should be in here long term.
// For now, I am just going to leave it like this.
id int32
}
// TypeCharoid is a postgres type for text.
var TypeCharoid = Type{id: 18}
// TypeData is a type containing raw postgres wire type information.
type TypeData struct {
TypeID int32
TypeLen int16
TypeModifier int32
}
// TypeEngine is a system for managing types.
// This is necessary for compound types like arrays which need ID generation.
type TypeEngine interface {
// TranslateType populates a column description with type information.
TranslateType(Type) (message.ColumnDescription, error)
}
// PrimitiveTypeEngine is a simple type engine that only works on primitive types.
type PrimitiveTypeEngine struct{}
// TranslateType translates a type to a column description.
func (pte PrimitiveTypeEngine) TranslateType(t Type) (message.ColumnDescription, error) {
return message.ColumnDescription{
TypeID: t.id, // just charoid for now; as far as I can tell most implementations do not really use this
TypeLen: -1, // vdsm had 4. . . but the spec says this should be negative
TypeModifier: -1,
Mode: 0, // send as text
}, nil
}

26
server/config.go
View file

@ -167,6 +167,25 @@ type Config struct {
MutexFraction int `toml:"mutex-fraction"`
} `toml:"profile"`
Postgres struct {
// Addr is the address to which to bind a postgres endpoint.
// If this is empty, no endpoint will be created.
Addr string `toml:"addr"`
// TLS configuration for postgres connections.
TLS TLSConfig `toml:"tls"`
StartupTimeout toml.Duration `toml:"startup-timeout"`
ReadTimeout toml.Duration `toml:"read-timeout"`
WriteTimeout toml.Duration `toml:"write-timout"`
MaxStartupSize uint32 `toml:"max-startup-size"`
// ConnectionLimit is the maximum number of postgres connections to allow simultaneously.
// Setting this to 0 disables the limit.
// This mostly exists because other DBs seem to have it.
ConnectionLimit uint16 `toml:"max-connections"`
} `toml:"postgres"`
// Txsrc determines which Tx implementation the holder/Index will use; one
// of the available transactional-storage engines. Choices are listed
// in the string constants below. Should be one of
@ -235,6 +254,13 @@ func NewConfig() *Config {
c.Profile.BlockRate = 10000000 // 1 sample per 10 ms
c.Profile.MutexFraction = 100 // 1% sampling
// Postgres config (off by default).
c.Postgres.MaxStartupSize = 8 * 1024 * 1024
c.Postgres.StartupTimeout = toml.Duration(5 * time.Second)
c.Postgres.ReadTimeout = toml.Duration(10 * time.Second)
c.Postgres.WriteTimeout = toml.Duration(10 * time.Second)
// we don't really need a connection limit
return c
}

399
server/pg.go Normal file
View file

@ -0,0 +1,399 @@
// Copyright 2020 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package server
import (
"context"
"crypto/tls"
"encoding/json"
"fmt"
"net"
"strconv"
"strings"
"time"
"github.com/pilosa/pilosa/v2"
"github.com/pilosa/pilosa/v2/logger"
"github.com/pilosa/pilosa/v2/pg"
pb "github.com/pilosa/pilosa/v2/proto"
"github.com/pkg/errors"
"golang.org/x/sync/errgroup"
)
// PostgresServer provides a postgres endpoint on pilosa.
type PostgresServer struct {
api *pilosa.API
logger logger.Logger
eg errgroup.Group
s pg.Server
stop context.CancelFunc
}
// NewPostgresServer creates a postgres server.
func NewPostgresServer(api *pilosa.API, logger logger.Logger, tls *tls.Config) *PostgresServer {
return &PostgresServer{
api: api,
logger: logger,
s: pg.Server{
QueryHandler: &queryDecodeHandler{
child: &pilosaQueryHandler{
api: api,
},
},
TypeEngine: pg.PrimitiveTypeEngine{},
StartupTimeout: 5 * time.Second,
ReadTimeout: 10 * time.Second,
WriteTimeout: 10 * time.Second,
MaxStartupSize: 8 * 1024 * 1024,
Logger: logger,
TLSConfig: tls,
},
}
}
// Start a postgres endpoint at the specified address.
func (s *PostgresServer) Start(addr string) error {
l, err := net.Listen("tcp", addr)
if err != nil {
return errors.Wrap(err, "creating listener")
}
s.logger.Printf("serving postgres wire protocol on %s", l.Addr())
ctx, cancel := context.WithCancel(context.Background())
s.stop = cancel
s.eg.Go(func() error { return s.s.Serve(ctx, l) })
return nil
}
func (s *PostgresServer) Close() error {
if s == nil {
return nil
}
if s.stop == nil {
return nil
}
s.stop()
s.logger.Printf("waiting for postgres connections to shut down")
return s.eg.Wait()
}
type pgPQLQuery struct {
index string
query string
}
func (q pgPQLQuery) String() string {
return fmt.Sprintf("[%s]%s", q.index, q.query)
}
func pgDecodePQL(str string) (q pg.Query, err error) {
defer func() {
err = errors.Wrap(err, "not a valid PQL-over-postgres query")
}()
if !strings.HasPrefix(str, "[") {
return nil, errors.New("missing index specification")
}
idx := strings.IndexRune(str, ']')
if idx == -1 {
return nil, errors.New("unclosed bracket in index specification")
}
return pgPQLQuery{
index: str[1:idx],
query: str[idx+1:],
}, nil
}
type pilosaQueryHandler struct {
api *pilosa.API
}
func pgWriteRow(w pg.QueryResultWriter, row *pilosa.Row) error {
err := w.WriteHeader(pg.ColumnInfo{
Name: "_id",
Type: pg.TypeCharoid,
})
if err != nil {
return errors.Wrap(err, "writing result header")
}
if row.Keys != nil {
for _, k := range row.Keys {
err = w.WriteRowText(k)
if err != nil {
return errors.Wrap(err, "writing key")
}
}
} else {
for _, col := range row.Columns() {
err = w.WriteRowText(strconv.FormatUint(col, 10))
if err != nil {
return errors.Wrap(err, "writing column ID")
}
}
}
return nil
}
func pgWriteRows(w pg.QueryResultWriter, rows pilosa.RowIdentifiers) error {
err := w.WriteHeader(pg.ColumnInfo{
Name: rows.Field(),
Type: pg.TypeCharoid,
})
if err != nil {
return errors.Wrap(err, "writing result header")
}
if rows.Keys != nil {
for _, k := range rows.Keys {
err = w.WriteRowText(k)
if err != nil {
return errors.Wrap(err, "writing key")
}
}
} else {
for _, row := range rows.Rows {
err = w.WriteRowText(strconv.FormatUint(row, 10))
if err != nil {
return errors.Wrap(err, "writing row ID")
}
}
}
return nil
}
func pgFormatVal(val interface{}) string {
switch val := val.(type) {
case bool:
return strconv.FormatBool(val)
case int64:
return strconv.FormatInt(val, 10)
case uint64:
return strconv.FormatUint(val, 10)
case string:
return val
default:
data, _ := json.Marshal(val)
return string(data)
}
}
func pgWriteExtractedTable(w pg.QueryResultWriter, tbl pilosa.ExtractedTable) error {
headers := make([]pg.ColumnInfo, len(tbl.Fields)+1)
headers[0] = pg.ColumnInfo{
Name: "_id",
Type: pg.TypeCharoid,
}
dataHeaders := headers[1:]
for i, f := range tbl.Fields {
dataHeaders[i] = pg.ColumnInfo{
Name: f.Name,
Type: pg.TypeCharoid,
}
}
err := w.WriteHeader(headers...)
if err != nil {
return errors.Wrap(err, "writing result header")
}
vals := make([]string, len(headers))
dataVals := vals[1:]
for _, col := range tbl.Columns {
if col.Column.Keyed {
vals[0] = col.Column.Key
} else {
vals[0] = strconv.FormatUint(col.Column.ID, 10)
}
for i, v := range col.Rows {
dataVals[i] = pgFormatVal(v)
}
err = w.WriteRowText(vals...)
if err != nil {
return errors.Wrap(err, "writing result row")
}
}
return nil
}
func pgWriteGroupCount(w pg.QueryResultWriter, counts []pilosa.GroupCount) error {
if len(counts) == 0 {
// Not enough information is available to construct the header.
// This is a significant flaw in the data type.
return nil
}
headers := make([]pg.ColumnInfo, len(counts[0].Group)+2)
for i, g := range counts[0].Group {
headers[i] = pg.ColumnInfo{
Name: g.Field,
Type: pg.TypeCharoid,
}
}
headers[len(headers)-2] = pg.ColumnInfo{
Name: "count",
Type: pg.TypeCharoid,
}
headers[len(headers)-1] = pg.ColumnInfo{
Name: "sum",
Type: pg.TypeCharoid,
}
err := w.WriteHeader(headers...)
if err != nil {
return errors.Wrap(err, "writing result header")
}
vals := make([]string, len(headers))
for _, gc := range counts {
for j, g := range gc.Group {
var v string
switch {
case g.Value != nil:
v = strconv.FormatInt(*g.Value, 10)
case g.RowKey != "":
v = g.RowKey
default:
v = strconv.FormatUint(g.RowID, 10)
}
vals[j] = v
}
vals[len(vals)-2] = strconv.FormatUint(gc.Count, 10)
vals[len(vals)-1] = strconv.FormatInt(gc.Sum, 10)
err := w.WriteRowText(vals...)
if err != nil {
return errors.Wrap(err, "writing group count result")
}
}
return nil
}
func pgWriteRowser(w pg.QueryResultWriter, result pb.ToRowser) error {
var data []string
return result.ToRows(func(row *pb.RowResponse) error {
if data == nil {
headers := make([]pg.ColumnInfo, len(row.Columns))
for i, h := range row.Headers {
headers[i] = pg.ColumnInfo{
Name: h.Name,
Type: pg.TypeCharoid,
}
}
err := w.WriteHeader(headers...)
if err != nil {
return errors.Wrap(err, "writing headers")
}
data = make([]string, len(headers))
}
for i, col := range row.Columns {
var v string
switch col := col.ColumnVal.(type) {
case *pb.ColumnResponse_BoolVal:
v = strconv.FormatBool(col.BoolVal)
case *pb.ColumnResponse_DecimalVal:
v = col.DecimalVal.String()
case *pb.ColumnResponse_Float64Val:
v = strconv.FormatFloat(col.Float64Val, 'g', -1, 64)
case *pb.ColumnResponse_Int64Val:
v = strconv.FormatInt(col.Int64Val, 10)
case *pb.ColumnResponse_Uint64Val:
v = strconv.FormatUint(col.Uint64Val, 10)
case *pb.ColumnResponse_StringVal:
v = col.StringVal
case *pb.ColumnResponse_StringArrayVal:
data, _ := json.Marshal(col.StringArrayVal.Vals)
v = string(data)
case *pb.ColumnResponse_Uint64ArrayVal:
data, _ := json.Marshal(col.Uint64ArrayVal.Vals)
v = string(data)
default:
return errors.Errorf("unable to process value of type %T", col)
}
data[i] = v
}
return w.WriteRowText(data...)
})
}
func pgWriteResult(w pg.QueryResultWriter, result interface{}) error {
switch result := result.(type) {
case *pilosa.Row:
return pgWriteRow(w, result)
case pilosa.RowIdentifiers:
return pgWriteRows(w, result)
case pilosa.ExtractedTable:
return pgWriteExtractedTable(w, result)
case []pilosa.GroupCount:
return pgWriteGroupCount(w, result)
case pb.ToRowser: // we should avoid protobuf where we can...
return pgWriteRowser(w, result)
default:
return errors.Errorf("result type %T not yet supported", result)
}
}
func (pqh *pilosaQueryHandler) HandleQuery(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error {
switch q := q.(type) {
case pgPQLQuery:
resp, err := pqh.api.Query(ctx, &pilosa.QueryRequest{
Index: q.index,
Query: q.query,
})
if err != nil {
return errors.Wrap(err, "executing query")
}
if len(resp.Results) != 1 {
return errors.Errorf("expected 1 query result but found %d", len(resp.Results))
}
return errors.Wrap(pgWriteResult(w, resp.Results[0]), "writing query result")
default:
return errors.Errorf("query type %T not yet supported (query: %s)", q, q)
}
}
type queryDecodeHandler struct {
child pg.QueryHandler
}
func (qdh *queryDecodeHandler) HandleQuery(ctx context.Context, w pg.QueryResultWriter, q pg.Query) error {
switch qv := q.(type) {
case pg.SimpleQuery:
if strings.HasPrefix(string(qv), "[") {
pqlQuery, err := pgDecodePQL(strings.TrimSuffix(string(qv), ";"))
if err != nil {
return errors.Wrap(err, "decoding query")
}
q = pqlQuery
}
}
return qdh.child.HandleQuery(ctx, w, q)
}

25
server/server.go
View file

@ -87,6 +87,7 @@ type Command struct {
listenURI *pilosa.URI
tlsConfig *tls.Config
closeTimeout time.Duration
pgserver *PostgresServer
serverOptions []pilosa.ServerOption
}
@ -171,6 +172,29 @@ func (m *Command) Start() (err error) {
}
}()
// Initialize postgres.
m.pgserver = nil
if m.Config.Postgres.Addr != "" {
var tlsConf *tls.Config
if m.Config.Postgres.TLS.CertificatePath != "" {
conf, err := GetTLSConfig(&m.Config.Postgres.TLS, m.logger.Logger())
if err != nil {
return errors.Wrap(err, "settuing up postgres TLS")
}
tlsConf = conf
}
m.pgserver = NewPostgresServer(m.API, m.logger, tlsConf)
m.pgserver.s.StartupTimeout = time.Duration(m.Config.Postgres.StartupTimeout)
m.pgserver.s.ReadTimeout = time.Duration(m.Config.Postgres.ReadTimeout)
m.pgserver.s.WriteTimeout = time.Duration(m.Config.Postgres.WriteTimeout)
m.pgserver.s.MaxStartupSize = m.Config.Postgres.MaxStartupSize
m.pgserver.s.ConnectionLimit = m.Config.Postgres.ConnectionLimit
err := m.pgserver.Start(m.Config.Postgres.Addr)
if err != nil {
return errors.Wrap(err, "starting postgres")
}
}
close(m.Started)
return nil
}
@ -515,6 +539,7 @@ func (m *Command) Close() error {
eg.Go(m.Handler.Close)
eg.Go(m.Server.Close)
eg.Go(m.API.Close)
eg.Go(m.pgserver.Close)
if m.gossipMemberSet != nil {
eg.Go(m.gossipMemberSet.Close)
}