Merge branch 'main' into shobsi-bq-tool-label

This reverts commit 0f60c7702a.

internal/tools/bigquery/bigqueryanalyzecontribution/bigqueryanalyzecontribution.go

@@ -289,6 +289,7 @@ func (t Tool) Invoke(ctx context.Context, params parameters.ParamValues, accessT
 	)
 
 	createModelQuery := bqClient.Query(createModelSQL)
+	createModelQuery.Labels = map[string]string{"genai-toolbox-tool": kind}
 
 	// Get session from provider if in protected mode.
 	// Otherwise, a new session will be created by the first query.
@@ -332,6 +333,7 @@ func (t Tool) Invoke(ctx context.Context, params parameters.ParamValues, accessT
 	getInsightsSQL := fmt.Sprintf("SELECT * FROM ML.GET_INSIGHTS(MODEL %s)", modelID)
 
 	getInsightsQuery := bqClient.Query(getInsightsSQL)
+	getInsightsQuery.Labels = map[string]string{"genai-toolbox-tool": kind}
 	getInsightsQuery.ConnectionProperties = []*bigqueryapi.ConnectionProperty{{Key: "session_id", Value: sessionID}}
 
 	job, err := getInsightsQuery.Run(ctx)

internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.go

@@ -303,6 +303,7 @@ func (t Tool) Invoke(ctx context.Context, params parameters.ParamValues, accessT
 
 	query := bqClient.Query(sql)
 	query.Location = bqClient.Location
+	query.Labels = map[string]string{"genai-toolbox-tool": kind}
 
 	query.ConnectionProperties = connProps
 

internal/tools/bigquery/bigqueryforecast/bigqueryforecast.go

@@ -279,6 +279,7 @@ func (t Tool) Invoke(ctx context.Context, params parameters.ParamValues, accessT
 	// JobStatistics.QueryStatistics.StatementType
 	query := bqClient.Query(sql)
 	query.Location = bqClient.Location
+	query.Labels = map[string]string{"genai-toolbox-tool": kind}
 	session, err := t.SessionProvider(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get BigQuery session: %w", err)

internal/tools/bigquery/bigquerysql/bigquerysql.go

@@ -230,6 +230,7 @@ func (t Tool) Invoke(ctx context.Context, params parameters.ParamValues, accessT
 	query := bqClient.Query(newStatement)
 	query.Parameters = highLevelParams
 	query.Location = bqClient.Location
+	query.Labels = map[string]string{"genai-toolbox-tool": kind}
 
 	connProps := []*bigqueryapi.ConnectionProperty{}
 	if t.SessionProvider != nil {