Repository: cycloidio/terracost (go)

Issue:
aws: optional "minimal" ingestion mode
## Abstract

One of the ways we can improve AWS ingestion time (#13) is to support a "minimal" mode. In this mode, only a subset of all the data would be ingested, correlating to the resources we support. This will let library users save time and space by ingesting only around 10% (in the case of EC2) of the pricing file.

## Implementation details

This can be done by having a new AWS ingestion option, e.g. `WithIngestionFilter`, that would accept a filtering function. The function would be called for each pricing record, marking the items that should (and shouldn't) be ingested. Terracost might then define two such functions out-of-the-box: a `DefaultFilter` that ingests everything and a `MinimalFilter` that only ingests records necessary for the library to work.

The "minimal" mode shouldn't be the default, as it will require re-ingestion each time new resources or attributes of already existing resources become supported.

Pull request:
aws: add WithIngestionFilter option and implement MinimalFilter
WithIngestionFilter option can be used to control the data that is ingested into the database. By using MinimalFilter, only the minimal amount of data is ingested, saving time and storage space.

This PR closes #22.

Patch:
diff --git a/aws/filter.go b/aws/filter.go
new file mode 100644
index 0000000..e66baae
--- /dev/null
+++ b/aws/filter.go
@@ -0,0 +1,69 @@
+package aws
+
+import (
+	"github.com/cycloidio/terracost/price"
+)
+
+// IngestionFilter allows control over what pricing data is ingested. Given a price.WithProduct the function returns
+// true if the record should be ingested, false if it should be skipped.
+type IngestionFilter func(pp *price.WithProduct) bool
+
+// DefaultFilter ingests all the records without filtering.
+func DefaultFilter(_ *price.WithProduct) bool {
+	return true
+}
+
+// MinimalFilter only ingests the supported records, skipping those that would never be used.
+func MinimalFilter(pp *price.WithProduct) bool {
+	switch pp.Product.Service {
+	case "AmazonEC2":
+		return minimalFilterEC2(pp)
+	case "AmazonRDS":
+		return minimalFilterRDS(pp)
+	default:
+		return false
+	}
+}
+
+// minimalFilterEC2 only ingests storage-related records as well as compute records that match supported attributes.
+func minimalFilterEC2(pp *price.WithProduct) bool {
+	switch pp.Product.Family {
+	case "Compute Instance":
+		allowedProductAttrs := map[string][]string{
+			"capacitystatus":  {"Used"},
+			"operatingSystem": {"Linux"},
+			"preInstalledSw":  {"NA"},
+			"tenancy":         {"Shared", "Dedicated"},
+		}
+		for k, vals := range allowedProductAttrs {
+			if !isValueAllowed(pp.Product.Attributes[k], vals) {
+				return false
+			}
+		}
+		return true
+	case "Storage", "System Operation":
+		return true
+	default:
+		return false
+	}
+}
+
+// minimalFilterRDS only ingests RDS records of supported product families.
+func minimalFilterRDS(pp *price.WithProduct) bool {
+	switch pp.Product.Family {
+	case "Database Instance", "Database Storage", "Provisioned IOPS":
+		return true
+	default:
+		return false
+	}
+}
+
+// isValueAllowed returns true if the allowed slice contains the value.
+func isValueAllowed(value string, allowed []string) bool {
+	for _, v := range allowed {
+		if value == v {
+			return true
+		}
+	}
+	return false
+}
diff --git a/aws/ingester.go b/aws/ingester.go
index 17ae156..482d347 100644
--- a/aws/ingester.go
+++ b/aws/ingester.go
@@ -39,6 +39,8 @@ type Ingester struct {
 	progressCh       chan<- progress.Progress
 	progressInterval time.Duration
 
+	ingestionFilter IngestionFilter
+
 	err error
 }
 
@@ -46,12 +48,13 @@ type Ingester struct {
 // WithRepositories function, other configuration options will use their default values.
 func NewIngester(service, region string, options ...Option) *Ingester {
 	ing := &Ingester{
-		httpClient: &http.Client{},
-		pricingURL: defaultPricingURL,
-		bufferSize: defaultBufferSize,
-		service:    service,
-		region:     region,
-		progressCh: nil,
+		httpClient:      &http.Client{},
+		pricingURL:      defaultPricingURL,
+		bufferSize:      defaultBufferSize,
+		service:         service,
+		region:          region,
+		progressCh:      nil,
+		ingestionFilter: DefaultFilter,
 	}
 	for _, opt := range options {
 		opt(ing)
@@ -141,7 +144,9 @@ func (ing *Ingester) Ingest(ctx context.Context, chSize int) <-chan *price.WithP
 				return
 			}
 
-			results <- pp
+			if ing.ingestionFilter(pp) {
+				results <- pp
+			}
 		}
 	}()
 
diff --git a/aws/options.go b/aws/options.go
index 44ba848..448fb07 100644
--- a/aws/options.go
+++ b/aws/options.go
@@ -45,3 +45,10 @@ func WithProgress(progressCh chan<- progress.Progress, interval time.Duration) O
 		ing.progressInterval = interval
 	}
 }
