support kz1 region

commit_hash:107d15f748ba862fca6f406dd3c246db78cf8351

Author: kabachook

.mapping.json

@@ -15,9 +15,14 @@
   "client/client.go":"security/cloudquery/cq-source-yc/client/client.go",
   "client/grpc.go":"security/cloudquery/cq-source-yc/client/grpc.go",
   "client/multiplexer.go":"security/cloudquery/cq-source-yc/client/multiplexer.go",
+  "client/multiplexer_test.go":"security/cloudquery/cq-source-yc/client/multiplexer_test.go",
+  "client/region.go":"security/cloudquery/cq-source-yc/client/region.go",
+  "client/region_test.go":"security/cloudquery/cq-source-yc/client/region_test.go",
   "client/resolvers.go":"security/cloudquery/cq-source-yc/client/resolvers.go",
   "client/resolvers_test.go":"security/cloudquery/cq-source-yc/client/resolvers_test.go",
   "client/resourcetype.go":"security/cloudquery/cq-source-yc/client/resourcetype.go",
+  "client/service.go":"security/cloudquery/cq-source-yc/client/service.go",
+  "client/service_test.go":"security/cloudquery/cq-source-yc/client/service_test.go",
   "client/spec.go":"security/cloudquery/cq-source-yc/client/spec.go",
   "client/testdata/yc_compute_instances.json":"security/cloudquery/cq-source-yc/client/testdata/yc_compute_instances.json",
   "client/testdata/yc_kubernetes_clusters.json":"security/cloudquery/cq-source-yc/client/testdata/yc_kubernetes_clusters.json",
@@ -182,6 +187,7 @@
   "plugin/client.go":"security/cloudquery/cq-source-yc/plugin/client.go",
   "plugin/plugin.go":"security/cloudquery/cq-source-yc/plugin/plugin.go",
   "plugin/tables.go":"security/cloudquery/cq-source-yc/plugin/tables.go",
+  "plugin/tables_test.go":"security/cloudquery/cq-source-yc/plugin/tables_test.go",
   "resources/access/access.go":"security/cloudquery/cq-source-yc/resources/access/access.go",
   "resources/access/access_policy.go":"security/cloudquery/cq-source-yc/resources/access/access_policy.go",
   "resources/access/access_policy_organizationmanager_organizations.go":"security/cloudquery/cq-source-yc/resources/access/access_policy_organizationmanager_organizations.go",

client/client.go

@@ -10,6 +10,7 @@ import (
 	"github.com/rs/zerolog"
 	"github.com/yandex-cloud/cq-source-yc/client/yc"
 	"github.com/yandex-cloud/cq-source-yc/client/yc/datalens"
+	"github.com/yandex-cloud/go-genproto/yandex/cloud/endpoint"
 	ycsdk "github.com/yandex-cloud/go-sdk"
 	ycsdkv2 "github.com/yandex-cloud/go-sdk/v2"
 )
