#53 Adding support for UserAgent and UserId

### What changes were proposed in this pull request? This patch adds proper support for the `user_agent` and `user_id` property of the connection string. In addition, it provides sensible default values for these parameters when connecting from a Go application. ### Why are the changes needed? Compatibility ### Does this PR introduce _any_ user-facing change? Compatibility ### How was this patch tested? Added tests Closes #105 from grundprinzip/53. Authored-by: Martin Grund <martin.grund@databricks.com> Signed-off-by: Martin Grund <martin.grund@databricks.com>

Author: grundprinzip

spark/client/channel/channel.go

@@ -24,9 +24,13 @@ import (
 "fmt"
 "net"
 "net/url"
+"os"
+"runtime"
 "strconv"
 "strings"
 
+"github.com/apache/spark-connect-go/v35/spark"
+
 "github.com/google/uuid"
 
 "google.golang.org/grpc/credentials/insecure"
@@ -56,6 +60,9 @@ type Builder interface {
 // SessionId identifies the client side session identifier. This value must be a UUID formatted
 // as a string.
 SessionId() string
+// UserAgent identifies the user agent string that is passed as part of the request. It contains
+// information about the operating system, Go version etc.
+UserAgent() string
 }
 
 // BaseBuilder is used to parse the different parameters of the connection
@@ -69,6 +76,7 @@ type BaseBuilder struct {
 user string
 headers map[string]string
 sessionId string
+userAgent string
 }
 
 func (cb *BaseBuilder) Host() string {
@@ -95,6 +103,10 @@ func (cb *BaseBuilder) SessionId() string {
 return cb.sessionId
 }
 
+func (cb *BaseBuilder) UserAgent() string {
+return cb.userAgent
+}
+
 // Build finalizes the creation of the gprc.ClientConn by creating a GRPC channel
 // with the necessary options extracted from the connection string. For
 // TLS connections, this function will load the system certificates.
@@ -178,6 +190,7 @@ func NewBuilder(connection string) (*BaseBuilder, error) {
 port: port,
 headers: map[string]string{},
 sessionId: uuid.NewString(),
+userAgent: "",
 }
 
 elements := strings.Split(u.Path, ";")
@@ -190,10 +203,33 @@ func NewBuilder(connection string) (*BaseBuilder, error) {
 cb.user = props[1]
 } else if props[0] == "session_id" {
 cb.sessionId = props[1]
+} else if props[0] == "user_agent" {
+cb.userAgent = props[1]
 } else {
 cb.headers[props[0]] = props[1]
 }
 }
 }
+
+// Set default user ID if not set.
+if cb.user == "" {
+cb.user = os.Getenv("USER")
+if cb.user == "" {
+cb.user = "na"
+}
+}
+
+// Update the user agent if it is not set or set to a custom value.
+val := os.Getenv("SPARK_CONNECT_USER_AGENT")
+if cb.userAgent == "" && val != "" {
+cb.userAgent = os.Getenv("SPARK_CONNECT_USER_AGENT")
+} else if cb.userAgent == "" {
+cb.userAgent = "_SPARK_CONNECT_GO"
+}
+
+// In addition, to the specified user agent, we need to append information about the
+// host encoded as user agent components.
+cb.userAgent = fmt.Sprintf("%s spark/%s os/%s go/%s", cb.userAgent, spark.Version(), runtime.GOOS, runtime.Version())
+
 return cb, nil
 }

spark/client/channel/channel_test.go

@@ -96,3 +96,19 @@ func TestChannelBuildConnect(t *testing.T) {
 assert.Nil(t, err, "no error for proper connection")
 assert.NotNil(t, conn)
 }
+
+func TestChannelBulder_UserAgent(t *testing.T) {
+cb, err := channel.NewBuilder("sc://localhost")
+assert.NoError(t, err)
+assert.True(t, strings.Contains(cb.UserAgent(), "_SPARK_CONNECT_GO"))
+assert.True(t, strings.Contains(cb.UserAgent(), "go/"))
+assert.True(t, strings.Contains(cb.UserAgent(), "spark/"))
+assert.True(t, strings.Contains(cb.UserAgent(), "os/"))
+
+cb, err = channel.NewBuilder("sc://localhost/;user_agent=custom")
+assert.NoError(t, err)
+assert.True(t, strings.Contains(cb.UserAgent(), "custom"))
+assert.True(t, strings.Contains(cb.UserAgent(), "go/"))
+assert.True(t, strings.Contains(cb.UserAgent(), "spark/"))
+assert.True(t, strings.Contains(cb.UserAgent(), "os/"))
+}

spark/client/client.go

@@ -56,8 +56,9 @@ func (s *sparkConnectClientImpl) newExecutePlanRequest(plan *proto.Plan) *proto.
 SessionId: s.sessionId,
 Plan: plan,
 UserContext: &proto.UserContext{
-UserId: "na",
+UserId: s.opts.UserId,
 },