+
+// WithIngestionFilter sets a custom IngestionFilter to control which pricing data records should be ingested.
+func WithIngestionFilter(filter IngestionFilter) Option {
+	return func(ing *Ingester) {
+		ing.ingestionFilter = filter
+	}
+}


Test patch:
diff --git a/aws/filter_test.go b/aws/filter_test.go
new file mode 100644
index 0000000..0118709
--- /dev/null
+++ b/aws/filter_test.go
@@ -0,0 +1,78 @@
+package aws
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/cycloidio/terracost/price"
+	"github.com/cycloidio/terracost/product"
+)
+
+func TestMinimalFilter(t *testing.T) {
+	t.Run("Allowed", func(t *testing.T) {
+		pps := []*price.WithProduct{
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Storage"}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "System Operation"}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Compute Instance", Attributes: map[string]string{
+				"capacitystatus":  "Used",
+				"operatingSystem": "Linux",
+				"preInstalledSw":  "NA",
+				"tenancy":         "Shared",
+			}}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Compute Instance", Attributes: map[string]string{
+				"capacitystatus":  "Used",
+				"operatingSystem": "Linux",
+				"preInstalledSw":  "NA",
+				"tenancy":         "Dedicated",
+			}}},
+			{Product: &product.Product{Service: "AmazonRDS", Family: "Database Instance"}},
+			{Product: &product.Product{Service: "AmazonRDS", Family: "Database Storage"}},
+			{Product: &product.Product{Service: "AmazonRDS", Family: "Provisioned IOPS"}},
+		}
+
+		for i, pp := range pps {
+			assert.True(t, MinimalFilter(pp), "case %d", i)
+		}
+	})
+
+	t.Run("Skipped", func(t *testing.T) {
+		pps := []*price.WithProduct{
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Compute Instance (bare metal)"}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "NAT Gateway"}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Fee"}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Compute Instance", Attributes: map[string]string{
+				"capacitystatus":  "Used",
+				"operatingSystem": "Linux",
+				"preInstalledSw":  "NA",
+				"tenancy":         "Host",
+			}}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Compute Instance", Attributes: map[string]string{
+				"capacitystatus":  "Used",
+				"operatingSystem": "SUSE",
+				"preInstalledSw":  "NA",
+				"tenancy":         "Shared",
+			}}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Compute Instance", Attributes: map[string]string{
+				"capacitystatus":  "Used",
+				"operatingSystem": "Linux",
+				"preInstalledSw":  "SQL Server",
+				"tenancy":         "Shared",
+			}}},
+			{Product: &product.Product{Service: "AmazonEC2", Family: "Compute Instance", Attributes: map[string]string{
+				"capacitystatus":  "AllocatedCapacityReservation",
+				"operatingSystem": "Linux",
+				"preInstalledSw":  "NA",
+				"tenancy":         "Shared",
+			}}},
+			{Product: &product.Product{Service: "AmazonRDS", Family: "CPU Credits"}},
+			{Product: &product.Product{Service: "AmazonRDS", Family: "Performance Insights"}},
+			{Product: &product.Product{Service: "AmazonRDS", Family: "RDSProxy"}},
+			{Product: &product.Product{Service: "AmazonRDS", Family: "Serverless"}},
+		}
+
+		for i, pp := range pps {
+			assert.False(t, MinimalFilter(pp), "case %d", i)
+		}
+	})
+}