diff --git a/client_test.go b/client_test.go index 803cd59..ec203d9 100644 --- a/client_test.go +++ b/client_test.go @@ -4,6 +4,7 @@ import ( "context" "database/sql/driver" "encoding/json" + "strings" "testing" "github.com/google/uuid" @@ -50,7 +51,7 @@ func TestCheckQueryID(t *testing.T) { assert.True(t, ok, "Expected ContextKeyQueryID to be present in the context") // Test case 2: Context already has ContextKeyQueryID - queryID := uuid.NewString() + queryID := strings.ReplaceAll(uuid.NewString(), "-", "") ctxWithQueryID := context.WithValue(ctx, ContextKeyQueryID, queryID) newCtxWithQueryID := checkQueryID(ctxWithQueryID) assert.Equal(t, queryID, newCtxWithQueryID.Value(ContextKeyQueryID), "Expected ContextKeyQueryID to remain unchanged in the context") @@ -64,7 +65,7 @@ func TestMakeHeadersQueryID(t *testing.T) { tenant: "default", sessionState: &SessionState{Role: "role1"}, } - queryId := uuid.NewString() + queryId := strings.ReplaceAll(uuid.NewString(), "-", "") ctx := context.WithValue(context.Background(), ContextKeyQueryID, queryId) headers, err := c.makeHeaders(ctx) assert.Nil(t, err) diff --git a/connection.go b/connection.go index 2d4ea76..eae3d91 100644 --- a/connection.go +++ b/connection.go @@ -6,6 +6,7 @@ import ( "fmt" "log" "os" + "strings" "sync/atomic" "github.com/google/uuid" @@ -192,7 +193,7 @@ func (dc *DatabendConn) ExecuteBatch() (err error) { // checkQueryID checks if query_id exists in context, if not, generate a new one func checkQueryID(ctx context.Context) context.Context { if _, ok := ctx.Value(ContextKeyQueryID).(string); !ok { - queryId := uuid.NewString() + queryId := strings.ReplaceAll(uuid.NewString(), "-", "") ctx = context.WithValue(ctx, ContextKeyQueryID, queryId) } return ctx