@@ -21,13 +22,18 @@ const (
 
 type Client struct {
 	hierarchy *yc.ResourceHierarchy
+	// services the installation offers, discovered once at startup. Nil when it
+	// was never discovered, which disables skipping altogether.
+	services serviceSet
 
 	OrganizationId          string
 	CloudId                 string
 	FolderId                string
 	MultiplexedResourceId   string
 	MultiplexedResourceType ResourceType
 
+	Region Region
+
 	Backend  state.Client
 	Logger   zerolog.Logger
 	SDK      *ycsdk.SDK
@@ -87,6 +93,8 @@ func (c *Client) WithMultiplexedResourceId(id string) *Client {
 }
 
 func New(ctx context.Context, logger zerolog.Logger, spec *Spec) (*Client, error) {
+	region, _ := RegionFromEndpoint(spec.Endpoint)
+
 	sdk, sdkv2, err := yc.Build(ctx, logger, yc.Config{
 		Endpoint:   spec.Endpoint,
 		UserAgent:  DefaultUserAgent,
@@ -96,6 +104,11 @@ func New(ctx context.Context, logger zerolog.Logger, spec *Spec) (*Client, error
 	if err != nil {
 		return nil, err
 	}
+	// Discover available services for `sdk.KnownServices`
+	_, err = sdk.ApiEndpoint().ApiEndpoint().List(ctx, &endpoint.ListApiEndpointsRequest{})
+	if err != nil {
+		return nil, fmt.Errorf("failed to discover services: %w", err)
+	}
 
 	// The middleware caches IAM tokens and refreshes them before expiry.
 	iamTokens := ycsdk.NewIAMTokenMiddleware(sdk, time.Now)
@@ -118,15 +131,22 @@ func New(ctx context.Context, logger zerolog.Logger, spec *Spec) (*Client, error
 		SDKv2:    sdkv2,
 		Datalens: dl,
 		Logger:   logger,
+		Region:   region,
 	}
 
 	hierarchy, err := yc.NewResourceHierarchy(ctx, logger, sdk, spec.OrganizationIDs, spec.CloudIDs, spec.FolderIDs)
 	if err != nil {
 		return nil, fmt.Errorf("fetch resource hierarchy: %w", err)
 	}
-
 	client.hierarchy = hierarchy
 
+	client.services = newServiceSet(sdk.KnownServices())
+
+	client.Logger.Debug().
+		Str("region", string(region)).
+		Interface("services", sdk.KnownServices()).
+		Msg("services offered by this installation")
+
 	if len(spec.OrganizationIDs) == 0 && len(spec.CloudIDs) == 0 && len(spec.FolderIDs) == 0 {
 		client.Logger.Warn().Msg("no organization_ids, cloud_ids, or folder_ids specified – assuming all resources nested in all orgs")
 	}

client/multiplexer.go

@@ -4,43 +4,77 @@ import (
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 )
 
-func OrganizationMultiplex(meta schema.ClientMeta) []schema.ClientMeta {
-	client := meta.(*Client)
-	hierarchyItems := client.hierarchy.OrganizationRows()
+// available reports whether a service can be reached in the installation this
+// client is connected to, logging the tables we skip because it cannot.
+func (c *Client) available(service Service) bool {
+	if c.serviceAvailable(service) {
+		return true
+	}
+	c.Logger.Info().
+		Str("service", string(service)).
+		Str("region", string(c.Region)).
+		Msg("skipping tables: this installation does not offer the service")
+	return false
+}
 
-	var l = make([]schema.ClientMeta, len(hierarchyItems))
-	for i, item := range hierarchyItems {
-		l[i] = client.WithOrganization(item.Organization).WithMultiplexedResourceId(item.Organization)
+// GlobalMultiplex is for tables whose RPC takes no folder, cloud or
+// organization: it keeps the single unscoped client the scheduler would use
+// anyway, but drops it in regions that do not have the service.
+func GlobalMultiplex(service Service) schema.Multiplexer {
+	return func(meta schema.ClientMeta) []schema.ClientMeta {
+		client := meta.(*Client)
+		if !client.available(service) {
+			return nil
+		}
+		return []schema.ClientMeta{client}
 	}
-	return l
 }
 
-func CloudMultiplex(meta schema.ClientMeta) []schema.ClientMeta {
-	client := meta.(*Client)
-	hierarchyItems := client.hierarchy.CloudRows()
+func OrganizationMultiplex(service Service) schema.Multiplexer {
+	return func(meta schema.ClientMeta) []schema.ClientMeta {
+		client := meta.(*Client)
+		if !client.available(service) {
+			return nil
+		}
+		hierarchyItems := client.hierarchy.OrganizationRows()
 
-	var l = make([]schema.ClientMeta, len(hierarchyItems))
-	for i, item := range hierarchyItems {
-		l[i] = client.WithOrganization(item.Organization).WithCloud(item.Cloud).WithMultiplexedResourceId(item.Cloud)
+		var l = make([]schema.ClientMeta, len(hierarchyItems))
+		for i, item := range hierarchyItems {
+			l[i] = client.WithOrganization(item.Organization).WithMultiplexedResourceId(item.Organization)
+		}
+		return l
 	}
-	return l
 }
 
-func FolderMultiplex(meta schema.ClientMeta) []schema.ClientMeta {
-	client := meta.(*Client)
-	hierarchyItems := client.hierarchy.FolderRows()
+func CloudMultiplex(service Service) schema.Multiplexer {
+	return func(meta schema.ClientMeta) []schema.ClientMeta {
+		client := meta.(*Client)
+		if !client.available(service) {
+			return nil
+		}
+		hierarchyItems := client.hierarchy.CloudRows()
 
-	var l = make([]schema.ClientMeta, len(hierarchyItems))
-	for i, item := range hierarchyItems {
-		l[i] = client.WithOrganization(item.Organization).WithCloud(item.Cloud).WithFolder(item.Folder).WithMultiplexedResourceId(item.Folder)
+		var l = make([]schema.ClientMeta, len(hierarchyItems))
+		for i, item := range hierarchyItems {
+			l[i] = client.WithOrganization(item.Organization).WithCloud(item.Cloud).WithMultiplexedResourceId(item.Cloud)
+		}
+		return l
 	}
-	return l
 }
 
-func PrependEmptyMultiplex(multiplexer schema.Multiplexer) schema.Multiplexer {
+func FolderMultiplex(service Service) schema.Multiplexer {
 	return func(meta schema.ClientMeta) []schema.ClientMeta {
 		client := meta.(*Client)
-		return append([]schema.ClientMeta{client}, multiplexer(meta)...)
+		if !client.available(service) {
+			return nil
+		}
+		hierarchyItems := client.hierarchy.FolderRows()
+
+		var l = make([]schema.ClientMeta, len(hierarchyItems))
+		for i, item := range hierarchyItems {
+			l[i] = client.WithOrganization(item.Organization).WithCloud(item.Cloud).WithFolder(item.Folder).WithMultiplexedResourceId(item.Folder)
+		}
+		return l
 	}
 }
 

client/multiplexer_test.go

@@ -0,0 +1,69 @@
+package client
+
+import (
+	"testing"
+
+	"github.com/cloudquery/plugin-sdk/v4/schema"
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/assert"
+)
+
+// A client whose installation does not offer the service must be turned away
+// before the multiplexer reaches for the resource hierarchy, which is why this
+// passes a client that has none.
+func TestMultiplexSkipsUnavailableService(t *testing.T) {
+	multiplexers := map[string]func(Service) schema.Multiplexer{
+		"organization": OrganizationMultiplex,
+		"cloud":        CloudMultiplex,
+		"folder":       FolderMultiplex,
+		"global":       GlobalMultiplex,
+	}
+
+	for name, multiplexer := range multiplexers {
+		t.Run(name, func(t *testing.T) {
+			c := &Client{
+				Region:   RegionKZ,
+				Logger:   zerolog.Nop(),
+				services: newServiceSet([]string{"compute"}),
+			}
+			assert.Empty(t, multiplexer(ServiceBaremetal)(c))
+		})
+	}
+}
+
+func TestGlobalMultiplexKeepsAvailableService(t *testing.T) {
+	c := &Client{
+		Region:   RegionKZ,
+		Logger:   zerolog.Nop(),
+		services: newServiceSet([]string{"compute"}),
+	}
+	assert.Equal(t, []schema.ClientMeta{c}, GlobalMultiplex(ServiceCompute)(c))
+}
+
+func TestServiceAvailable(t *testing.T) {
+	discovered := &Client{
+		Region:   RegionKZ,
+		services: newServiceSet([]string{"compute", "vpc"}),
+	}
+	assert.True(t, discovered.serviceAvailable(ServiceCompute))
+	assert.False(t, discovered.serviceAvailable(ServiceBaremetal))
+	assert.False(t, discovered.serviceAvailable("no-such-service"))
+
+	t.Run("undiscovered_client_filters_nothing", func(t *testing.T) {
+		c := &Client{Region: RegionKZ}
+		assert.True(t, c.serviceAvailable(ServiceBaremetal))
+		assert.True(t, c.serviceAvailable("no-such-service"))
+	})
+
+	t.Run("datalens_follows_the_region", func(t *testing.T) {
+		// DataLens is an HTTP API, so discovery never lists it and the region
+		// has to answer for it
+		ru := &Client{Region: RegionRU, services: newServiceSet([]string{"compute"})}
+		kz := &Client{Region: RegionKZ, services: newServiceSet([]string{"compute"})}
+		unknown := &Client{services: newServiceSet([]string{"compute"})}
+
+		assert.True(t, ru.serviceAvailable(ServiceDataLens))
+		assert.False(t, kz.serviceAvailable(ServiceDataLens))
+		assert.True(t, unknown.serviceAvailable(ServiceDataLens))
+	})
+}

client/region.go

@@ -0,0 +1,44 @@
+package client
+
+import (
+	"strings"
+)
+
+type Region string
+
+const (
+	RegionRU Region = "ru-central1"
+	RegionKZ Region = "kz1"
+)
+
+var RegionEndpoints = map[Region]string{
+	RegionRU: "api.cloud.yandex.net:443",
+	RegionKZ: "api.yandexcloud.kz:443",
+}
+
+var regionByHost = func() map[string]Region {
+	result := make(map[string]Region, len(RegionEndpoints))
+	for region, endpoint := range RegionEndpoints {
+		result[endpointHost(endpoint)] = region
+	}
+	return result
+}()
+
+func RegionFromEndpoint(endpoint string) (Region, bool) {
+	region, ok := regionByHost[endpointHost(endpoint)]
+	return region, ok
+}
+
+// endpointHost strips the scheme, the port and the path off an endpoint,
+// leaving the bare lowercase host.
+func endpointHost(endpoint string) string {
+	host := endpoint
+	if _, after, ok := strings.Cut(host, "://"); ok {
+		host = after
+	}
+	host, _, _ = strings.Cut(host, "/")
+	if i := strings.LastIndex(host, ":"); i >= 0 {
+		host = host[:i]
+	}
+	return strings.ToLower(host)
+}

client/region_test.go

@@ -0,0 +1,43 @@
+package client_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/yandex-cloud/cq-source-yc/client"
+)
+
+func TestRegionFromEndpoint(t *testing.T) {
+	tests := []struct {
+		endpoint string
+		want     client.Region
+		wantOk   bool
+	}{
+		{endpoint: "api.cloud.yandex.net:443", want: client.RegionRU, wantOk: true},
+		{endpoint: "api.yandexcloud.kz:443", want: client.RegionKZ, wantOk: true},
+		{endpoint: "api.cloud.yandex.net", want: client.RegionRU, wantOk: true},
+		{endpoint: "API.YandexCloud.KZ:443", want: client.RegionKZ, wantOk: true},
+		{endpoint: "https://api.yandexcloud.kz:443/", want: client.RegionKZ, wantOk: true},
+		{endpoint: "api.cloud-preprod.yandex.net:443"},
+		{endpoint: "localhost:7777"},
+		{endpoint: ""},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.endpoint, func(t *testing.T) {
+			region, ok := client.RegionFromEndpoint(tt.endpoint)
+			assert.Equal(t, tt.wantOk, ok)
+			assert.Equal(t, tt.want, region)
+		})
+	}
+}
+
+// Every region we claim to know must be reachable from its own endpoint,
+// otherwise the availability data can never be applied to it.
+func TestRegionEndpointsRoundTrip(t *testing.T) {
+	for region, endpoint := range client.RegionEndpoints {
+		got, ok := client.RegionFromEndpoint(endpoint)
+		assert.True(t, ok, "endpoint %q of region %q is not recognised", endpoint, region)
+		assert.Equal(t, region, got)
+	}
+}