feat(connector): look up connector by project and domain (#7378)

Signed-off-by: Kevin Su <pingsutw@gmail.com>
Signed-off-by: Kevin Su <pingsutw@apache.org>

Author: pingsutw

flyteplugins/go/tasks/plugins/webapi/connector/client.go

@@ -132,7 +132,7 @@ func updateRegistry(
 					IsConnectorApp:      isConnectorApp,
 				}
 				supportedCategoryName := supportedCategory.GetName()
-				registryKey := RegistryKey{domain: connectorDeployment.Domain, taskTypeName: supportedCategoryName, taskTypeVersion: supportedCategory.GetVersion()}
+				registryKey := RegistryKey{project: connectorDeployment.Project, domain: connectorDeployment.Domain, taskTypeName: supportedCategoryName, taskTypeVersion: supportedCategory.GetVersion()}
 				newConnectorRegistry[registryKey] = connector
 				connectorSupportedTaskCategories[supportedCategoryName] = struct{}{}
 			}
@@ -162,7 +162,7 @@ func getConnectorRegistry(ctx context.Context, cs *ClientSet) Registry {
 	// If the connector doesn't implement the metadata service, we construct the registry based on the configuration
 	for taskType, connectorDeploymentID := range cfg.ConnectorForTaskTypes {
 		if connectorDeployment, ok := cfg.ConnectorDeployments[connectorDeploymentID]; ok {
-			registryKey := RegistryKey{domain: connectorDeployment.Domain, taskTypeName: taskType, taskTypeVersion: defaultTaskTypeVersion}
+			registryKey := RegistryKey{project: connectorDeployment.Project, domain: connectorDeployment.Domain, taskTypeName: taskType, taskTypeVersion: defaultTaskTypeVersion}
 			connector := &Connector{
 				ConnectorDeployment: connectorDeployment,
 				ConnectorID:         connectorDeploymentID,
@@ -174,7 +174,7 @@ func getConnectorRegistry(ctx context.Context, cs *ClientSet) Registry {
 
 	// Ensure that the old configuration is backward compatible
 	for _, taskType := range cfg.SupportedTaskTypes {
-		registryKey := RegistryKey{domain: cfg.DefaultConnector.Domain, taskTypeName: taskType, taskTypeVersion: defaultTaskTypeVersion}
+		registryKey := RegistryKey{project: cfg.DefaultConnector.Project, domain: cfg.DefaultConnector.Domain, taskTypeName: taskType, taskTypeVersion: defaultTaskTypeVersion}
 		if _, ok := newConnectorRegistry[registryKey]; !ok {
 			connector := &Connector{
 				ConnectorDeployment: &cfg.DefaultConnector,

flyteplugins/go/tasks/plugins/webapi/connector/config.go

@@ -101,7 +101,10 @@ type Deployment struct {
 	// DefaultTimeout gives the default RPC timeout if a more specific one is not defined in Timeouts; if neither DefaultTimeout nor Timeouts is defined for an operation, RPC timeout will not be enforced
 	DefaultTimeout config.Duration `json:"defaultTimeout,omitempty" yaml:"defaultTimeout,omitempty"`
 
-	// The tasks in this domain will be handled by this agent
+	// This connector will handle the tasks in this project
+	Project string `json:"project,omitempty" yaml:"project,omitempty"`
+
+	// This connector will handle the tasks in this domain
 	Domain string `json:"domain,omitempty" yaml:"domain,omitempty"`
 }
 

flyteplugins/go/tasks/plugins/webapi/connector/plugin.go

@@ -54,6 +54,7 @@ func (p *ConnectorService) SetSupportedTaskType(taskTypes []string) {
 }
 
 type RegistryKey struct {
+	project         string
 	domain          string
 	taskTypeName    string
 	taskTypeVersion int32
@@ -98,6 +99,7 @@ type ResourceMetaWrapper struct {
 	OutputPrefix          string
 	ConnectorResourceMeta []byte
 	TaskCategory          *connectorPb.TaskCategory
+	Project               string
 	Domain                string
 	Connection            *flyteIdl.Connection
 }
@@ -152,7 +154,9 @@ func (p *Plugin) Create(ctx context.Context, taskCtx webapi.TaskExecutionContext
 	outputPrefix := taskCtx.OutputWriter().GetOutputPrefixPath().String()
 
 	taskCategory := connectorPb.TaskCategory{Name: taskTemplate.GetType(), Version: taskTemplate.GetTaskTypeVersion()}
-	connector, err := p.getFinalConnector(&taskCategory, p.cfg, taskTemplate.GetId().GetDomain())
+	project := taskTemplate.GetId().GetProject()
+	domain := taskTemplate.GetId().GetDomain()
+	connector, err := p.getFinalConnector(&taskCategory, p.cfg, project, domain)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -215,13 +219,14 @@ func (p *Plugin) Create(ctx context.Context, taskCtx webapi.TaskExecutionContext
 		ConnectorResourceMeta: res.GetResourceMeta(),
 		TaskCategory:          &taskCategory,
 		Connection:            &connection,
-		Domain:                taskTemplate.GetId().GetDomain(),
+		Project:               project,
+		Domain:                domain,
 	}, nil, nil
 }
 
 func (p *Plugin) Get(ctx context.Context, taskCtx webapi.GetContext) (latest webapi.Resource, err error) {
 	metadata := taskCtx.ResourceMeta().(ResourceMetaWrapper)
-	connector, err := p.getFinalConnector(metadata.TaskCategory, p.cfg, metadata.Domain)
+	connector, err := p.getFinalConnector(metadata.TaskCategory, p.cfg, metadata.Project, metadata.Domain)
 	if err != nil {
 		return nil, err
 	}
@@ -260,7 +265,7 @@ func (p *Plugin) Delete(ctx context.Context, taskCtx webapi.DeleteContext) error
 		return nil
 	}
 	metadata := taskCtx.ResourceMeta().(ResourceMetaWrapper)
-	connector, err := p.getFinalConnector(metadata.TaskCategory, p.cfg, metadata.Domain)
+	connector, err := p.getFinalConnector(metadata.TaskCategory, p.cfg, metadata.Project, metadata.Domain)
 	if err != nil {
 		return err
 	}
@@ -386,22 +391,32 @@ func (p *Plugin) watchConnectors(ctx context.Context, connectorService *Connecto
 	}, p.cfg.PollInterval.Duration, ctx.Done())
 }
 
-func (p *Plugin) getFinalConnector(taskCategory *connectorPb.TaskCategory, cfg *Config, domain string) (*Connector, error) {
+func (p *Plugin) getFinalConnector(taskCategory *connectorPb.TaskCategory, cfg *Config, project, domain string) (*Connector, error) {
 	p.mu.RLock()
 	defer p.mu.RUnlock()
 
-	registryKey := RegistryKey{domain: domain, taskTypeName: taskCategory.GetName(), taskTypeVersion: taskCategory.GetVersion()}
+	taskTypeName := taskCategory.GetName()
+	taskTypeVersion := taskCategory.GetVersion()
+
+	registryKey := RegistryKey{project: project, domain: domain, taskTypeName: taskTypeName, taskTypeVersion: taskTypeVersion}
+	if connector, exists := p.registry[registryKey]; exists {
+		return connector, nil
+	}
+	logger.Debugf(context.Background(), "No connector found for task type [%s] and version [%d] in project [%s] domain [%s].", taskTypeName, taskTypeVersion, project, domain)
+
+	// Fall back to a connector registered for the domain across all projects.
+	registryKey = RegistryKey{project: "", domain: domain, taskTypeName: taskTypeName, taskTypeVersion: taskTypeVersion}
 	if connector, exists := p.registry[registryKey]; exists {
 		return connector, nil
 	}
-	logger.Debugf(context.Background(), "No connector found for task type [%s] and version [%d] in domain [%s].", taskCategory.GetName(), taskCategory.GetVersion(), domain)
+	logger.Debugf(context.Background(), "No connector found for task type [%s] and version [%d] in domain [%s].", taskTypeName, taskTypeVersion, domain)
 
-	// Use the connector that supports across all domains.
-	registryKey = RegistryKey{domain: "", taskTypeName: taskCategory.GetName(), taskTypeVersion: taskCategory.GetVersion()}
+	// Fall back to a connector that supports across all projects and domains.
+	registryKey = RegistryKey{project: "", domain: "", taskTypeName: taskTypeName, taskTypeVersion: taskTypeVersion}
 	if connector, exists := p.registry[registryKey]; exists {
 		return connector, nil
 	}
-	logger.Debugf(context.Background(), "No connector found for task type [%s] and version [%d] in any domain.", taskCategory.GetName(), taskCategory.GetVersion())
+	logger.Debugf(context.Background(), "No connector found for task type [%s] and version [%d] in any project or domain.", taskTypeName, taskTypeVersion)
 
 	if len(cfg.DefaultConnector.Endpoint) != 0 {
 		return &Connector{

flyteplugins/go/tasks/plugins/webapi/connector/plugin_test.go

@@ -67,16 +67,16 @@ func TestPlugin(t *testing.T) {
 		ray := &connectorPb.TaskCategory{Name: "ray", Version: defaultTaskTypeVersion}
 		foo := &connectorPb.TaskCategory{Name: "foo", Version: defaultTaskTypeVersion}
 		bar := &connectorPb.TaskCategory{Name: "bar", Version: defaultTaskTypeVersion}
-		connector, err := plugin.getFinalConnector(spark, &cfg, "")
+		connector, err := plugin.getFinalConnector(spark, &cfg, "", "")
 		assert.NoError(t, err)
 		assert.Equal(t, connector.ConnectorDeployment.Endpoint, "localhost:80")
-		connector, err = plugin.getFinalConnector(foo, &cfg, "")
+		connector, err = plugin.getFinalConnector(foo, &cfg, "", "")
 		assert.NoError(t, err)
 		assert.Equal(t, connector.ConnectorDeployment.Endpoint, cfg.DefaultConnector.Endpoint)
-		connector, err = plugin.getFinalConnector(bar, &cfg, "")
+		connector, err = plugin.getFinalConnector(bar, &cfg, "", "")
 		assert.NoError(t, err)
 		assert.Equal(t, connector.ConnectorDeployment.Endpoint, cfg.DefaultConnector.Endpoint)
-		connector, err = plugin.getFinalConnector(ray, &cfg, "production")
+		connector, err = plugin.getFinalConnector(ray, &cfg, "", "production")
 		assert.NoError(t, err)
 		assert.Equal(t, connector.ConnectorDeployment.Endpoint, rayConnector.ConnectorDeployment.Endpoint)
 	})