Presto plugin executor

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/golang/protobuf v1.3.3
 	github.com/googleapis/gnostic v0.4.1 // indirect
 	github.com/hashicorp/golang-lru v0.5.4
-	github.com/lyft/flyteidl v0.17.6
+	github.com/lyft/flyteidl v0.17.9
 	github.com/lyft/flytestdlib v0.3.2
 	github.com/magiconair/properties v1.8.1
 	github.com/mitchellh/mapstructure v1.1.2

go.sum

@@ -296,8 +296,8 @@ github.com/lyft/api v0.0.0-20191031200350-b49a72c274e0 h1:NGL46+1RYcCXb3sShp0nQq
 github.com/lyft/api v0.0.0-20191031200350-b49a72c274e0/go.mod h1:/L5qH+AD540e7Cetbui1tuJeXdmNhO8jM6VkXeDdDhQ=
 github.com/lyft/apimachinery v0.0.0-20191031200210-047e3ea32d7f h1:PGuAMDzAen0AulUfaEhNQMYmUpa41pAVo3zHI+GJsCM=
 github.com/lyft/apimachinery v0.0.0-20191031200210-047e3ea32d7f/go.mod h1:llRdnznGEAqC3DcNm6yEj472xaFVfLM7hnYofMb12tQ=
-github.com/lyft/flyteidl v0.17.6 h1:O0qpT6ya45e/92+E84uGOYa0ZsaFoE5ZfPoyJ6e1bEQ=
-github.com/lyft/flyteidl v0.17.6/go.mod h1:/zQXxuHO11u/saxTTZc8oYExIGEShXB+xCB1/F1Cu20=
+github.com/lyft/flyteidl v0.17.9 h1:JXT9PovHqS9V3YN74x9zWT0kvIEL48c2uNoujF1KMes=
+github.com/lyft/flyteidl v0.17.9/go.mod h1:/zQXxuHO11u/saxTTZc8oYExIGEShXB+xCB1/F1Cu20=
 github.com/lyft/flytestdlib v0.3.0 h1:nIkX4MlyYdcLLzaF35RI2P5BhARt+qMgHoFto8eVNzU=
 github.com/lyft/flytestdlib v0.3.0/go.mod h1:LJPPJlkFj+wwVWMrQT3K5JZgNhZi2mULsCG4ZYhinhU=
 github.com/lyft/flytestdlib v0.3.2 h1:bY6Y+Fg6Jdc7zY4GAYuR7t2hjWwynIdmRvtLcRNaGnw=

go/tasks/plugins/presto/client/presto_client.go

@@ -0,0 +1,66 @@
+package client
+
+import (
+	"context"
+	"net/http"
+	"net/url"
+
+	"github.com/lyft/flyteplugins/go/tasks/plugins/svc"
+
+	"time"
+
+	"github.com/lyft/flyteplugins/go/tasks/plugins/presto/config"
+)
+
+const (
+	httpRequestTimeoutSecs = 30
+	//AcceptHeaderKey          = "Accept"
+	//ContentTypeHeaderKey     = "Content-Type"
+	//ContentTypeJSON          = "application/json"
+	//ContentTypeTextPlain     = "text/plain"
+	//PrestoCatalogHeader      = "X-Presto-Catalog"
+	//PrestoRoutingGroupHeader = "X-Presto-Routing-Group"
+	//PrestoSchemaHeader       = "X-Presto-Schema"
+	//PrestoSourceHeader       = "X-Presto-Source"
+	//PrestoUserHeader         = "X-Presto-User"
+)
+
+type prestoClient struct {
+	client      *http.Client
+	environment *url.URL
+}
+
+type PrestoExecuteArgs struct {
+	RoutingGroup string `json:"routing_group,omitempty"`
+	Catalog      string `json:"catalog,omitempty"`
+	Schema       string `json:"schema,omitempty"`
+	Source       string `json:"source,omitempty"`
+}
+type PrestoExecuteResponse struct {
+	ID      string
+	Status  svc.CommandStatus
+	NextURI string
+}
+
+func (p *prestoClient) ExecuteCommand(
+	ctx context.Context,
+	queryStr string,
+	extraArgs interface{}) (interface{}, error) {
+
+	return PrestoExecuteResponse{}, nil
+}
+
+func (p *prestoClient) KillCommand(ctx context.Context, commandID string) error {
+	return nil
+}
+
+func (p *prestoClient) GetCommandStatus(ctx context.Context, commandID string) (svc.CommandStatus, error) {
+	return NewPrestoStatus(ctx, "UNKNOWN"), nil
+}
+
+func NewPrestoClient(cfg *config.Config) svc.ServiceClient {
+	return &prestoClient{
+		client:      &http.Client{Timeout: httpRequestTimeoutSecs * time.Second},
+		environment: cfg.Environment.ResolveReference(&cfg.Environment.URL),
+	}
+}

go/tasks/plugins/presto/client/presto_status.go

