Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update deps typing #23

Merged
merged 7 commits into from
May 23, 2024
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Run mypy on costmangement function
Iain-S committed May 23, 2024
commit c0dffffdbafa07e23fc10706837c1a9004c789b7
90 changes: 55 additions & 35 deletions usage_function/costmanagement/__init__.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,21 @@
"""An Azure function to collect cost-management information."""
import logging
from datetime import datetime, timedelta
from typing import Final

import azure.functions as func
import requests
from azure.identity import DefaultAzureCredential
from azure.mgmt.costmanagement import CostManagementClient
from azure.mgmt.costmanagement.models import QueryDefinition
from azure.mgmt.costmanagement.models import (
QueryDefinition,
QueryGrouping,
QueryDataset,
TimeframeType,
ExportType,
QueryAggregation,
QueryTimePeriod,
)

import utils.settings
from utils import models
@@ -18,30 +27,37 @@
format="%(levelname)s %(asctime)s: %(name)s - %(message)s",
datefmt="%d/%m/%Y %I:%M:%S %p",
)
logger = logging.getLogger(__name__)
logger: Final = logging.getLogger(__name__)

RETRY_ATTEMPTS = 5
RETRY_ATTEMPTS: Final = 5
# We should only need one set of credentials
CREDENTIALS = DefaultAzureCredential()
CREDENTIALS: Final = DefaultAzureCredential()

# The constant parts of the Azure Cost Management API query. To be appended
# with to/from times.
QUERY_TEMPLATE = {
"type": "ActualCost",
"timeframe": "Custom",
"dataset": {
"granularity": "None",
"aggregation": {
# TODO What's the difference between Cost and PreTaxCost, and which should
# we use?
"totalCost": {"name": "Cost", "function": "Sum"},
},
"grouping": [
{"type": "Dimension", "name": "SubscriptionId"},
{"type": "Dimension", "name": "SubscriptionName"},
],
QUERY_TYPE: Final = ExportType.ACTUAL_COST
QUERY_TIMEFRAME: Final = TimeframeType.CUSTOM
QUERY_DATASET: Final = QueryDataset(
granularity=None,
grouping=[
QueryGrouping(
type="Dimension",
name="SubscriptionId",
),
QueryGrouping(
type="Dimension",
name="SubscriptionName",
),
],
aggregation={
# TODO What's the difference between Cost and PreTaxCost,
# and which should we use?
"totalCost": QueryAggregation(
name="Cost",
function="Sum",
)
},
}
)


def _truncate_date(date_time):
@@ -90,24 +106,28 @@ def get_all_usage(start_datetime: datetime, end_datetime: datetime, mgmt_group:
window_start = covered_to + timedelta(days=1)
window_end = min(window_start + max_timeperiod, end_datetime)
parameters = QueryDefinition(
time_period={"from_property": window_start, "to": window_end},
**QUERY_TEMPLATE,
time_period=QueryTimePeriod(from_property=window_start, to=window_end),
dataset=QUERY_DATASET,
type=QUERY_TYPE,
timeframe=QUERY_TIMEFRAME,
)
new_data = cm_client.query.usage(scope=scope, parameters=parameters)
if new_data.next_link:
# TODO How to deal with paging?
msg = (
"Cost management query returned multiple pages of results. "
"The function app is not prepared to deal with this."
)
raise NotImplementedError(msg)
for amount, sub_id, name, currency in new_data.rows:
key = (sub_id, name, currency)
if key in data:
data[key] += amount
else:
data[key] = amount
covered_to = window_end
if new_data:
if new_data.next_link:
# TODO How to deal with paging?
msg = (
"Cost management query returned multiple pages of results. "
"The function app is not prepared to deal with this."
)
raise NotImplementedError(msg)
if new_data.rows:
for amount, sub_id, name, currency in new_data.rows:
key = (sub_id, name, currency)
if key in data:
data[key] += amount
else:
data[key] = amount
covered_to = window_end

# Convert data to a list of CMUsage objects.
all_usage = models.AllCMUsage(
1 change: 1 addition & 0 deletions usage_function/tests/mypy.ini
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
[mypy]
ignore_missing_imports = True
check_untyped_defs = True
plugins = pydantic.mypy

[pydantic-mypy]
8 changes: 4 additions & 4 deletions usage_function/tests/test_function_app.py
Original file line number Diff line number Diff line change
@@ -6,13 +6,13 @@
from uuid import UUID

from azure.mgmt.costmanagement.models import (
QueryDefinition,
QueryTimePeriod,
ExportType,
TimeframeType,
QueryAggregation,
QueryDataset,
QueryDefinition,
QueryGrouping,
QueryAggregation,
QueryTimePeriod,
TimeframeType,
)
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
6 changes: 3 additions & 3 deletions usage_function/tests/test_utils.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""Tests for function app utils."""
import logging
from datetime import datetime, timedelta, date
from datetime import date, datetime, timedelta
from typing import Final
from unittest import TestCase, main
from unittest.mock import MagicMock, call, patch
@@ -10,10 +10,10 @@
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import HttpUrl, TypeAdapter

import utils.usage
import utils.models
import utils.logutils
import utils.models
import utils.settings
import utils.usage

HTTP_ADAPTER: Final = TypeAdapter(HttpUrl)

4 changes: 2 additions & 2 deletions usage_function/utils/auth.py
Original file line number Diff line number Diff line change
@@ -36,9 +36,9 @@ def create_access_token(self):
access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)

expire = datetime.utcnow() + access_token_expires
token_claims.update({"exp": expire})
token_claims.update({"exp": str(expire)})

return jwt.encode(token_claims, self.private_key, algorithm="RS256")
return jwt.encode(token_claims, self.private_key, algorithm="RS256") # type: ignore

def __call__(self, r: requests.PreparedRequest) -> requests.PreparedRequest:
"""Add the bearer token to the request."""
64 changes: 0 additions & 64 deletions usage_function/utils/wrapper.py

This file was deleted.