Merge pull request #1336 from snyk/feat/telemetry_iac_source_count

eliecharra · web-flow · commit 506439d7bf50 · 2022-02-09T15:05:01.000+01:00
Add iac source count to telemetry
diff --git a/pkg/driftctl.go b/pkg/driftctl.go
@@ -38,7 +38,7 @@ type ScanOptions struct {
 
 type DriftCTL struct {
 	remoteSupplier           resource.Supplier
-	iacSupplier              resource.Supplier
+	iacSupplier              resource.IaCSupplier
 	alerter                  alerter.AlerterInterface
 	analyzer                 *analyser.Analyzer
 	resourceFactory          resource.ResourceFactory
@@ -50,7 +50,7 @@ type DriftCTL struct {
 }
 
 func NewDriftCTL(remoteSupplier resource.Supplier,
-	iacSupplier resource.Supplier,
+	iacSupplier resource.IaCSupplier,
 	alerter *alerter.Alerter,
 	analyzer *analyser.Analyzer,
 	resFactory resource.ResourceFactory,
@@ -189,6 +189,7 @@ func (d DriftCTL) scan() (remoteResources []*resource.Resource, resourcesFromSta
 	if err != nil {
 		return nil, nil, err
 	}
+	d.store.Bucket(memstore.TelemetryBucket).Set("iac_source_count", d.iacSupplier.SourceCount())
 
 	logrus.Info("Start scanning cloud provider")
 	d.scanProgress.Start()
diff --git a/pkg/driftctl_test.go b/pkg/driftctl_test.go
@@ -62,8 +62,9 @@ func runTest(t *testing.T, cases TestCases) {
 				res.Sch = schema
 			}
 
-			stateSupplier := &resource.MockSupplier{}
+			stateSupplier := &resource.MockIaCSupplier{}
 			stateSupplier.On("Resources").Return(c.stateResources, nil)
+			stateSupplier.On("SourceCount").Return(uint(2))
 
 			if c.remoteResources == nil {
 				c.remoteResources = []*resource.Resource{}
@@ -138,6 +139,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 0, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 0, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 		},
 		{
@@ -155,6 +157,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 			options: func(t *testing.T) *pkg.ScanOptions {
 				return &pkg.ScanOptions{Deep: true}
@@ -173,6 +176,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 0, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 		},
 		{
@@ -226,6 +230,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 		},
 		{
@@ -264,6 +269,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 		},
 		{
@@ -304,6 +310,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 		},
 		{
@@ -344,6 +351,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 		},
 		{
@@ -434,6 +442,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 4, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 2, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 			options: func(t *testing.T) *pkg.ScanOptions {
 				return &pkg.ScanOptions{
@@ -530,6 +539,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 6, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 2, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 			options: func(t *testing.T) *pkg.ScanOptions {
 				return &pkg.ScanOptions{
@@ -628,6 +638,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) {
 				assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources"))
 				assert.Equal(t, 0, store.Bucket(memstore.TelemetryBucket).Get("total_managed"))
 				assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration"))
+				assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count"))
 			},
 			options: func(t *testing.T) *pkg.ScanOptions {
 				filterStr := "Id=='role-test-1'"
diff --git a/pkg/iac/supplier/IacChainSupplier.go b/pkg/iac/supplier/IacChainSupplier.go
@@ -10,7 +10,7 @@ import (
 )
 
 type IacChainSupplier struct {
-	suppliers []resource.Supplier
+	suppliers []resource.IaCSupplier
 	runner    *parallel.ParallelRunner
 }
 
@@ -20,7 +20,15 @@ func NewIacChainSupplier() *IacChainSupplier {
 	}
 }
 
-func (r *IacChainSupplier) AddSupplier(supplier resource.Supplier) {
+func (r *IacChainSupplier) SourceCount() uint {
+	count := uint(0)
+	for _, supplier := range r.suppliers {
+		count += supplier.SourceCount()
+	}
+	return count
+}
+
+func (r *IacChainSupplier) AddSupplier(supplier resource.IaCSupplier) {
 	r.suppliers = append(r.suppliers, supplier)
 }
 
diff --git a/pkg/iac/supplier/IacChainSupplier_test.go b/pkg/iac/supplier/IacChainSupplier_test.go
@@ -11,22 +11,22 @@ import (
 func TestIacChainSupplier_Resources(t *testing.T) {
 	tests := []struct {
 		name          string
-		initSuppliers func(suppliers *[]resource.Supplier)
+		initSuppliers func(suppliers *[]resource.IaCSupplier)
 		want          []*resource.Resource
 		wantErr       bool
 	}{
 		{
 			name: "All failed",
-			initSuppliers: func(suppliers *[]resource.Supplier) {
-				sup := &resource.MockSupplier{}
+			initSuppliers: func(suppliers *[]resource.IaCSupplier) {
+				sup := &resource.MockIaCSupplier{}
 				sup.On("Resources").Return(nil, errors.New("1"))
 				*suppliers = append(*suppliers, sup)
 
-				sup = &resource.MockSupplier{}
+				sup = &resource.MockIaCSupplier{}
 				sup.On("Resources").Return(nil, errors.New("2"))
 				*suppliers = append(*suppliers, sup)
 
-				sup = &resource.MockSupplier{}
+				sup = &resource.MockIaCSupplier{}
 				sup.On("Resources").Return(nil, errors.New("3"))
 				*suppliers = append(*suppliers, sup)
 			},
@@ -35,16 +35,16 @@ func TestIacChainSupplier_Resources(t *testing.T) {
 		},
 		{
 			name: "Partial failed",
-			initSuppliers: func(suppliers *[]resource.Supplier) {
-				sup := &resource.MockSupplier{}
+			initSuppliers: func(suppliers *[]resource.IaCSupplier) {
+				sup := &resource.MockIaCSupplier{}
 				sup.On("Resources").Return(nil, errors.New("1"))
 				*suppliers = append(*suppliers, sup)
 
-				sup = &resource.MockSupplier{}
+				sup = &resource.MockIaCSupplier{}
 				sup.On("Resources").Return(nil, errors.New("2"))
 				*suppliers = append(*suppliers, sup)
 
-				sup = &resource.MockSupplier{}
+				sup = &resource.MockIaCSupplier{}
 				sup.On("Resources").Return([]*resource.Resource{
 					&resource.Resource{
 						Id:    "ID",
@@ -67,7 +67,7 @@ func TestIacChainSupplier_Resources(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			r := NewIacChainSupplier()
-			suppliers := make([]resource.Supplier, 0)
+			suppliers := make([]resource.IaCSupplier, 0)
 			tt.initSuppliers(&suppliers)
 
 			for _, supplier := range suppliers {
diff --git a/pkg/iac/supplier/supplier.go b/pkg/iac/supplier/supplier.go
@@ -37,7 +37,7 @@ func GetIACSupplier(configs []config.SupplierConfig,
 	progress output.Progress,
 	alerter *alerter.Alerter,
 	factory resource.ResourceFactory,
-	filter filter.Filter) (resource.Supplier, error) {
+	filter filter.Filter) (resource.IaCSupplier, error) {
 
 	chainSupplier := NewIacChainSupplier()
 	for _, config := range configs {
@@ -47,7 +47,7 @@ func GetIACSupplier(configs []config.SupplierConfig,
 
 		deserializer := resource.NewDeserializer(factory)
 
-		var supplier resource.Supplier
+		var supplier resource.IaCSupplier
 		var err error
 		switch config.Key {
 		case state.TerraformStateReaderSupplier:
diff --git a/pkg/iac/terraform/state/terraform_state_reader.go b/pkg/iac/terraform/state/terraform_state_reader.go
@@ -41,6 +41,7 @@ type TerraformStateReader struct {
 	progress       output.Progress
 	filter         filter.Filter
 	alerter        *alerter.Alerter
+	sourceCount    uint
 }
 
 func (r *TerraformStateReader) initReader() error {
@@ -49,7 +50,16 @@ func (r *TerraformStateReader) initReader() error {
 }
 
 func NewReader(config config.SupplierConfig, library *terraform.ProviderLibrary, backendOpts *backend.Options, progress output.Progress, alerter *alerter.Alerter, deserializer *resource.Deserializer, filter filter.Filter) (*TerraformStateReader, error) {
-	reader := TerraformStateReader{library: library, config: config, deserializer: deserializer, backendOptions: backendOpts, progress: progress, alerter: alerter, filter: filter}
+	reader := TerraformStateReader{
+		library:        library,
+		config:         config,
+		deserializer:   deserializer,
+		backendOptions: backendOpts,
+		progress:       progress,
+		alerter:        alerter,
+		filter:         filter,
+		sourceCount:    0,
+	}
 	err := reader.initReader()
 	if err != nil {
 		return nil, err
@@ -214,8 +224,13 @@ func (r *TerraformStateReader) Resources() ([]*resource.Resource, error) {
 	return r.retrieveMultiplesStates()
 }
 
+func (r *TerraformStateReader) SourceCount() uint {
+	return r.sourceCount
+}
+
 func (r *TerraformStateReader) retrieveForState(path string) ([]*resource.Resource, error) {
 	r.config.Path = path
+	r.sourceCount += 1
 	logrus.WithFields(logrus.Fields{
 		"path":    r.config.Path,
 		"backend": r.config.Backend,
diff --git a/pkg/resource/mock_IaCSupplier.go b/pkg/resource/mock_IaCSupplier.go
diff --git a/pkg/resource/supplier.go b/pkg/resource/supplier.go
@@ -1,10 +1,16 @@
 package resource
 
-// Resource Supplier supply the list of resource.Resource, its the front to retrieve remote resources
+// Supplier supply the list of resource.Resource, it's the main interface to retrieve remote resources
 type Supplier interface {
 	Resources() ([]*Resource, error)
 }
 
+// IaCSupplier supply the list of resource.Resource, it's the main interface to retrieve state resources
+type IaCSupplier interface {
+	Supplier
+	SourceCount() uint
+}
+
 type StoppableSupplier interface {
 	Supplier
 	Stop()
diff --git a/pkg/telemetry/telemetry.go b/pkg/telemetry/telemetry.go
@@ -20,6 +20,7 @@ type telemetry struct {
 	TotalManaged   int    `json:"total_managed"`
 	Duration       uint   `json:"duration"`
 	ProviderName   string `json:"provider_name"`
+	IaCSourceCount uint   `json:"iac_source_count"`
 }
 
 type Telemetry struct {
@@ -59,6 +60,10 @@ func (te Telemetry) SendTelemetry(store memstore.Bucket) {
 		t.ProviderName = val
 	}
 
+	if val, ok := store.Get("iac_source_count").(uint); ok {
+		t.IaCSourceCount = val
+	}
+
 	body, err := json.Marshal(t)
 	if err != nil {
 		logrus.Debug(err)
diff --git a/pkg/telemetry/telemetry_test.go b/pkg/telemetry/telemetry_test.go
@@ -44,12 +44,14 @@ func TestSendTelemetry(t *testing.T) {
 				TotalManaged:   1,
 				Duration:       123,
 				ProviderName:   "aws",
+				IaCSourceCount: 2,
 			},
 			setStoreValues: func(s memstore.Bucket, a *analyser.Analysis) {
 				s.Set("total_resources", a.Summary().TotalResources)
 				s.Set("total_managed", a.Summary().TotalManaged)
 				s.Set("duration", uint(a.Duration.Seconds()+0.5))
 				s.Set("provider_name", "aws")
+				s.Set("iac_source_count", uint(2))
 			},
 		},
 		{

Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@ import (`
`10`	`10`	`)`
`11`	`11`
`12`	`12`	`type IacChainSupplier struct {`
`13`		`- suppliers []resource.Supplier`
	`13`	`+ suppliers []resource.IaCSupplier`
`14`	`14`	`runner *parallel.ParallelRunner`
`15`	`15`	`}`
`16`	`16`
`@@ -20,7 +20,15 @@ func NewIacChainSupplier() *IacChainSupplier {`
`20`	`20`	`}`
`21`	`21`	`}`
`22`	`22`
`23`		`-func (r *IacChainSupplier) AddSupplier(supplier resource.Supplier) {`
	`23`	`+func (r *IacChainSupplier) SourceCount() uint {`
	`24`	`+ count := uint(0)`
	`25`	`+ for _, supplier := range r.suppliers {`
	`26`	`+ count += supplier.SourceCount()`
	`27`	`+ }`
	`28`	`+ return count`
	`29`	`+}`
	`30`	`+`
	`31`	`+func (r *IacChainSupplier) AddSupplier(supplier resource.IaCSupplier) {`
`24`	`32`	`r.suppliers = append(r.suppliers, supplier)`
`25`	`33`	`}`
`26`	`34`