From 24adfee4e133383b2ee6f0ca23c503f7ef943f94 Mon Sep 17 00:00:00 2001 From: Rafi Shamim Date: Tue, 9 Feb 2021 17:23:27 -0500 Subject: [PATCH] sql/pgwire: send placeholder BackendKeyData Some tools expect this message to be returned at connection time and will not connect without it. CockroachDB does not support pgwire cancellation, but we can still send a placeholder value here, and continue ignoring cancellation requests like we already do. Added a pgwire cancellation test to make sure nothing broke, as well as a pgwire-level test to make sure the BackendKeyData is sent to the client. Release note (sql change): When a connection is established, CockroachDB will now return a placeholder BackendKeyData message in the response. This is for compatibility with some tools, but using the BackendKeyData to cancel a query will still have no effect, just the same as before. --- pkg/sql/pgwire/conn.go | 33 +++++++++++++++++++----- pkg/sql/pgwire/conn_test.go | 44 ++++++++++++++++++++++++++++++++ pkg/sql/pgwire/pgwirebase/msg.go | 1 + pkg/testutils/pgtest/pgtest.go | 12 +++++++++ 4 files changed, 83 insertions(+), 7 deletions(-) diff --git a/pkg/sql/pgwire/conn.go b/pkg/sql/pgwire/conn.go index 6c80ce14abd9..0bc2d7349088 100644 --- a/pkg/sql/pgwire/conn.go +++ b/pkg/sql/pgwire/conn.go @@ -302,10 +302,8 @@ func (c *conn) serveImpl( dummyCh := make(chan error) close(dummyCh) procCh = dummyCh - // An initial readyForQuery message is part of the handshake. - c.msgBuilder.initMsg(pgwirebase.ServerMsgReady) - c.msgBuilder.writeByte(byte(sql.IdleTxnBlock)) - if err := c.msgBuilder.finishMsg(c.conn); err != nil { + + if err := c.sendReadyForQuery(); err != nil { reserved.Close(ctx) return } @@ -688,14 +686,35 @@ func (c *conn) sendInitialConnData( if err := c.sendParamStatus("is_superuser", superUserVal); err != nil { return sql.ConnectionHandler{}, err } + if err := c.sendReadyForQuery(); err != nil { + return sql.ConnectionHandler{}, err + } + return connHandler, nil +} - // An initial readyForQuery message is part of the handshake. +// sendReadyForQuery sends the final messages of the connection handshake. +// This includes a placeholder BackendKeyData message and a ServerMsgReady +// message indicating that there is no active transaction. +func (c *conn) sendReadyForQuery() error { + // Send the client a dummy BackendKeyData message. This is necessary for + // compatibility with tools that require this message. This information is + // normally used by clients to send a CancelRequest message: + // https://www.postgresql.org/docs/9.6/static/protocol-flow.html#AEN112861 + // CockroachDB currently ignores all CancelRequests. + c.msgBuilder.initMsg(pgwirebase.ServerMsgBackendKeyData) + c.msgBuilder.putInt32(0) + c.msgBuilder.putInt32(0) + if err := c.msgBuilder.finishMsg(c.conn); err != nil { + return err + } + + // An initial ServerMsgReady message is part of the handshake. c.msgBuilder.initMsg(pgwirebase.ServerMsgReady) c.msgBuilder.writeByte(byte(sql.IdleTxnBlock)) if err := c.msgBuilder.finishMsg(c.conn); err != nil { - return sql.ConnectionHandler{}, err + return err } - return connHandler, nil + return nil } // An error is returned iff the statement buffer has been closed. In that case, diff --git a/pkg/sql/pgwire/conn_test.go b/pkg/sql/pgwire/conn_test.go index acb14f50ea9c..18c1083c6428 100644 --- a/pkg/sql/pgwire/conn_test.go +++ b/pkg/sql/pgwire/conn_test.go @@ -1521,3 +1521,47 @@ func TestSetSessionArguments(t *testing.T) { t.Fatal(err) } } + +// TestCancelQuery uses the pgwire-level query cancellation protocol provided +// by lib/pq to make sure that canceling a query has no effect, and makes sure +// the dummy BackendKeyData does not cause problems. +func TestCancelQuery(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + cancelCtx, cancel := context.WithCancel(context.Background()) + args := base.TestServerArgs{ + Knobs: base.TestingKnobs{ + SQLExecutor: &sql.ExecutorTestingKnobs{ + BeforeExecute: func(ctx context.Context, stmt string) { + if strings.Contains(stmt, "pg_sleep") { + cancel() + } + }, + }, + }, + } + s, _, _ := serverutils.StartServer(t, args) + defer s.Stopper().Stop(cancelCtx) + + pgURL, cleanupFunc := sqlutils.PGUrl( + t, s.ServingSQLAddr(), "TestCancelQuery" /* prefix */, url.User(security.RootUser), + ) + defer cleanupFunc() + + db, err := gosql.Open("postgres", pgURL.String()) + require.NoError(t, err) + defer db.Close() + + // Cancellation has no effect on ongoing query. + if _, err := db.QueryContext(cancelCtx, "select pg_sleep(0.2)"); err != nil { + t.Fatalf("unexpected error: %s", err) + } + + // Context is already canceled, so error should come before execution. + if _, err := db.QueryContext(cancelCtx, "select 1"); err == nil { + t.Fatal("expected error") + } else if err.Error() != "context canceled" { + t.Fatalf("unexpected error: %s", err) + } +} diff --git a/pkg/sql/pgwire/pgwirebase/msg.go b/pkg/sql/pgwire/pgwirebase/msg.go index 0eb8f478c80a..fa50ba48e4f9 100644 --- a/pkg/sql/pgwire/pgwirebase/msg.go +++ b/pkg/sql/pgwire/pgwirebase/msg.go @@ -37,6 +37,7 @@ const ( ClientMsgTerminate ClientMessageType = 'X' ServerMsgAuth ServerMessageType = 'R' + ServerMsgBackendKeyData ServerMessageType = 'K' ServerMsgBindComplete ServerMessageType = '2' ServerMsgCommandComplete ServerMessageType = 'C' ServerMsgCloseComplete ServerMessageType = '3' diff --git a/pkg/testutils/pgtest/pgtest.go b/pkg/testutils/pgtest/pgtest.go index 7f3c23bed893..513ffd8dc747 100644 --- a/pkg/testutils/pgtest/pgtest.go +++ b/pkg/testutils/pgtest/pgtest.go @@ -64,10 +64,22 @@ func NewPGTest(ctx context.Context, addr, user string) (*PGTest, error) { } msgs, err := p.Until(false /* keepErrMsg */, &pgproto3.ReadyForQuery{}) foundCrdb := false + var backendKeyData *pgproto3.BackendKeyData for _, msg := range msgs { if s, ok := msg.(*pgproto3.ParameterStatus); ok && s.Name == "crdb_version" { foundCrdb = true } + if d, ok := msg.(*pgproto3.BackendKeyData); ok { + // We inspect the BackendKeyData outside of the loop since we only + // want to do the assertions if foundCrdb==true. + backendKeyData = d + } + } + if backendKeyData == nil { + return nil, errors.Errorf("did not receive BackendKeyData") + } + if foundCrdb && (backendKeyData.ProcessID != 0 || backendKeyData.SecretKey != 0) { + return nil, errors.Errorf("unexpected BackendKeyData: %+v", d) } p.isCockroachDB = foundCrdb success = err == nil