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

@@ -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,

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

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'

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