Add API to get domain (#5443)

Signed-off-by: zychen5186 <[email protected]>
flyteorg · Jun 20, 2024 · f87a049 · f87a049
1 parent 16e7780
commit f87a049
Show file tree

Hide file tree

Showing 30 changed files with 1,844 additions and 701 deletions.
diff --git a/flyteadmin/pkg/manager/impl/project_manager.go b/flyteadmin/pkg/manager/impl/project_manager.go
@@ -42,18 +42,6 @@ func (m *ProjectManager) CreateProject(ctx context.Context, request admin.Projec
 	return &admin.ProjectRegisterResponse{}, nil
 }
 
-func (m *ProjectManager) getDomains() []*admin.Domain {
-	configDomains := m.config.ApplicationConfiguration().GetDomainsConfig()
-	var domains = make([]*admin.Domain, len(*configDomains))
-	for index, configDomain := range *configDomains {
-		domains[index] = &admin.Domain{
-			Id:   configDomain.ID,
-			Name: configDomain.Name,
-		}
-	}
-	return domains
-}
-
 func (m *ProjectManager) ListProjects(ctx context.Context, request admin.ProjectListRequest) (*admin.Projects, error) {
 	spec := util.FilterSpec{
 		RequestFilters: request.Filters,
@@ -76,7 +64,6 @@ func (m *ProjectManager) ListProjects(ctx context.Context, request admin.Project
 		return nil, errors.NewFlyteAdminErrorf(codes.InvalidArgument,
 			"invalid pagination token %s for ListProjects", request.Token)
 	}
-
 	// And finally, query the database
 	listProjectsInput := repoInterfaces.ListResourceInput{
 		Limit:         int(request.Limit),
@@ -88,7 +75,7 @@ func (m *ProjectManager) ListProjects(ctx context.Context, request admin.Project
 	if err != nil {
 		return nil, err
 	}
-	projects := transformers.FromProjectModels(projectModels, m.getDomains())
+	projects := transformers.FromProjectModels(projectModels, m.GetDomains(ctx, admin.GetDomainRequest{}).Domains)
 
 	var token string
 	if len(projects) == int(request.Limit) {
@@ -135,11 +122,25 @@ func (m *ProjectManager) GetProject(ctx context.Context, request admin.ProjectGe
 	if err != nil {
 		return nil, err
 	}
-	projectResponse := transformers.FromProjectModel(projectModel, m.getDomains())
+	projectResponse := transformers.FromProjectModel(projectModel, m.GetDomains(ctx, admin.GetDomainRequest{}).Domains)
 
 	return &projectResponse, nil
 }
 
+func (m *ProjectManager) GetDomains(ctx context.Context, request admin.GetDomainRequest) *admin.GetDomainsResponse {
+	configDomains := m.config.ApplicationConfiguration().GetDomainsConfig()
+	var domains = make([]*admin.Domain, len(*configDomains))
+	for index, configDomain := range *configDomains {
+		domains[index] = &admin.Domain{
+			Id:   configDomain.ID,
+			Name: configDomain.Name,
+		}
+	}
+	return &admin.GetDomainsResponse{
+		Domains: domains,
+	}
+}
+
 func NewProjectManager(db repoInterfaces.Repository, config runtimeInterfaces.Configuration) interfaces.ProjectInterface {
 	return &ProjectManager{
 		db:     db,

diff --git a/flyteadmin/pkg/manager/interfaces/project.go b/flyteadmin/pkg/manager/interfaces/project.go
@@ -12,4 +12,5 @@ type ProjectInterface interface {
 	ListProjects(ctx context.Context, request admin.ProjectListRequest) (*admin.Projects, error)
 	UpdateProject(ctx context.Context, request admin.Project) (*admin.ProjectUpdateResponse, error)
 	GetProject(ctx context.Context, request admin.ProjectGetRequest) (*admin.Project, error)
+	GetDomains(ctx context.Context, request admin.GetDomainRequest) *admin.GetDomainsResponse
 }
diff --git a/flyteadmin/pkg/manager/mocks/project.go b/flyteadmin/pkg/manager/mocks/project.go
@@ -10,12 +10,14 @@ type CreateProjectFunc func(ctx context.Context, request admin.ProjectRegisterRe
 type ListProjectFunc func(ctx context.Context, request admin.ProjectListRequest) (*admin.Projects, error)
 type UpdateProjectFunc func(ctx context.Context, request admin.Project) (*admin.ProjectUpdateResponse, error)
 type GetProjectFunc func(ctx context.Context, request admin.ProjectGetRequest) (*admin.Project, error)
+type GetDomainsFunc func(ctx context.Context, request admin.GetDomainRequest) *admin.GetDomainsResponse
 
 type MockProjectManager struct {
 	listProjectFunc   ListProjectFunc
 	createProjectFunc CreateProjectFunc
 	updateProjectFunc UpdateProjectFunc
 	getProjectFunc    GetProjectFunc
+	getDomainsFunc    GetDomainsFunc
 }
 
 func (m *MockProjectManager) SetCreateProject(createProjectFunc CreateProjectFunc) {
@@ -58,3 +60,10 @@ func (m *MockProjectManager) GetProject(ctx context.Context, request admin.Proje
 	}
 	return nil, nil
 }
+
+func (m *MockProjectManager) GetDomains(ctx context.Context, request admin.GetDomainRequest) *admin.GetDomainsResponse {
+	if m.getDomainsFunc != nil {
+		return m.getDomainsFunc(ctx, request)
+	}
+	return nil
+}
diff --git a/flyteadmin/pkg/rpc/adminservice/metrics.go b/flyteadmin/pkg/rpc/adminservice/metrics.go
@@ -64,6 +64,12 @@ type projectEndpointMetrics struct {
 	get      util.RequestMetrics
 }
 
+type domainEndpointMetrics struct {
+	scope promutils.Scope
+
+	get util.RequestMetrics
+}
+
 type attributeEndpointMetrics struct {
 	scope promutils.Scope
 
@@ -117,6 +123,7 @@ type AdminMetrics struct {
 	namedEntityEndpointMetrics             namedEntityEndpointMetrics
 	nodeExecutionEndpointMetrics           nodeExecutionEndpointMetrics
 	projectEndpointMetrics                 projectEndpointMetrics
+	domainEndpointMetrics                  domainEndpointMetrics
 	projectAttributesEndpointMetrics       attributeEndpointMetrics
 	projectDomainAttributesEndpointMetrics attributeEndpointMetrics
 	workflowAttributesEndpointMetrics      attributeEndpointMetrics
@@ -179,6 +186,10 @@ func InitMetrics(adminScope promutils.Scope) AdminMetrics {
 			update:   util.NewRequestMetrics(adminScope, "update_project"),
 			get:      util.NewRequestMetrics(adminScope, "get_project"),
 		},
+		domainEndpointMetrics: domainEndpointMetrics{
+			scope: adminScope,
+			get:   util.NewRequestMetrics(adminScope, "get_domain"),
+		},
 		projectAttributesEndpointMetrics: attributeEndpointMetrics{
 			scope:  adminScope,
 			update: util.NewRequestMetrics(adminScope, "update_project_attrs"),

diff --git a/flyteadmin/pkg/rpc/adminservice/project.go b/flyteadmin/pkg/rpc/adminservice/project.go
@@ -81,3 +81,17 @@ func (m *AdminService) GetProject(ctx context.Context, request *admin.ProjectGet
 	m.Metrics.projectEndpointMetrics.get.Success()
 	return response, nil
 }
+
+func (m *AdminService) GetDomains(ctx context.Context, request *admin.GetDomainRequest) (*admin.GetDomainsResponse, error) {
+	defer m.interceptPanic(ctx, request)
+	if request == nil {
+		return nil, status.Errorf(codes.InvalidArgument, "Incorrect request, nil requests not allowed")
+	}
+	var response *admin.GetDomainsResponse
+	m.Metrics.domainEndpointMetrics.get.Time(func() {
+		response = m.ProjectManager.GetDomains(ctx, *request)
+	})
+
+	m.Metrics.domainEndpointMetrics.get.Success()
+	return response, nil
+}
diff --git a/flyteadmin/tests/project_test.go b/flyteadmin/tests/project_test.go
@@ -227,3 +227,17 @@ func TestUpdateProjectLabels_BadLabels(t *testing.T) {
 	// Assert that update went through without an error.
 	assert.EqualError(t, err, "rpc error: code = InvalidArgument desc = invalid label value [#bar]: [a valid label must be an empty string or consist of alphanumeric characters, '-', '_' or '.', and must start and end with an alphanumeric character (e.g. 'MyValue',  or 'my_value',  or '12345', regex used for validation is '(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])?')]")
 }
+
+func TestGetDomains(t *testing.T) {
+	ctx := context.Background()
+	client, conn := GetTestAdminServiceClient()
+	defer conn.Close()
+
+	domains, err := client.GetDomains(ctx, &admin.GetDomainRequest{})
+	assert.Nil(t, err)
+	assert.NotEmpty(t, domains.Domains)
+	for _, domain := range domains.Domains {
+		assert.Contains(t, []string{"development", "domain", "staging", "production"}, domain.Id)
+		assert.Contains(t, []string{"development", "domain", "staging", "production"}, domain.Name)
+	}
+}
diff --git a/flytectl/pkg/ext/domain_fetcher.go b/flytectl/pkg/ext/domain_fetcher.go
@@ -0,0 +1,15 @@
+package ext
+
+import (
+	"context"
+
+	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
+)
+
+func (a *AdminFetcherExtClient) GetDomains(ctx context.Context) (*admin.GetDomainsResponse, error) {
+	domains, err := a.AdminServiceClient().GetDomains(ctx, &admin.GetDomainRequest{})
+	if err != nil {
+		return nil, err
+	}
+	return domains, nil
+}
diff --git a/flytectl/pkg/ext/domain_fetcher_test.go b/flytectl/pkg/ext/domain_fetcher_test.go
@@ -0,0 +1,35 @@
+package ext
+
+import (
+	"testing"
+
+	"github.com/flyteorg/flyte/flyteidl/clients/go/admin/mocks"
+	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestAdminFetcherExtClient_GetDomains(t *testing.T) {
+	domain1 := &admin.Domain{
+		Id:   "development",
+		Name: "development",
+	}
+	domain2 := &admin.Domain{
+		Id:   "staging",
+		Name: "staging",
+	}
+	domain3 := &admin.Domain{
+		Id:   "production",
+		Name: "production",
+	}
+	domains := &admin.GetDomainsResponse{
+		Domains: []*admin.Domain{domain1, domain2, domain3},
+	}
+
+	adminClient := new(mocks.AdminServiceClient)
+	adminFetcherExt := AdminFetcherExtClient{AdminClient: adminClient}
+
+	adminClient.OnGetDomainsMatch(mock.Anything, mock.Anything).Return(domains, nil)
+	_, err := adminFetcherExt.GetDomains(ctx)
+	assert.Nil(t, err)
+}
diff --git a/flytectl/pkg/ext/fetcher.go b/flytectl/pkg/ext/fetcher.go
@@ -74,6 +74,8 @@ type AdminFetcherExtInterface interface {
 
 	// GetProjectByID fetches a single project by its identifier. If project does not exist, an error will be returned
 	GetProjectByID(ctx context.Context, projectID string) (*admin.Project, error)
+
+	GetDomains(ctx context.Context) (*admin.GetDomainsResponse, error)
 }
 
 // AdminFetcherExtClient is used for interacting with extended features used for fetching data from admin service

diff --git a/flytectl/pkg/ext/mocks/admin_fetcher_ext_interface.go b/flytectl/pkg/ext/mocks/admin_fetcher_ext_interface.go
diff --git a/flyteidl/clients/go/admin/mocks/AdminServiceClient.go b/flyteidl/clients/go/admin/mocks/AdminServiceClient.go
diff --git a/flyteidl/clients/go/admin/mocks/AdminServiceServer.go b/flyteidl/clients/go/admin/mocks/AdminServiceServer.go