You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2024/07/04 23:50:02 UTC
(spark-connect-go) branch master updated: [SPARK-48754] Refactor Channel builder to support other implementations
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark-connect-go.git
The following commit(s) were added to refs/heads/master by this push:
new 05ffb58 [SPARK-48754] Refactor Channel builder to support other implementations
05ffb58 is described below
commit 05ffb58fe1673af369bb2e86318f11a5a1ff5a43
Author: Martin Grund <ma...@databricks.com>
AuthorDate: Fri Jul 5 08:49:57 2024 +0900
[SPARK-48754] Refactor Channel builder to support other implementations
### What changes were proposed in this pull request?
Until now, there could only exist one implementation of a channel builder making it very hard to provide custom implementations of custom authentication and authorization handlers without directly manipulating the Spark Connect code.
This refactoring makes the channel builder an actual interface and allowing consumers to plugin their channel builder during session creation.
```golang
ctx := context.Background()
spark, err := session.NewSessionBuilder().ChannelBuilder(cb).Build(ctx)
```
This will then use the provided channel builder to build the GRPC connection.
### Why are the changes needed?
Compatibility.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Existing tests.
Closes #30 from grundprinzip/SPARK-48754-v2.
Authored-by: Martin Grund <ma...@databricks.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
Makefile | 2 +-
cmd/spark-connect-example-raw-grpc-client/main.go | 4 +-
cmd/spark-connect-example-spark-session/main.go | 9 +-
spark/client/channel/channel.go | 107 +++++++++++++++-------
spark/client/channel/channel_test.go | 31 ++++---
spark/client/channel/compat.go | 6 +-
spark/sql/datatype.go | 39 +++-----
spark/sql/session/sparksession.go | 22 +++--
spark/sql/session/sparksession_test.go | 28 +++---
9 files changed, 142 insertions(+), 106 deletions(-)
diff --git a/Makefile b/Makefile
index 716e28d..01fd6fb 100644
--- a/Makefile
+++ b/Makefile
@@ -21,7 +21,7 @@ GOFILES_NOVENDOR := $(shell find . -name vendor -prune -o -type f -name
GOFILES_BUILD := $(shell find . -type f -name '*.go' -not -name '*_test.go')
PROTOFILES := $(shell find . -name vendor -prune -o -type f -name '*.proto' -print)
-ALLGOFILES := $(shell find . -type f -name '*.go')
+ALLGOFILES := $(shell find . -type f -name '*.go' -not -name '*.pb.go')
DATE := $(shell date -u -d "@$(SOURCE_DATE_EPOCH)" '+%FT%T%z' 2>/dev/null || date -u '+%FT%T%z')
BUILDFLAGS_NOPIE :=
diff --git a/cmd/spark-connect-example-raw-grpc-client/main.go b/cmd/spark-connect-example-raw-grpc-client/main.go
index 21e43ab..f85316b 100644
--- a/cmd/spark-connect-example-raw-grpc-client/main.go
+++ b/cmd/spark-connect-example-raw-grpc-client/main.go
@@ -28,9 +28,7 @@ import (
"google.golang.org/grpc/credentials/insecure"
)
-var (
- remote = flag.String("remote", "localhost:15002", "the remote address of Spark Connect server to connect to")
-)
+var remote = flag.String("remote", "localhost:15002", "the remote address of Spark Connect server to connect to")
func main() {
ctx := context.Background()
diff --git a/cmd/spark-connect-example-spark-session/main.go b/cmd/spark-connect-example-spark-session/main.go
index f23697f..a9c17b0 100644
--- a/cmd/spark-connect-example-spark-session/main.go
+++ b/cmd/spark-connect-example-spark-session/main.go
@@ -19,16 +19,15 @@ package main
import (
"context"
"flag"
- "github.com/apache/spark-connect-go/v35/spark/sql/session"
"log"
+ "github.com/apache/spark-connect-go/v35/spark/sql/session"
+
"github.com/apache/spark-connect-go/v35/spark/sql"
)
-var (
- remote = flag.String("remote", "sc://localhost:15002",
- "the remote address of Spark Connect server to connect to")
-)
+var remote = flag.String("remote", "sc://localhost:15002",
+ "the remote address of Spark Connect server to connect to")
func main() {
flag.Parse()
diff --git a/spark/client/channel/channel.go b/spark/client/channel/channel.go
index 3e0f75a..94de111 100644
--- a/spark/client/channel/channel.go
+++ b/spark/client/channel/channel.go
@@ -27,6 +27,8 @@ import (
"strconv"
"strings"
+ "google.golang.org/grpc/credentials/insecure"
+
"github.com/apache/spark-connect-go/v35/spark/sparkerrors"
"golang.org/x/oauth2"
"google.golang.org/grpc"
@@ -37,27 +39,61 @@ import (
// Reserved header parameters that must not be injected as variables.
var reservedParams = []string{"user_id", "token", "use_ssl"}
-// Builder is used to parse the different parameters of the connection
+// Builder is the interface that is used to implement different patterns that
+// create the GRPC connection.
+//
+// This allows other consumers to plugin custom authentication and authorization
+// handlers without having to extend directly the Spark Connect code.
+type Builder interface {
+ Build(ctx context.Context) (*grpc.ClientConn, error)
+ Headers() map[string]string
+ Host() string
+ Port() int
+ Token() string
+ User() string
+}
+
+// BaseBuilder is used to parse the different parameters of the connection
// string according to the specification documented here:
//
// https://github.com/apache/spark/blob/master/connector/connect/docs/client-connection-string.md
-type Builder struct {
- Host string
- Port int
- Token string
- User string
- Headers map[string]string
+type BaseBuilder struct {
+ host string
+ port int
+ token string
+ user string
+ headers map[string]string
+}
+
+func (cb BaseBuilder) Host() string {
+ return cb.host
+}
+
+func (cb BaseBuilder) Port() int {
+ return cb.port
+}
+
+func (cb BaseBuilder) Token() string {
+ return cb.token
+}
+
+func (cb BaseBuilder) User() string {
+ return cb.user
+}
+
+func (cb BaseBuilder) Headers() map[string]string {
+ return cb.headers
}
// 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.
-func (cb *Builder) Build(ctx context.Context) (*grpc.ClientConn, error) {
+func (cb *BaseBuilder) Build(ctx context.Context) (*grpc.ClientConn, error) {
var opts []grpc.DialOption
- opts = append(opts, grpc.WithAuthority(cb.Host))
- if cb.Token == "" {
- opts = append(opts, grpc.WithInsecure())
+ opts = append(opts, grpc.WithAuthority(cb.host))
+ if cb.token == "" {
+ opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
} else {
// Note: On the Windows platform, use of x509.SystemCertPool() requires
// go version 1.18 or higher.
@@ -69,15 +105,14 @@ func (cb *Builder) Build(ctx context.Context) (*grpc.ClientConn, error) {
RootCAs: systemRoots,
})
opts = append(opts, grpc.WithTransportCredentials(cred))
-
- t := oauth2.Token{
- AccessToken: cb.Token,
+ ts := oauth2.StaticTokenSource(&oauth2.Token{
+ AccessToken: cb.token,
TokenType: "bearer",
- }
- opts = append(opts, grpc.WithPerRPCCredentials(oauth.NewOauthAccess(&t)))
+ })
+ opts = append(opts, grpc.WithPerRPCCredentials(oauth.TokenSource{TokenSource: ts}))
}
- remote := fmt.Sprintf("%v:%v", cb.Host, cb.Port)
+ remote := fmt.Sprintf("%v:%v", cb.host, cb.port)
conn, err := grpc.DialContext(ctx, remote, opts...)
if err != nil {
return nil, sparkerrors.WithType(fmt.Errorf("failed to connect to remote %s: %w", remote, err), sparkerrors.ConnectionError)
@@ -85,27 +120,29 @@ func (cb *Builder) Build(ctx context.Context) (*grpc.ClientConn, error) {
return conn, nil
}
-// NewBuilder creates a new instance of the Builder. This constructor effectively
+// NewBuilder creates a new instance of the BaseBuilder. This constructor effectively
// parses the connection string and extracts the relevant parameters directly.
-func NewBuilder(connection string) (*Builder, error) {
-
+func NewBuilder(connection string) (Builder, error) {
u, err := url.Parse(connection)
if err != nil {
return nil, err
}
+ if u.Hostname() == "" {
+ return nil, sparkerrors.WithType(errors.New("URL must contain a hostname"), sparkerrors.InvalidInputError)
+ }
+
if u.Scheme != "sc" {
return nil, sparkerrors.WithType(errors.New("URL schema must be set to `sc`"), sparkerrors.InvalidInputError)
}
- var port = 15002
- var host = u.Host
+ port := 15002
+ host := u.Host
// Check if the host part of the URL contains a port and extract.
if strings.Contains(u.Host, ":") {
- hostStr, portStr, err := net.SplitHostPort(u.Host)
- if err != nil {
- return nil, err
- }
+ // We can ignore the error here already since the url parsing
+ // raises the error about invalid port.
+ hostStr, portStr, _ := net.SplitHostPort(u.Host)
host = hostStr
if len(portStr) != 0 {
port, err = strconv.Atoi(portStr)
@@ -117,13 +154,15 @@ func NewBuilder(connection string) (*Builder, error) {
// Validate that the URL path is empty or follows the right format.
if u.Path != "" && !strings.HasPrefix(u.Path, "/;") {
- return nil, sparkerrors.WithType(fmt.Errorf("the URL path (%v) must be empty or have a proper parameter syntax", u.Path), sparkerrors.InvalidInputError)
+ return nil, sparkerrors.WithType(
+ fmt.Errorf("the URL path (%v) must be empty or have a proper parameter syntax", u.Path),
+ sparkerrors.InvalidInputError)
}
- cb := &Builder{
- Host: host,
- Port: port,
- Headers: map[string]string{},
+ cb := &BaseBuilder{
+ host: host,
+ port: port,
+ headers: map[string]string{},
}
elements := strings.Split(u.Path, ";")
@@ -131,11 +170,11 @@ func NewBuilder(connection string) (*Builder, error) {
props := strings.Split(e, "=")
if len(props) == 2 {
if props[0] == "token" {
- cb.Token = props[1]
+ cb.token = props[1]
} else if props[0] == "user_id" {
- cb.User = props[1]
+ cb.user = props[1]
} else {
- cb.Headers[props[0]] = props[1]
+ cb.headers[props[0]] = props[1]
}
}
}
diff --git a/spark/client/channel/channel_test.go b/spark/client/channel/channel_test.go
index 4073527..f60b813 100644
--- a/spark/client/channel/channel_test.go
+++ b/spark/client/channel/channel_test.go
@@ -18,10 +18,11 @@ package channel_test
import (
"context"
- "github.com/apache/spark-connect-go/v35/spark/client/channel"
"strings"
"testing"
+ "github.com/apache/spark-connect-go/v35/spark/client/channel"
+
"github.com/apache/spark-connect-go/v35/spark/sparkerrors"
"github.com/stretchr/testify/assert"
)
@@ -37,12 +38,14 @@ func TestBasicChannelBuilder(t *testing.T) {
func TestBasicChannelParsing(t *testing.T) {
_, err := channel.NewBuilder("abc://asdada:1333")
-
assert.False(t, strings.Contains(err.Error(), "scheme"), "Channel build should fail with wrong scheme")
- cb, err := channel.NewBuilder("sc://empty")
+ _, err = channel.NewBuilder("sc://:1333")
+ assert.False(t, strings.Contains(err.Error(), "scheme"), "Should not have an error for a proper URL")
+
+ cb, err := channel.NewBuilder("sc://empty")
assert.Nilf(t, err, "Valid path should not fail: %v", err)
- assert.Equalf(t, 15002, cb.Port, "Default port must be set, but got %v", cb.Port)
+ assert.Equalf(t, 15002, cb.Port(), "Default port must be set, but got %v", cb.Port)
_, err = channel.NewBuilder("sc://empty:port")
assert.NotNilf(t, err, "Port must be a valid integer %v", err)
@@ -55,19 +58,19 @@ func TestBasicChannelParsing(t *testing.T) {
assert.ErrorIs(t, err, sparkerrors.InvalidInputError)
cb, err = channel.NewBuilder(goodChannelURL)
- assert.Equal(t, "host", cb.Host)
- assert.Equal(t, 15002, cb.Port)
- assert.Len(t, cb.Headers, 1)
- assert.Equal(t, "c", cb.Headers["x-other-header"])
- assert.Equal(t, "a", cb.User)
- assert.Equal(t, "b", cb.Token)
+ assert.Equal(t, "host", cb.Host())
+ assert.Equal(t, 15002, cb.Port())
+ assert.Len(t, cb.Headers(), 1)
+ assert.Equal(t, "c", cb.Headers()["x-other-header"])
+ assert.Equal(t, "a", cb.User())
+ assert.Equal(t, "b", cb.Token())
cb, err = channel.NewBuilder("sc://localhost:443/;token=token;user_id=user_id;cluster_id=a")
assert.NoError(t, err)
- assert.Equal(t, 443, cb.Port)
- assert.Equal(t, "localhost", cb.Host)
- assert.Equal(t, "token", cb.Token)
- assert.Equal(t, "user_id", cb.User)
+ assert.Equal(t, 443, cb.Port())
+ assert.Equal(t, "localhost", cb.Host())
+ assert.Equal(t, "token", cb.Token())
+ assert.Equal(t, "user_id", cb.User())
}
func TestChannelBuildConnect(t *testing.T) {
diff --git a/spark/client/channel/compat.go b/spark/client/channel/compat.go
index 33f8a3e..3b4a735 100644
--- a/spark/client/channel/compat.go
+++ b/spark/client/channel/compat.go
@@ -1,6 +1,6 @@
package channel
-// ChannelBuilder re-exports Builder as its previous name for compatibility.
+// ChannelBuilder re-exports BaseBuilder as its previous name for compatibility.
//
-// Deprecated: use Builder instead.
-type ChannelBuilder = Builder
+// Deprecated: use BaseBuilder instead.
+type ChannelBuilder = BaseBuilder
diff --git a/spark/sql/datatype.go b/spark/sql/datatype.go
index ecbadf3..398aa2c 100644
--- a/spark/sql/datatype.go
+++ b/spark/sql/datatype.go
@@ -25,92 +25,79 @@ type DataType interface {
TypeName() string
}
-type BooleanType struct {
-}
+type BooleanType struct{}
func (t BooleanType) TypeName() string {
return getDataTypeName(t)
}
-type ByteType struct {
-}
+type ByteType struct{}
func (t ByteType) TypeName() string {
return getDataTypeName(t)
}
-type ShortType struct {
-}
+type ShortType struct{}
func (t ShortType) TypeName() string {
return getDataTypeName(t)
}
-type IntegerType struct {
-}
+type IntegerType struct{}
func (t IntegerType) TypeName() string {
return getDataTypeName(t)
}
-type LongType struct {
-}
+type LongType struct{}
func (t LongType) TypeName() string {
return getDataTypeName(t)
}
-type FloatType struct {
-}
+type FloatType struct{}
func (t FloatType) TypeName() string {
return getDataTypeName(t)
}
-type DoubleType struct {
-}
+type DoubleType struct{}
func (t DoubleType) TypeName() string {
return getDataTypeName(t)
}
-type DecimalType struct {
-}
+type DecimalType struct{}
func (t DecimalType) TypeName() string {
return getDataTypeName(t)
}
-type StringType struct {
-}
+type StringType struct{}
func (t StringType) TypeName() string {
return getDataTypeName(t)
}
-type BinaryType struct {
-}
+type BinaryType struct{}
func (t BinaryType) TypeName() string {
return getDataTypeName(t)
}
-type TimestampType struct {
-}
+type TimestampType struct{}
func (t TimestampType) TypeName() string {
return getDataTypeName(t)
}
-type TimestampNtzType struct {
-}
+type TimestampNtzType struct{}
func (t TimestampNtzType) TypeName() string {
return getDataTypeName(t)
}
-type DateType struct {
-}
+type DateType struct{}
func (t DateType) TypeName() string {
return getDataTypeName(t)
diff --git a/spark/sql/session/sparksession.go b/spark/sql/session/sparksession.go
index 43f3f61..8a45fb0 100644
--- a/spark/sql/session/sparksession.go
+++ b/spark/sql/session/sparksession.go
@@ -19,6 +19,7 @@ package session
import (
"context"
"fmt"
+
"github.com/apache/spark-connect-go/v35/spark/client/channel"
proto "github.com/apache/spark-connect-go/v35/internal/generated"
@@ -41,6 +42,7 @@ func NewSessionBuilder() *SparkSessionBuilder {
type SparkSessionBuilder struct {
connectionString string
+ channelBuilder channel.Builder
}
// Remote sets the connection string for remote connection
@@ -49,21 +51,27 @@ func (s *SparkSessionBuilder) Remote(connectionString string) *SparkSessionBuild
return s
}
-func (s *SparkSessionBuilder) Build(ctx context.Context) (SparkSession, error) {
+func (s *SparkSessionBuilder) ChannelBuilder(cb channel.Builder) *SparkSessionBuilder {
+ s.channelBuilder = cb
+ return s
+}
- cb, err := channel.NewBuilder(s.connectionString)
- if err != nil {
- return nil, sparkerrors.WithType(fmt.Errorf("failed to connect to remote %s: %w", s.connectionString, err), sparkerrors.ConnectionError)
+func (s *SparkSessionBuilder) Build(ctx context.Context) (SparkSession, error) {
+ if s.channelBuilder == nil {
+ cb, err := channel.NewBuilder(s.connectionString)
+ if err != nil {
+ return nil, sparkerrors.WithType(fmt.Errorf("failed to connect to remote %s: %w", s.connectionString, err), sparkerrors.ConnectionError)
+ }
+ s.channelBuilder = cb
}
-
- conn, err := cb.Build(ctx)
+ conn, err := s.channelBuilder.Build(ctx)
if err != nil {
return nil, sparkerrors.WithType(fmt.Errorf("failed to connect to remote %s: %w", s.connectionString, err), sparkerrors.ConnectionError)
}
// Add metadata to the request.
meta := metadata.MD{}
- for k, v := range cb.Headers {
+ for k, v := range s.channelBuilder.Headers() {
meta[k] = append(meta[k], v)
}
diff --git a/spark/sql/session/sparksession_test.go b/spark/sql/session/sparksession_test.go
index 9faced2..002c030 100644
--- a/spark/sql/session/sparksession_test.go
+++ b/spark/sql/session/sparksession_test.go
@@ -151,7 +151,7 @@ func TestSQLCallsExecutePlanWithSQLOnClient(t *testing.T) {
func TestNewSessionBuilderCreatesASession(t *testing.T) {
ctx := context.Background()
- spark, err := NewSessionBuilder().Remote("sc:connection").Build(ctx)
+ spark, err := NewSessionBuilder().Remote("sc://connection").Build(ctx)
assert.NoError(t, err)
assert.NotNil(t, spark)
}
@@ -195,22 +195,24 @@ func TestWriteResultStreamsArrowResultToCollector(t *testing.T) {
session := &sparkSessionImpl{
client: &connectServiceClient{
- executePlanClient: &sql.ExecutePlanClient{&mocks.ProtoClient{
- RecvResponses: []*proto.ExecutePlanResponse{
- {
- ResponseType: &proto.ExecutePlanResponse_SqlCommandResult_{
- SqlCommandResult: &proto.ExecutePlanResponse_SqlCommandResult{},
+ executePlanClient: &sql.ExecutePlanClient{
+ &mocks.ProtoClient{
+ RecvResponses: []*proto.ExecutePlanResponse{
+ {
+ ResponseType: &proto.ExecutePlanResponse_SqlCommandResult_{
+ SqlCommandResult: &proto.ExecutePlanResponse_SqlCommandResult{},
+ },
},
- },
- {
- ResponseType: &proto.ExecutePlanResponse_ArrowBatch_{
- ArrowBatch: &proto.ExecutePlanResponse_ArrowBatch{
- RowCount: 1,
- Data: buf.Bytes(),
+ {
+ ResponseType: &proto.ExecutePlanResponse_ArrowBatch_{
+ ArrowBatch: &proto.ExecutePlanResponse_ArrowBatch{
+ RowCount: 1,
+ Data: buf.Bytes(),
+ },
},
},
},
- }},
+ },
},
t: t,
},
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org