refactor multiplex logic for quota_limits

commit_hash:c8e3bb5105249287739cc8ae63d342d41796ead6

Author: kabachook

plugin/tables.go

@@ -107,6 +107,7 @@ func PluginAutoGeneratedTables() schema.Tables {
 		organizationmanager.OsloginSettings(),
 		organizationmanager.Users(),
 		privatelink.PrivateEndpoints(),
+		quotamanager.QuotaLimits(),
 		quotamanager.QuotaServices(),
 		resourcemanager.Clouds(),
 		resourcemanager.Folders(),

resources/quotamanager/limits.go

@@ -4,14 +4,38 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"slices"
+	"sync"
 
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 	"github.com/cloudquery/plugin-sdk/v4/transformers"
 	"github.com/yandex-cloud/cq-source-yc/client"
 	"github.com/yandex-cloud/go-genproto/yandex/cloud/quotamanager/v1"
 )
 
+type servicesCacher struct {
+	responses sync.Map
+}
+
+func (sc *servicesCacher) Get(ctx context.Context, client *client.Client, resourceType client.ResourceType) ([]*quotamanager.Service, error) {
+	if val, ok := sc.responses.Load(resourceType); ok {
+		return val.([]*quotamanager.Service), nil
+	}
+
+	result, err := client.SDK.QuotaManager().QuotaLimit().ListServices(ctx, &quotamanager.ListServicesRequest{
+		ResourceType: string(resourceType),
+		// HACK: if services > 1000 this would work incorrectly
+		PageSize: 1000,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	sc.responses.Store(resourceType, result.Services)
+	return result.Services, nil
+}
+
+var cacher = servicesCacher{}
+
 type QuotaLimit struct {
 	Resource *quotamanager.Resource
 	*quotamanager.QuotaLimit
@@ -26,59 +50,51 @@ func QuotaLimits() *schema.Table {
 			transformers.WithUnwrapStructFields("Resource", "QuotaLimit"),
 			transformers.WithPrimaryKeys("Resource.Id", "QuotaId"),
 		),
+		Multiplex: client.CombineMultiplex(client.OrganizationMultiplex, client.CloudMultiplex),
 	}
 }
 
 func fetchQuotaLimits(ctx context.Context, meta schema.ClientMeta, parent *schema.Resource, res chan<- interface{}) error {
-	service, ok := parent.Item.(*Service)
-	if !ok {
-		return fmt.Errorf("failed to extract service id from parent %T (%s)", parent.Item, parent.Item)
+	c := meta.(*client.Client)
+
+	if c.MultiplexedResourceId == "" {
+		return fmt.Errorf("client must be multiplexed by resource hierarchy entity (multiplexed resource id is empty)")
+	}
+	if c.MultiplexedResourceType == "" {
+		return fmt.Errorf("client must be multiplexed by resource hierarchy entity (multiplexed resource type is empty) ")
 	}
 
-	var joinedErr error
-	// HACK: artificial child + multiplex: service (from parent) -> resource (from multiplex).
-	// When using `schema.Tables/Relates`, the child table doesn't use multiplexing (feature), so we are doing it manually.
-	// e.g. we want to resolve all `Service{Id: "vpc"}` quotas for `quotamanager.Resource{Id: "bg...", Type: "resource-manager.cloud"}`
-	for _, cmeta := range client.CombineMultiplex(client.OrganizationMultiplex, client.CloudMultiplex)(meta) {
-		c := cmeta.(*client.Client)
-
-		resourceId := c.MultiplexedResourceId
-		if resourceId == "" {
-			return fmt.Errorf("client must be multiplexed by resource hierarchy entity (multiplexed resource id is empty)")
-		}
+	// HACK: artificial (resource, service) multiplex: (Organizations + Clouds) ✕ service (from API call).
+	// Bad concurrency as CQ SDK limits apply per multiplexer entity – resource in our case,
+	// so concurrency key would be (resource), hence if services count is big, it would be slow
 
-		resourceType := c.MultiplexedResourceType
-		if resourceType == "" {
-			return fmt.Errorf("client must be multiplexed by resource hierarchy entity (multiplexed resource type is empty) ")
-		}
+	// Don't want to call rarely-changing response API per each org or cloud
+	services, err := cacher.Get(ctx, c, c.MultiplexedResourceType)
+	if err != nil {
+		return fmt.Errorf("failed to fetch quota-enabled services for %s", c.MultiplexedResourceType)
+	}
 
-		if slices.Index(quotaEnabledResourceTypes, resourceType) == -1 {
-			return fmt.Errorf("%T (%s) is not a quota-enabled service", resourceType, resourceType)
-		}
+	resource := &quotamanager.Resource{
+		Id:   c.MultiplexedResourceId,
+		Type: string(c.MultiplexedResourceType),
+	}
 
+	var joinedErr error
+	for _, service := range services {
 		it := c.SDK.QuotaManager().QuotaLimit().QuotaLimitIterator(ctx, &quotamanager.ListQuotaLimitsRequest{
-			Resource: &quotamanager.Resource{
-				Id:   resourceId,
-				Type: string(resourceType),
-			},
-			Service: service.Id,
+			Resource: resource,
+			Service:  service.Id,
 		})
 
 		for it.Next() {
-			ql := it.Value()
-
 			res <- &QuotaLimit{
-				Resource: &quotamanager.Resource{
-					Id:   resourceId,
-					Type: string(resourceType),
-				},
-				QuotaLimit: ql,
+				Resource:   resource,
+				QuotaLimit: it.Value(),
 			}
 		}
 
 		if err := it.Error(); err != nil {
 			// continue iterating because of artificial multiplex
-			c.Logger.Err(err)
 			joinedErr = errors.Join(joinedErr, err)
 		}
 	}

resources/quotamanager/services.go

@@ -21,12 +21,10 @@ func QuotaServices() *schema.Table {
 		Description: `https://yandex.cloud/ru/docs/quota-manager/api-ref/grpc/QuotaLimit/listServices#yandex.cloud.quotamanager.v1.Service`,
 		Resolver:    fetchQuotaServices,
 		Transform:   client.TransformWithStruct(&Service{}, transformers.WithUnwrapAllEmbeddedStructs(), transformers.WithPrimaryKeys("Id", "ResourceType")),
-		Relations: schema.Tables{
-			QuotaLimits(),
-		},
 	}
 }
 
+// NOTE: no API to get this list, so hardcode
 var quotaEnabledResourceTypes = []client.ResourceType{client.ResourceTypeOrganization, client.ResourceTypeCloud}
 
 func fetchQuotaServices(ctx context.Context, meta schema.ClientMeta, _ *schema.Resource, res chan<- interface{}) error {