Merge pull request #341 from dolthub/max/trace-ctx

[no-release-notes] tracing infra
dolthub · Jun 26, 2024 · e40badc · e40badc
2 parents 55a46c5 + 4083c07
commit e40badc
Show file tree

Hide file tree

Showing 15 changed files with 370 additions and 333 deletions.
diff --git a/go/mysql/auth_server.go b/go/mysql/auth_server.go
@@ -18,6 +18,7 @@ package mysql
 
 import (
 	"bytes"
+	"context"
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/sha1"
@@ -221,7 +222,7 @@ func authServerDialogSwitchData() []byte {
 func AuthServerReadPacketString(c *Conn) (string, error) {
 	// Read a packet, the password is the payload, as a
 	// zero terminated string.
-	data, err := c.ReadPacket()
+	data, err := c.ReadPacket(context.Background())
 	if err != nil {
 		return "", err
 	}

diff --git a/go/mysql/client.go b/go/mysql/client.go
@@ -173,7 +173,7 @@ func (c *Conn) Ping() error {
 	if err := c.writePacket([]byte{ComPing}); err != nil {
 		return NewSQLError(CRServerGone, SSUnknownSQLState, "%v", err)
 	}
-	data, err := c.readEphemeralPacket()
+	data, err := c.readEphemeralPacket(context.Background())
 	if err != nil {
 		return NewSQLError(CRServerLost, SSUnknownSQLState, "%v", err)
 	}
@@ -215,7 +215,7 @@ func parseCharacterSet(cs string) (uint8, error) {
 // Returns a SQLError.
 func (c *Conn) clientHandshake(params *ConnParams) error {
 	// Wait for the server initial handshake packet, and parse it.
-	data, err := c.readPacket()
+	data, err := c.readPacket(context.Background())
 	if err != nil {
 		return NewSQLError(CRServerLost, "", "initial packet read failed: %v", err)
 	}
@@ -336,7 +336,7 @@ func (c *Conn) clientHandshake(params *ConnParams) error {
 		}
 
 		// Wait for response, should be OK.
-		response, err := c.readPacket()
+		response, err := c.readPacket(context.Background())
 		if err != nil {
 			return NewSQLError(CRServerLost, SSUnknownSQLState, "%v", err)
 		}
@@ -358,7 +358,7 @@ func (c *Conn) clientHandshake(params *ConnParams) error {
 // handleAuthResponse parses server's response after client sends the password for authentication
 // and handles next steps for AuthSwitchRequestPacket and AuthMoreDataPacket.
 func (c *Conn) handleAuthResponse(params *ConnParams) error {
-	response, err := c.readPacket()
+	response, err := c.readPacket(context.Background())
 	if err != nil {
 		return NewSQLError(CRServerLost, SSUnknownSQLState, "%v", err)
 	}
@@ -759,7 +759,7 @@ func (c *Conn) requestPublicKey() (rsaKey *rsa.PublicKey, err error) {
 		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "error sending public key request packet: %v", err)
 	}
 
-	response, err := c.readPacket()
+	response, err := c.readPacket(context.Background())
 	if err != nil {
 		return nil, NewSQLError(CRServerLost, SSUnknownSQLState, "%v", err)
 	}