fix: context key type (#119)

* fix: context key type

* add version

connection.go

@@ -183,10 +183,9 @@ 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 {
-		return ctx
+	if _, ok := ctx.Value(ContextKeyQueryID).(string); !ok {
+		queryId := uuid.NewString()
+		ctx = context.WithValue(ctx, ContextKeyQueryID, queryId)
 	}
-	queryId := uuid.NewString()
-	ctx = context.WithValue(ctx, ContextKeyQueryID, queryId)
 	return ctx
 }

restful.go

@@ -29,7 +29,6 @@ const (
 	AuthMethodAccessToken  AuthMethod = "accessToken"
 )
 
-type ContextKey string
 type RequestType int
 
 // request type
@@ -41,10 +40,10 @@ const (
 )
 
 const (
-	ContextKeyQueryID  ContextKey = "X-DATABEND-QUERY-ID"
-	ContextUserAgentID ContextKey = "USER-AGENT"
-	EMPTY_FIELD_AS     string     = "empty_field_as"
-	PURGE              string     = "purge"
+	ContextKeyQueryID  string = "X-DATABEND-QUERY-ID"
+	ContextUserAgentID string = "USER-AGENT"
+	EMPTY_FIELD_AS     string = "empty_field_as"
+	PURGE              string = "purge"
 )
 
 type PresignedResponse struct {

restful_test.go

@@ -42,6 +42,20 @@ func TestMakeHeadersAccessToken(t *testing.T) {
 	assert.NotEmptyf(t, headers["X-Databend-Query-Id"], "Query ID is not empty")
 }
 
+func TestCheckQueryID(t *testing.T) {
+	// Test case 1: Context does not have ContextKeyQueryID
+	ctx := context.Background()
+	newCtx := checkQueryID(ctx)
+	_, ok := newCtx.Value(ContextKeyQueryID).(string)
+	assert.True(t, ok, "Expected ContextKeyQueryID to be present in the context")
+
+	// Test case 2: Context already has ContextKeyQueryID
+	queryID := 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")
+}
+
 func TestMakeHeadersQueryID(t *testing.T) {
 	c := APIClient{
 		user:         "root",

version.go

@@ -1,3 +1,3 @@
 package godatabend
 
-var version = "0.6.0"
+var version = "0.6.2"