@@ -0,0 +1,41 @@
+package client
+
+import (
+	"context"
+	"strings"
+
+	"github.com/lyft/flyteplugins/go/tasks/plugins/svc"
+	"github.com/lyft/flytestdlib/logger"
+)
+
+// This type is meant only to encapsulate the response coming from Presto as a type, it is
+// not meant to be stored locally.
+const (
+	PrestoStatusUnknown   svc.CommandStatus = "UNKNOWN"
+	PrestoStatusQueued    svc.CommandStatus = "QUEUED"
+	PrestoStatusRunning   svc.CommandStatus = "RUNNING"
+	PrestoStatusFinished  svc.CommandStatus = "FINISHED"
+	PrestoStatusFailed    svc.CommandStatus = "FAILED"
+	PrestoStatusCancelled svc.CommandStatus = "CANCELLED"
+)
+
+var PrestoStatuses = map[svc.CommandStatus]struct{}{
+	PrestoStatusUnknown:   {},
+	PrestoStatusQueued:    {},
+	PrestoStatusRunning:   {},
+	PrestoStatusFinished:  {},
+	PrestoStatusFailed:    {},
+	PrestoStatusCancelled: {},
+}
+
+func NewPrestoStatus(ctx context.Context, state string) svc.CommandStatus {
+	upperCased := strings.ToUpper(state)
+	if strings.Contains(upperCased, "FAILED") {
+		return PrestoStatusFailed
+	} else if _, ok := PrestoStatuses[svc.CommandStatus(upperCased)]; ok {
+		return svc.CommandStatus(upperCased)
+	} else {
+		logger.Warnf(ctx, "Invalid Presto Status found: %v", state)
+		return PrestoStatusUnknown
+	}
+}

go/tasks/plugins/presto/config/config.go

@@ -0,0 +1,67 @@
+package config
+
+//go:generate pflags Config --default-var=defaultConfig
+
+import (
+	"context"
+	"net/url"
+
+	"github.com/lyft/flytestdlib/config"
+	"github.com/lyft/flytestdlib/logger"
+
+	pluginsConfig "github.com/lyft/flyteplugins/go/tasks/config"
+)
+
+const prestoConfigSectionKey = "presto"
+
+func URLMustParse(s string) config.URL {
+	r, err := url.Parse(s)
+	if err != nil {
+		logger.Panicf(context.TODO(), "Bad Presto URL Specified as default, error: %s", err)
+	}
+	if r == nil {
+		logger.Panicf(context.TODO(), "Nil Presto URL specified.", err)
+	}
+	return config.URL{URL: *r}
+}
+
+type RoutingGroupConfig struct {
+	Name                             string  `json:"primaryLabel" pflag:",The name of a given Presto routing group"`
+	Limit                            int     `json:"limit" pflag:",Resource quota (in the number of outstanding requests) of the routing group"`
+	ProjectScopeQuotaProportionCap   float64 `json:"projectScopeQuotaProportionCap" pflag:",A floating point number between 0 and 1, specifying the maximum proportion of quotas allowed to allocate to a project in the routing group"`
+	NamespaceScopeQuotaProportionCap float64 `json:"namespaceScopeQuotaProportionCap" pflag:",A floating point number between 0 and 1, specifying the maximum proportion of quotas allowed to allocate to a namespace in the routing group"`
+}
+
+var (
+	defaultConfig = Config{
+		Environment:         URLMustParse("https://prestoproxy-internal.lyft.net:443"),
+		DefaultRoutingGroup: "adhoc",
+		Workers:             15,
+		LruCacheSize:        2000,
+		AwsS3ShardFormatter: "s3://lyft-modelbuilder/{}/",
+		AwsS3ShardCount:     2,
+		RoutingGroupConfigs: []RoutingGroupConfig{{Name: "adhoc", Limit: 250}},
+	}
+
+	prestoConfigSection = pluginsConfig.MustRegisterSubSection(prestoConfigSectionKey, &defaultConfig)
+)
+
+// Presto plugin configs
+type Config struct {
+	Environment         config.URL           `json:"endpoint" pflag:",Endpoint for Presto to use"`
+	DefaultRoutingGroup string               `json:"defaultRoutingGroup" pflag:",Default Presto routing group"`
+	Workers             int                  `json:"workers" pflag:",Number of parallel workers to refresh the cache"`
+	LruCacheSize        int                  `json:"lruCacheSize" pflag:",Size of the AutoRefreshCache"`
+	AwsS3ShardFormatter string               `json:"awsS3ShardFormatter" pflag:", S3 bucket prefix where Presto results will be stored"`
+	AwsS3ShardCount     int                  `json:"awsS3ShardStringLength" pflag:", Number of characters for the S3 bucket shard prefix"`
+	RoutingGroupConfigs []RoutingGroupConfig `json:"clusterConfigs" pflag:"-,A list of cluster configs. Each of the configs corresponds to a service cluster"`
+}
+
+// Retrieves the current config value or default.
+func GetPrestoConfig() *Config {
+	return prestoConfigSection.GetConfig().(*Config)
+}
+
+func SetPrestoConfig(cfg *Config) error {
+	return prestoConfigSection.SetConfig(cfg)
+}