github.com/matrixorigin/matrixone@v1.2.0/pkg/udf/pythonservice/client.go (about)

     1  // Copyright 2023 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package pythonservice
    16  
    17  import (
    18  	"context"
    19  	"io"
    20  	"math"
    21  	"sync"
    22  
    23  	"google.golang.org/grpc"
    24  	"google.golang.org/grpc/credentials/insecure"
    25  
    26  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    27  	"github.com/matrixorigin/matrixone/pkg/udf"
    28  )
    29  
    30  type Client struct {
    31  	cfg   ClientConfig
    32  	sc    udf.ServiceClient
    33  	mutex sync.Mutex
    34  }
    35  
    36  func NewClient(cfg ClientConfig) (*Client, error) {
    37  	err := cfg.Validate()
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  	return &Client{cfg: cfg}, nil
    42  }
    43  
    44  func (c *Client) init() error {
    45  	if c.sc == nil {
    46  		err := func() error {
    47  			c.mutex.Lock()
    48  			defer c.mutex.Unlock()
    49  			if c.sc == nil {
    50  				conn, err := grpc.Dial(c.cfg.ServerAddress,
    51  					grpc.WithTransportCredentials(insecure.NewCredentials()),
    52  					grpc.WithDefaultCallOptions(
    53  						grpc.MaxCallSendMsgSize(math.MaxInt32),
    54  						grpc.MaxCallRecvMsgSize(math.MaxInt32),
    55  					),
    56  				)
    57  				if err != nil {
    58  					return err
    59  				}
    60  				c.sc = udf.NewServiceClient(conn)
    61  			}
    62  			return nil
    63  		}()
    64  		if err != nil {
    65  			return err
    66  		}
    67  	}
    68  	return nil
    69  }
    70  
    71  func (c *Client) Run(ctx context.Context, request *udf.Request, pkgReader udf.PkgReader) (*udf.Response, error) {
    72  	if request.Udf.Language != udf.LanguagePython {
    73  		return nil, moerr.NewInvalidArg(ctx, "udf language", request.Udf.Language)
    74  	}
    75  
    76  	if request.Type != udf.RequestType_DataRequest {
    77  		return nil, moerr.NewInvalidInput(ctx, "type of the first udf request must be 'DataRequest'")
    78  	}
    79  
    80  	err := c.init()
    81  	if err != nil {
    82  		return nil, err
    83  	}
    84  
    85  	stream, err := c.sc.Run(ctx)
    86  	if err != nil {
    87  		return nil, err
    88  	}
    89  	defer stream.CloseSend()
    90  	err = stream.Send(request)
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  	response, err := stream.Recv()
    95  	if err != nil {
    96  		return nil, err
    97  	}
    98  	switch response.Type {
    99  	case udf.ResponseType_DataResponse:
   100  		return response, nil
   101  	case udf.ResponseType_PkgRequest:
   102  		var reader io.ReadCloser
   103  		reader, err = pkgReader.Get(ctx, request.Udf.Body)
   104  		if reader != nil {
   105  			defer reader.Close()
   106  		}
   107  		if err != nil {
   108  			return nil, err
   109  		}
   110  
   111  		pkgRequest := &udf.Request{
   112  			Udf:     request.Udf,
   113  			Type:    udf.RequestType_PkgResponse,
   114  			Context: request.Context,
   115  		}
   116  		for {
   117  			pkg := &udf.Package{}
   118  			buffer := make([]byte, 5*1024*1024)
   119  			var offset int
   120  			for {
   121  				var n int
   122  				n, err = reader.Read(buffer[offset:])
   123  				if err != nil {
   124  					if err == io.EOF {
   125  						pkg.Last = true
   126  					} else {
   127  						return nil, err
   128  					}
   129  				}
   130  				offset += n
   131  
   132  				if pkg.Last || len(buffer) == offset {
   133  					pkg.Data = buffer[:offset]
   134  					break
   135  				}
   136  			}
   137  
   138  			pkgRequest.Udf.ImportPkg = pkg
   139  			err = stream.Send(pkgRequest)
   140  			if err != nil {
   141  				return nil, err
   142  			}
   143  
   144  			response, err = stream.Recv()
   145  			if err != nil {
   146  				return nil, err
   147  			}
   148  
   149  			if pkg.Last {
   150  				return response, nil
   151  			}
   152  		}
   153  	default:
   154  		return nil, moerr.NewInternalError(ctx, "error udf response type")
   155  	}
   156  }
   157  
   158  func (c *Client) Language() string {
   159  	return udf.LanguagePython
   160  }