fix(aws): include marketplace fees (#4028)

Signed-off-by: Kush Agarwal <agrawalkush783@gmail.com>
Co-authored-by: Warwick <warwick.peatey@ibm.com>

pkg/cloud/aws/athenaintegration.go

@@ -57,8 +57,30 @@ var AthenaNetSPPricingCoalesce = fmt.Sprintf("COALESCE(%s, %s, 0)", AthenaNetSPP
 const AthenaDateColumn = "line_item_usage_start_date"
 const AthenaDateTruncColumn = "DATE_TRUNC('day'," + AthenaDateColumn + ") as usage_date"
 
+// AthenaBillingEntityColumn distinguishes standard AWS charges ('AWS') from AWS
+// Marketplace charges ('AWS Marketplace') on a CUR line item.
+const AthenaBillingEntityColumn = "bill_billing_entity"
+const AthenaMarketplaceBillingEntity = "AWS Marketplace"
+
 const AthenaWhereDateFmt = `line_item_usage_start_date >= date '%s' AND line_item_usage_start_date < date '%s'`
-const AthenaWhereUsage = "(line_item_line_item_type = 'Usage' OR line_item_line_item_type = 'DiscountedUsage' OR line_item_line_item_type = 'SavingsPlanCoveredUsage' OR line_item_line_item_type = 'EdpDiscount' OR line_item_line_item_type = 'PrivateRateDiscount')"
+
+// AthenaWhereUsageBase filters to usage-driving line item types only. It references no
+// optional CUR columns, so it is always safe to use regardless of which columns a given
+// CUR export includes.
+const AthenaWhereUsageBase = "(line_item_line_item_type = 'Usage' OR line_item_line_item_type = 'DiscountedUsage' OR line_item_line_item_type = 'SavingsPlanCoveredUsage' OR line_item_line_item_type = 'EdpDiscount' OR line_item_line_item_type = 'PrivateRateDiscount')"
+
+// AthenaWhereUsage extends AthenaWhereUsageBase with AWS Marketplace 'Fee' line items
+// (flat-rate/subscription charges for third-party SaaS products). Marketplace is scoped
+// to bill_billing_entity = 'AWS Marketplace' so this does not also pull in
+// non-Marketplace 'Fee' rows, such as Reserved Instance upfront purchases, which are
+// outside the scope of this Marketplace-specific fix. CUR 2.0 exports can disable any
+// column, including bill_billing_entity, so callers must only use this filter when
+// AthenaBillingEntityColumn is confirmed present (see getCloudCost) -- otherwise the
+// query will fail with COLUMN_NOT_FOUND and fall back to AthenaWhereUsageBase instead.
+var AthenaWhereUsage = fmt.Sprintf(
+	"(line_item_line_item_type = 'Usage' OR line_item_line_item_type = 'DiscountedUsage' OR line_item_line_item_type = 'SavingsPlanCoveredUsage' OR line_item_line_item_type = 'EdpDiscount' OR line_item_line_item_type = 'PrivateRateDiscount' OR (line_item_line_item_type = 'Fee' AND %s = '%s'))",
+	AthenaBillingEntityColumn, AthenaMarketplaceBillingEntity,
+)
 
 // AthenaQueryIndexes is a struct for holding the context of a query
 type AthenaQueryIndexes struct {
@@ -186,13 +208,13 @@ func (ai *AthenaIntegration) getCloudCost(start, end time.Time, limit int) (*ope
 	whereDate := fmt.Sprintf(AthenaWhereDateFmt, start.Format("2006-01-02"), end.Format("2006-01-02"))
 	wherePartitions := ai.GetPartitionWhere(start, end, isCUR20(allColumns))
 
-	// Query for all line items with a resource_id or from AWS Marketplace, which did not end before
-	// the range or start after it. This captures all costs with any amount of
-	// overlap with the range, for which we will only extract the relevant costs
+	// Query for all line items whose usage start date falls within the given range and
+	// partition, restricted to usage-driving line item types and, when the
+	// bill_billing_entity column exists, AWS Marketplace fees (see GetWhereUsage).
 	whereConjuncts := []string{
 		wherePartitions,
 		whereDate,
-		AthenaWhereUsage,
+		ai.GetWhereUsage(allColumns),
 	}
 	columnStr := strings.Join(selectColumns, ", ")
 	whereClause := strings.Join(whereConjuncts, " AND ")
@@ -246,6 +268,19 @@ func (ai *AthenaIntegration) GetListCostColumn() string {
 	return fmt.Sprintf("SUM(%s) as list_cost", listCostBuilder.String())
 }
 
+// GetWhereUsage returns the usage-type filter to apply to the CUR query. When the CUR
+// export includes bill_billing_entity, AWS Marketplace 'Fee' line items are included
+// alongside the usual usage-driving types (see AthenaWhereUsage). CUR 2.0 exports can
+// disable any column, so when bill_billing_entity is absent this falls back to
+// AthenaWhereUsageBase -- referencing a missing column would otherwise fail the entire
+// query with COLUMN_NOT_FOUND, not just omit Marketplace fees.
+func (ai *AthenaIntegration) GetWhereUsage(allColumns map[string]bool) string {
+	if allColumns[AthenaBillingEntityColumn] {
+		return AthenaWhereUsage
+	}
+	return AthenaWhereUsageBase
+}
+
 func (ai *AthenaIntegration) GetNetCostColumn(allColumns map[string]bool) string {
 	netCostColumn := ""
 	if allColumns[AthenaNetPricingColumn] { // if Net pricing exists

pkg/cloud/aws/athenaintegration_coverage_test.go

@@ -7,6 +7,86 @@ import (
 	"github.com/opencost/opencost/pkg/cloud"
 )
 
+func TestAthenaWhereUsage(t *testing.T) {
+	// Regression test for https://github.com/opencost/opencost/issues/4022 -
+	// AWS Marketplace subscription/fee charges (line_item_line_item_type = 'Fee')
+	// must be included, but only when billed through AWS Marketplace, so that
+	// non-Marketplace 'Fee' rows (e.g. Reserved Instance upfront purchases,
+	// already captured via 'DiscountedUsage' amortization) are not swept in and
+	// double counted.
+	expected := "(line_item_line_item_type = 'Usage' OR line_item_line_item_type = 'DiscountedUsage' OR line_item_line_item_type = 'SavingsPlanCoveredUsage' OR line_item_line_item_type = 'EdpDiscount' OR line_item_line_item_type = 'PrivateRateDiscount' OR (line_item_line_item_type = 'Fee' AND bill_billing_entity = 'AWS Marketplace'))"
+	if AthenaWhereUsage != expected {
+		t.Errorf("AthenaWhereUsage = %v, want %v", AthenaWhereUsage, expected)
+	}
+
+	if !strings.Contains(AthenaWhereUsage, "line_item_line_item_type = 'Fee' AND bill_billing_entity = 'AWS Marketplace'") {
+		t.Errorf("AthenaWhereUsage should include AWS Marketplace 'Fee' rows scoped to bill_billing_entity, got: %v", AthenaWhereUsage)
+	}
+
+	// A bare, unscoped 'Fee' disjunct would also match non-Marketplace fees like RI
+	// upfront purchases, so it must never appear on its own.
+	if strings.Contains(AthenaWhereUsage, "line_item_line_item_type = 'Fee')") {
+		t.Errorf("AthenaWhereUsage should not include an unscoped 'Fee' clause, got: %v", AthenaWhereUsage)
+	}
+}
+
+func TestAthenaWhereUsageBase(t *testing.T) {
+	// AthenaWhereUsageBase must reference only mandatory CUR columns (never
+	// bill_billing_entity, which CUR 2.0 exports can disable) so it is always a safe
+	// fallback when that column is absent.
+	expected := "(line_item_line_item_type = 'Usage' OR line_item_line_item_type = 'DiscountedUsage' OR line_item_line_item_type = 'SavingsPlanCoveredUsage' OR line_item_line_item_type = 'EdpDiscount' OR line_item_line_item_type = 'PrivateRateDiscount')"
+	if AthenaWhereUsageBase != expected {
+		t.Errorf("AthenaWhereUsageBase = %v, want %v", AthenaWhereUsageBase, expected)
+	}
+
+	if strings.Contains(AthenaWhereUsageBase, "bill_billing_entity") || strings.Contains(AthenaWhereUsageBase, "Fee") {
+		t.Errorf("AthenaWhereUsageBase must not reference bill_billing_entity or 'Fee', got: %v", AthenaWhereUsageBase)
+	}
+}
+
+func TestAthenaIntegration_GetWhereUsage(t *testing.T) {
+	ai := &AthenaIntegration{}
+
+	// Regression test for https://github.com/opencost/opencost/pull/4028#discussion -
+	// bill_billing_entity is a CUR 2.0-disableable column, like resource_tags,
+	// line_item_usage_account_name, and bill_payer_account_name elsewhere in this file.
+	// GetWhereUsage must check allColumns before referencing it, or CUR exports that
+	// omit the column would fail every query with COLUMN_NOT_FOUND instead of merely
+	// missing Marketplace fees.
+	t.Run("bill_billing_entity present", func(t *testing.T) {
+		allColumns := map[string]bool{
+			AthenaBillingEntityColumn: true,
+		}
+		got := ai.GetWhereUsage(allColumns)
+		if got != AthenaWhereUsage {
+			t.Errorf("GetWhereUsage() = %v, want AthenaWhereUsage (%v)", got, AthenaWhereUsage)
+		}
+		if !strings.Contains(got, "bill_billing_entity = 'AWS Marketplace'") {
+			t.Errorf("GetWhereUsage() should include the Marketplace fee clause when bill_billing_entity is present, got: %v", got)
+		}
+	})
+
+	t.Run("bill_billing_entity absent", func(t *testing.T) {
+		allColumns := map[string]bool{
+			"line_item_line_item_type": true,
+		}
+		got := ai.GetWhereUsage(allColumns)
+		if got != AthenaWhereUsageBase {
+			t.Errorf("GetWhereUsage() = %v, want AthenaWhereUsageBase (%v)", got, AthenaWhereUsageBase)
+		}
+		if strings.Contains(got, "bill_billing_entity") {
+			t.Errorf("GetWhereUsage() must not reference bill_billing_entity when it is absent from allColumns, got: %v", got)
+		}
+	})
+
+	t.Run("empty allColumns", func(t *testing.T) {
+		got := ai.GetWhereUsage(map[string]bool{})
+		if got != AthenaWhereUsageBase {
+			t.Errorf("GetWhereUsage() = %v, want AthenaWhereUsageBase (%v)", got, AthenaWhereUsageBase)
+		}
+	})
+}
+
 func TestAthenaIntegration_GetListCostColumn(t *testing.T) {
 	ai := &AthenaIntegration{}
 	expected := "SUM(CASE line_item_line_item_type WHEN 'EdpDiscount' THEN 0 WHEN 'PrivateRateDiscount' THEN 0 ELSE line_item_unblended_cost END) as list_cost"