Add cloud method to get reporting data

Adds a method to `azure_cloud_provider` to query the Cost Management API
for usage data per invoice.  For now, this query is relatively static.
We're always calling the API at the billing invoice section scope,
with the widest timeframe possible (one year), and with the same
requested dataset. As the scope of the application's reporting needs
changes, this function may change to be more general and/or revert back
to the SDK.

atst/domain/csp/cloud/azure_cloud_provider.py

@@ -24,6 +24,7 @@ from .models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
+    CostManagementQueryCSPResult,
     KeyVaultCredentials,
     ManagementGroupCSPResponse,
     ProductPurchaseCSPPayload,
@@ -32,6 +33,7 @@ from .models import (
     ProductPurchaseVerificationCSPResult,
     PrincipalAdminRoleCSPPayload,
     PrincipalAdminRoleCSPResult,
+    ReportingCSPPayload,
     TaskOrderBillingCreationCSPPayload,
     TaskOrderBillingCreationCSPResult,
     TaskOrderBillingVerificationCSPPayload,
@@ -1070,3 +1072,41 @@ class AzureCloudProvider(CloudProviderInterface):
         hashed = sha256_hex(tenant_id)
         raw_creds = self.get_secret(hashed)
         return KeyVaultCredentials(**json.loads(raw_creds))
+
+    def get_reporting_data(self, payload: ReportingCSPPayload):
+        """
+        Queries the Cost Management API for an invoice section's raw reporting data
+        
+        We query at the invoiceSection scope. The full scope path is passed in
+        with the payload at the `invoice_section_id` key.
+        """
+        creds = self._source_tenant_creds(payload.tenant_id)
+        token = self._get_sp_token(
+            payload.tenant_id, creds.tenant_sp_client_id, creds.tenant_sp_key
+        )
+
+        if not token:
+            raise AuthenticationException("Could not retrieve tenant access token")
+
+        headers = {"Authorization": f"Bearer {token}"}
+
+        request_body = {
+            "type": "Usage",
+            "timeframe": "Custom",
+            "timePeriod": {"from": payload.from_date, "to": payload.to_date,},
+            "dataset": {
+                "granularity": "Daily",
+                "aggregation": {"totalCost": {"name": "PreTaxCost", "function": "Sum"}},
+                "grouping": [{"type": "Dimension", "name": "InvoiceId"}],
+            },
+        }
+        cost_mgmt_url = (
+            f"/providers/Microsoft.CostManagement/query?api-version=2019-11-01"
+        )
+        result = self.sdk.requests.post(
+            f"{self.sdk.cloud.endpoints.resource_manager}{payload.invoice_section_id}{cost_mgmt_url}",
+            json=request_body,
+            headers=headers,
+        )
+        if result.ok:
+            return CostManagementQueryCSPResult(**result.json())

atst/domain/csp/cloud/mock_cloud_provider.py

@@ -25,12 +25,15 @@ from .models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
+    CostManagementQueryCSPResult,
+    CostManagementQueryProperties,
     ProductPurchaseCSPPayload,
     ProductPurchaseCSPResult,
     ProductPurchaseVerificationCSPPayload,
     ProductPurchaseVerificationCSPResult,
     PrincipalAdminRoleCSPPayload,
     PrincipalAdminRoleCSPResult,
+    ReportingCSPPayload,
     SubscriptionCreationCSPPayload,
     SubscriptionCreationCSPResult,
     SubscriptionVerificationCSPPayload,
@@ -487,3 +490,25 @@ class MockCloudProvider(CloudProviderInterface):
 
     def update_tenant_creds(self, tenant_id, secret):
         return secret
+
+    def get_reporting_data(self, payload: ReportingCSPPayload):
+        self._maybe_raise(self.NETWORK_FAILURE_PCT, self.NETWORK_EXCEPTION)
+        self._maybe_raise(self.SERVER_FAILURE_PCT, self.SERVER_EXCEPTION)
+        self._maybe_raise(self.UNAUTHORIZED_RATE, self.AUTHORIZATION_EXCEPTION)
+        object_id = str(uuid4())
+
+        properties = CostManagementQueryProperties(
+            **dict(
+                columns=[
+                    {"name": "PreTaxCost", "type": "Number"},
+                    {"name": "UsageDate", "type": "Number"},
+                    {"name": "InvoiceId", "type": "String"},
+                    {"name": "Currency", "type": "String"},
+                ],
+                rows=[],
+            )
+        )
+
+        return CostManagementQueryCSPResult(
+            **dict(name=object_id, properties=properties,)
+        )

atst/domain/csp/cloud/models.py

@@ -499,3 +499,34 @@ class UserCSPPayload(BaseCSPPayload):
 
 class UserCSPResult(AliasModel):
     id: str
+
+
+class QueryColumn(AliasModel):
+    name: str
+    type: str
+
+
+class CostManagementQueryProperties(AliasModel):
+    columns: List[QueryColumn]
+    rows: List[Optional[list]]
+
+
+class CostManagementQueryCSPResult(AliasModel):
+    name: str
+    properties: CostManagementQueryProperties
+
+
+class ReportingCSPPayload(BaseCSPPayload):
+    invoice_section_id: str
+    from_date: str
+    to_date: str
+
+    @root_validator(pre=True)
+    def extract_invoice_section(cls, values):
+        try:
+            values["invoice_section_id"] = values["billing_profile_properties"][
+                "invoice_sections"
+            ][0]["invoice_section_id"]
+            return values
+        except (KeyError, IndexError):
+            raise ValueError("Invoice section ID not present in payload")

tests/domain/cloud/test_azure_csp.py

@@ -5,6 +5,8 @@ from uuid import uuid4
 import pytest
 from tests.factories import ApplicationFactory, EnvironmentFactory
 from tests.mock_azure import AUTH_CREDENTIALS, mock_azure
+import pendulum
+import pydantic
 
 from atst.domain.csp.cloud import AzureCloudProvider
 from atst.domain.csp.cloud.models import (
@@ -20,10 +22,12 @@ from atst.domain.csp.cloud.models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
+    CostManagementQueryCSPResult,
     ProductPurchaseCSPPayload,
     ProductPurchaseCSPResult,
     ProductPurchaseVerificationCSPPayload,
     ProductPurchaseVerificationCSPResult,
+    ReportingCSPPayload,
     SubscriptionCreationCSPPayload,
     SubscriptionCreationCSPResult,
     SubscriptionVerificationCSPPayload,
@@ -718,3 +722,77 @@ def test_create_subscription_verification(mock_azure: AzureCloudProvider):
             payload
         )
         assert result.subscription_id == "60fbbb72-0516-4253-ab18-c92432ba3230"
+
+
+def test_get_reporting_data(mock_azure: AzureCloudProvider):
+    mock_result = Mock()
+    mock_result.json.return_value = {
+        "eTag": None,
+        "id": "providers/Microsoft.Billing/billingAccounts/52865e4c-52e8-5a6c-da6b-c58f0814f06f:7ea5de9d-b8ce-4901-b1c5-d864320c7b03_2019-05-31/billingProfiles/XQDJ-6LB4-BG7-TGB/invoiceSections/P73M-XC7J-PJA-TGB/providers/Microsoft.CostManagement/query/e82d0cda-2ffb-4476-a98a-425c83c216f9",
+        "location": None,
+        "name": "e82d0cda-2ffb-4476-a98a-425c83c216f9",
+        "properties": {
+            "columns": [
+                {"name": "PreTaxCost", "type": "Number"},
+                {"name": "UsageDate", "type": "Number"},
+                {"name": "InvoiceId", "type": "String"},
+                {"name": "Currency", "type": "String"},
+            ],
+            "nextLink": None,
+            "rows": [],
+        },
+        "sku": None,
+        "type": "Microsoft.CostManagement/query",
+    }
+    mock_result.ok = True
+    mock_azure.sdk.requests.post.return_value = mock_result
+    mock_azure = mock_get_secret(mock_azure)
+
+    # Subset of a profile's CSP data that we care about for reporting
+    csp_data = {
+        "tenant_id": "6d2d2d6c-a6d6-41e1-8bb1-73d11475f8f4",
+        "billing_profile_properties": {
+            "invoice_sections": [
+                {
+                    "invoice_section_id": "providers/Microsoft.Billing/billingAccounts/52865e4c-52e8-5a6c-da6b-c58f0814f06f:7ea5de9d-b8ce-4901-b1c5-d864320c7b03_2019-05-31/billingProfiles/XQDJ-6LB4-BG7-TGB/invoiceSections/P73M-XC7J-PJA-TGB",
+                }
+            ],
+        },
+    }
+
+    data: CostManagementQueryCSPResult = mock_azure.get_reporting_data(
+        ReportingCSPPayload(
+            from_date=pendulum.now().subtract(years=1).add(days=1).format("YYYY-MM-DD"),
+            to_date=pendulum.now().format("YYYY-MM-DD"),
+            **csp_data,
+        )
+    )
+
+    assert isinstance(data, CostManagementQueryCSPResult)
+    assert data.name == "e82d0cda-2ffb-4476-a98a-425c83c216f9"
+    assert len(data.properties.columns) == 4
+
+
+def test_get_reporting_data_malformed_payload(mock_azure: AzureCloudProvider):
+    mock_result = Mock()
+    mock_result.ok = True
+    mock_azure.sdk.requests.post.return_value = mock_result
+    mock_azure = mock_get_secret(mock_azure)
+
+    # Malformed csp_data payloads that should throw pydantic validation errors
+    index_error = {
+        "tenant_id": "6d2d2d6c-a6d6-41e1-8bb1-73d11475f8f4",
+        "billing_profile_properties": {"invoice_sections": [],},
+    }
+    key_error = {
+        "tenant_id": "6d2d2d6c-a6d6-41e1-8bb1-73d11475f8f4",
+        "billing_profile_properties": {"invoice_sections": [{}],},
+    }
+
+    for malformed_payload in [key_error, index_error]:
+        with pytest.raises(pydantic.ValidationError):
+            assert mock_azure.get_reporting_data(
+                ReportingCSPPayload(
+                    from_date="foo", to_date="bar", **malformed_payload,
+                )
+            )