+ClientType: &s.opts.UserAgent,
 // Operation ID is needed for being able to reattach.
 OperationId: &operationId,
 RequestOptions: []*proto.ExecutePlanRequest_RequestOption{
@@ -109,16 +110,22 @@ func (s *sparkConnectClientImpl) ExecutePlan(ctx context.Context, plan *proto.Pl
 return NewExecuteResponseStream(c, s.sessionId, *request.OperationId, s.opts), nil
 }
 
-func (s *sparkConnectClientImpl) AnalyzePlan(ctx context.Context, plan *proto.Plan) (*proto.AnalyzePlanResponse, error) {
-request := proto.AnalyzePlanRequest{
+// Creates a new AnalyzePlanRequest with the necessary metadata.
+func (s *sparkConnectClientImpl) newAnalyzePlanStub() proto.AnalyzePlanRequest {
+return proto.AnalyzePlanRequest{
 SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_Schema_{
-Schema: &proto.AnalyzePlanRequest_Schema{
-Plan: plan,
-},
-},
 UserContext: &proto.UserContext{
-UserId: "na",
+UserId: s.opts.UserId,
+},
+ClientType: &s.opts.UserAgent,
+}
+}
+
+func (s *sparkConnectClientImpl) AnalyzePlan(ctx context.Context, plan *proto.Plan) (*proto.AnalyzePlanResponse, error) {
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_Schema_{
+Schema: &proto.AnalyzePlanRequest_Schema{
+Plan: plan,
 },
 }
 // Append the other items to the request.
@@ -149,17 +156,11 @@ func (s *sparkConnectClientImpl) Explain(ctx context.Context, plan *proto.Plan,
 return nil, sparkerrors.WithType(fmt.Errorf("unsupported explain mode %v",
 explainMode), sparkerrors.InvalidArgumentError)
 }
-
-request := proto.AnalyzePlanRequest{
-SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_Explain_{
-Explain: &proto.AnalyzePlanRequest_Explain{
-Plan: plan,
-ExplainMode: mode,
-},
-},
-UserContext: &proto.UserContext{
-UserId: "na",
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_Explain_{
+Explain: &proto.AnalyzePlanRequest_Explain{
+Plan: plan,
+ExplainMode: mode,
 },
 }
 // Append the other items to the request.
@@ -174,17 +175,11 @@ func (s *sparkConnectClientImpl) Explain(ctx context.Context, plan *proto.Plan,
 
 func (s *sparkConnectClientImpl) Persist(ctx context.Context, plan *proto.Plan, storageLevel utils.StorageLevel) error {
 protoLevel := utils.ToProtoStorageLevel(storageLevel)
-
-request := proto.AnalyzePlanRequest{
-SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_Persist_{
-Persist: &proto.AnalyzePlanRequest_Persist{
-Relation: plan.GetRoot(),
-StorageLevel: protoLevel,
-},
-},
-UserContext: &proto.UserContext{
-UserId: "na",
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_Persist_{
+Persist: &proto.AnalyzePlanRequest_Persist{
+Relation: plan.GetRoot(),
+StorageLevel: protoLevel,
 },
 }
 // Append the other items to the request.
@@ -198,15 +193,10 @@ func (s *sparkConnectClientImpl) Persist(ctx context.Context, plan *proto.Plan,
 }
 
 func (s *sparkConnectClientImpl) Unpersist(ctx context.Context, plan *proto.Plan) error {
-request := proto.AnalyzePlanRequest{
-SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_Unpersist_{
-Unpersist: &proto.AnalyzePlanRequest_Unpersist{
-Relation: plan.GetRoot(),
-},
-},
-UserContext: &proto.UserContext{
-UserId: "na",
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_Unpersist_{
+Unpersist: &proto.AnalyzePlanRequest_Unpersist{
+Relation: plan.GetRoot(),
 },
 }
 // Append the other items to the request.
@@ -220,15 +210,10 @@ func (s *sparkConnectClientImpl) Unpersist(ctx context.Context, plan *proto.Plan
 }
 
 func (s *sparkConnectClientImpl) GetStorageLevel(ctx context.Context, plan *proto.Plan) (*utils.StorageLevel, error) {
-request := proto.AnalyzePlanRequest{
-SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_GetStorageLevel_{
-GetStorageLevel: &proto.AnalyzePlanRequest_GetStorageLevel{
-Relation: plan.GetRoot(),
-},
-},
-UserContext: &proto.UserContext{
-UserId: "na",
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_GetStorageLevel_{
+GetStorageLevel: &proto.AnalyzePlanRequest_GetStorageLevel{
+Relation: plan.GetRoot(),
 },
 }
 // Append the other items to the request.
@@ -245,14 +230,9 @@ func (s *sparkConnectClientImpl) GetStorageLevel(ctx context.Context, plan *prot
 }
 
 func (s *sparkConnectClientImpl) SparkVersion(ctx context.Context) (string, error) {
-request := proto.AnalyzePlanRequest{
-SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_SparkVersion_{
-SparkVersion: &proto.AnalyzePlanRequest_SparkVersion{},
-},
-UserContext: &proto.UserContext{
-UserId: "na",
-},
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_SparkVersion_{
+SparkVersion: &proto.AnalyzePlanRequest_SparkVersion{},
 }
 // Append the other items to the request.
 ctx = metadata.NewOutgoingContext(ctx, s.metadata)
@@ -265,15 +245,10 @@ func (s *sparkConnectClientImpl) SparkVersion(ctx context.Context) (string, erro
 }
 
 func (s *sparkConnectClientImpl) DDLParse(ctx context.Context, sql string) (*types.StructType, error) {
-request := proto.AnalyzePlanRequest{
-SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_DdlParse{
-DdlParse: &proto.AnalyzePlanRequest_DDLParse{
-DdlString: sql,
-},
-},
-UserContext: &proto.UserContext{
-UserId: "na",
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_DdlParse{
+DdlParse: &proto.AnalyzePlanRequest_DDLParse{
+DdlString: sql,
 },
 }
 // Append the other items to the request.
@@ -287,16 +262,11 @@ func (s *sparkConnectClientImpl) DDLParse(ctx context.Context, sql string) (*typ
 }
 
 func (s *sparkConnectClientImpl) SameSemantics(ctx context.Context, plan1 *proto.Plan, plan2 *proto.Plan) (bool, error) {
-request := proto.AnalyzePlanRequest{
-SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_SameSemantics_{
-SameSemantics: &proto.AnalyzePlanRequest_SameSemantics{
-TargetPlan: plan1,
-OtherPlan: plan2,
-},
-},
-UserContext: &proto.UserContext{
-UserId: "na",
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_SameSemantics_{
+SameSemantics: &proto.AnalyzePlanRequest_SameSemantics{
+TargetPlan: plan1,
+OtherPlan: plan2,
 },
 }
 // Append the other items to the request.
@@ -310,15 +280,10 @@ func (s *sparkConnectClientImpl) SameSemantics(ctx context.Context, plan1 *proto
 }
 
 func (s *sparkConnectClientImpl) SemanticHash(ctx context.Context, plan *proto.Plan) (int32, error) {
-request := proto.AnalyzePlanRequest{
-SessionId: s.sessionId,
-Analyze: &proto.AnalyzePlanRequest_SemanticHash_{
-SemanticHash: &proto.AnalyzePlanRequest_SemanticHash{
-Plan: plan,
-},
-},
-UserContext: &proto.UserContext{
-UserId: "na",
+request := s.newAnalyzePlanStub()
+request.Analyze = &proto.AnalyzePlanRequest_SemanticHash_{
+SemanticHash: &proto.AnalyzePlanRequest_SemanticHash{
+Plan: plan,
 },
 }
 // Append the other items to the request.
@@ -337,8 +302,9 @@ func (s *sparkConnectClientImpl) Config(ctx context.Context,
 request := &proto.ConfigRequest{
 Operation: operation,
 UserContext: &proto.UserContext{
-UserId: "na",
+UserId: s.opts.UserId,
 },
+ClientType: &s.opts.UserAgent,
 }
 request.SessionId = s.sessionId
 resp, err := s.client.Config(ctx, request)

spark/client/options/options.go

@@ -17,8 +17,16 @@ package options
 
 type SparkClientOptions struct {
 ReattachExecution bool
+UserAgent string
+UserId string
 }
 
 var DefaultSparkClientOptions = SparkClientOptions{
 ReattachExecution: false,
 }
+
+func NewSparkClientOptions(reattach bool) SparkClientOptions {
+return SparkClientOptions{
+ReattachExecution: reattach,
+}
+}

spark/sql/sparksession.go

@@ -93,9 +93,15 @@ func (s *SparkSessionBuilder) Build(ctx context.Context) (SparkSession, error) {
 }
 
 sessionId := uuid.NewString()
+
+// Update the options according to the configuration.
+opts := options.NewSparkClientOptions(options.DefaultSparkClientOptions.ReattachExecution)
+opts.UserAgent = s.channelBuilder.UserAgent()
+opts.UserId = s.channelBuilder.User()
+
 return &sparkSessionImpl{
 sessionId: sessionId,
-client: client.NewSparkExecutor(conn, meta, sessionId, options.DefaultSparkClientOptions),
+client: client.NewSparkExecutor(conn, meta, sessionId, opts),
 }, nil
 }
 

spark/version.go

@@ -0,0 +1,20 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package spark
+
+func Version() string {
+return "3.5.x"
+}