Skip to content

Commit

Permalink
change over to export API
Browse files Browse the repository at this point in the history
  • Loading branch information
wpbrown committed Sep 4, 2020
1 parent b898c9f commit d55fb2a
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 138 deletions.
2 changes: 0 additions & 2 deletions requirements.in
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
azure-cli-core
azure-mgmt-billing
azure-mgmt-consumption
azure-storage-blob
click
16 changes: 6 additions & 10 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,40 +2,37 @@
# This file is autogenerated by pip-compile
# To update, run:
#
# pip-compile
# pip-compile requirements.in
#
adal==1.2.4 # via azure-cli-core, msrestazure
applicationinsights==0.11.9 # via azure-cli-telemetry
argcomplete==1.11.1 # via azure-cli-core, knack
azure-cli-core==2.8.0 # via -r requirements.in
azure-cli-nspkg==3.0.4 # via azure-cli-core, azure-cli-telemetry
azure-cli-telemetry==1.0.4 # via azure-cli-core
azure-common==1.1.25 # via azure-mgmt-billing, azure-mgmt-consumption, azure-mgmt-resource
azure-core==1.6.0 # via azure-mgmt-core, azure-storage-blob
azure-common==1.1.25 # via azure-mgmt-billing, azure-mgmt-resource
azure-core==1.6.0 # via azure-mgmt-core
azure-mgmt-billing==0.2.0 # via -r requirements.in
azure-mgmt-consumption==3.0.0 # via -r requirements.in
azure-mgmt-core==1.0.0 # via azure-cli-core
azure-mgmt-nspkg==3.0.2 # via azure-mgmt-billing
azure-mgmt-resource==10.0.0 # via azure-cli-core
azure-nspkg==3.0.2 # via azure-cli-nspkg, azure-mgmt-nspkg
azure-storage-blob==12.3.2 # via -r requirements.in
bcrypt==3.1.7 # via paramiko
certifi==2020.6.20 # via msrest, requests
cffi==1.14.0 # via bcrypt, cryptography, pynacl
chardet==3.0.4 # via requests
click==7.1.2 # via -r requirements.in
colorama==0.4.3 # via azure-cli-core, knack
cryptography==2.9.2 # via adal, azure-storage-blob, paramiko, pyjwt, pyopenssl
cryptography==2.9.2 # via adal, paramiko, pyjwt, pyopenssl
humanfriendly==8.2 # via azure-cli-core
idna==2.10 # via requests
importlib-metadata==1.7.0 # via argcomplete
isodate==0.6.0 # via msrest
jmespath==0.10.0 # via azure-cli-core, knack
knack==0.7.1 # via azure-cli-core
msal-extensions==0.1.3 # via azure-cli-core
msal==1.0.0 # via azure-cli-core, msal-extensions
msrest==0.6.17 # via azure-cli-core, azure-mgmt-consumption, azure-mgmt-resource, azure-storage-blob, msrestazure
msrestazure==0.6.4 # via azure-cli-core, azure-mgmt-billing, azure-mgmt-consumption, azure-mgmt-resource
msrest==0.6.17 # via azure-cli-core, azure-mgmt-resource, msrestazure
msrestazure==0.6.4 # via azure-cli-core, azure-mgmt-billing, azure-mgmt-resource
oauthlib==3.1.0 # via requests-oauthlib
paramiko==2.7.1 # via azure-cli-core
pkginfo==1.5.0.1 # via azure-cli-core
Expand All @@ -52,4 +49,3 @@ requests==2.24.0 # via adal, azure-cli-core, azure-core, msal, msrest,
six==1.15.0 # via azure-cli-core, azure-core, bcrypt, cryptography, isodate, knack, msrestazure, pynacl, pyopenssl, python-dateutil
tabulate==0.8.7 # via knack
urllib3==1.25.9 # via requests
zipp==3.1.0 # via importlib-metadata
125 changes: 63 additions & 62 deletions src/azmpcli/__main__.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,14 @@
import datetime
import itertools
import json
import time
import uuid
from typing import Any, Collection, Iterable, List

import click
from azure.common.client_factory import get_client_from_cli_profile
from azure.common.credentials import get_azure_cli_credentials, get_cli_profile
from azure.common.credentials import get_cli_profile
from azure.mgmt.billing import BillingManagementClient
from azure.mgmt.billing.models import BillingPeriod
from azure.mgmt.consumption import ConsumptionManagementClient
from azure.storage.blob import BlobServiceClient

from . import _patch # noqa: F401

Expand Down Expand Up @@ -40,20 +38,60 @@ def get_billing_periods(client: BillingManagementClient, names: Iterable[str]) -
return selected_periods


def generate_usage_blob_data(
client: ConsumptionManagementClient, billing_account_name: str, billing_period: str
def generate_onetime_export(
client: BillingManagementClient,
billing_account_name: str,
billing_period: BillingPeriod,
storage_account_resource_id: str,
) -> str:
download_operation = client.usage_details.download(
"/providers/Microsoft.Billing/billingAccounts/{}/providers/Microsoft.Billing/"
"billingPeriods/{}".format(billing_account_name, billing_period),
metric="amortizedcost",
service_client = client._client
name = f"onetime{str(uuid.uuid1()).replace('-', '')}"
url = service_client.format_url(
"/providers/Microsoft.Billing/billingAccounts/{enrollmentId}"
"/providers/Microsoft.CostManagement/exports/{name}",
enrollmentId=billing_account_name,
name=name,
)
while not download_operation.done():
download_operation.wait(30)
print("Generate data status: {}".format(download_operation.status()))
download_result = download_operation.result()
print("Got URL to blob: {}".format(download_result.download_url))
return download_result.download_url
query_parameters = {"api-version": "2020-06-01"}
header_parameters = {"Content-Type": "application/json"}
content = {
"properties": {
"definition": {
"dataSet": {"granularity": "Daily"},
"timePeriod": {
"from": f"{billing_period.billing_period_start_date.strftime('%Y-%m-%d')}T00:00:00Z",
"to": f"{billing_period.billing_period_end_date.strftime('%Y-%m-%d')}T23:59:59Z",
},
"timeframe": "Custom",
"type": "AmortizedCost",
},
"deliveryInfo": {
"destination": {
"container": "usage-final",
"resourceId": storage_account_resource_id,
"rootFolderPath": "export",
}
},
"format": "Csv",
"schedule": {"status": "Inactive"},
}
}
request = service_client.put(url, query_parameters, header_parameters, content)
response = service_client.send(request, stream=False)
if response.status_code != 201:
raise Exception("Failed to create export.")
result = json.loads(response.content)
return result["id"]


def start_onetime_export(client: BillingManagementClient, export_resource_id: str) -> None:
service_client = client._client
url = service_client.format_url("/{resourceId}/run", resourceId=export_resource_id)
query_parameters = {"api-version": "2020-06-01"}
request = service_client.post(url, query_parameters)
response = service_client.send(request, stream=False)
if response.status_code != 200:
raise Exception("Failed to start export.")


def get_billing_accounts(client: BillingManagementClient) -> List[str]:
Expand All @@ -79,23 +117,19 @@ def get_azure_cli_credentials_non_default_sub(resource: str, subscription: str)


@click.command()
@click.option("-s", "--storage", "storage_account_name", help="Storage account name.", required=True)
@click.option(
"--storage-subscription",
"storage_account_subscription",
help="CLI account subscription to access storage (not required).",
"-s", "--storage", "storage_account_resource_id", help="Storage account resource id.", required=True
)
@click.option(
"-a", "--account", "billing_account_name", help="EA billing account number.", show_default="Auto-detect"
)
@click.argument("billing_periods", nargs=-1)
def cli(
storage_account_name: str,
storage_account_resource_id: str,
billing_account_name: str,
billing_periods: Collection[str],
storage_account_subscription: str,
) -> None:
billing_client = get_client_from_cli_profile(BillingManagementClient)
billing_client: BillingManagementClient = get_client_from_cli_profile(BillingManagementClient)

if billing_account_name is None:
accounts = get_billing_accounts(billing_client)
Expand Down Expand Up @@ -123,46 +157,13 @@ def cli(
period.billing_period_end_date.strftime("%Y%m%d"),
)

print("Generating usage data (this can take 5 to 10 minutes)...")
cm_client = get_client_from_cli_profile(ConsumptionManagementClient)
generated_blob_url = generate_usage_blob_data(cm_client, billing_account_name, period.name)

blob_account_url = "https://{}.blob.core.windows.net/".format(storage_account_name)
storage_resource = "https://storage.azure.com/"
if storage_account_subscription is None:
credential, _ = get_azure_cli_credentials(resource=storage_resource)
else:
credential = get_azure_cli_credentials_non_default_sub(
resource=storage_resource, subscription=storage_account_subscription
)
service = BlobServiceClient(account_url=blob_account_url, credential=credential)
container = service.get_container_client("usage-final")
blob = container.get_blob_client("export/finalamortized/{}/manual_load.csv".format(export_label))
blob.start_copy_from_url(generated_blob_url)
while True:
props = blob.get_blob_properties()
if props.copy is None:
time.sleep(5)
continue
if props.copy.status == "pending":
print(
"Blob is transferring... ",
props.copy.status,
props.copy.progress,
props.copy.status_description,
)
time.sleep(10)
continue

print(
"Transfer ended... ",
props.copy.status,
props.copy.progress,
props.copy.status_description if props.copy.status_description else "",
)
break
print(f"Create onetime export for {export_label}...")
resource_id = generate_onetime_export(
billing_client, billing_account_name, period, storage_account_resource_id
)
start_onetime_export(billing_client, resource_id)

print("Data load complete.")
print("Queued all exports.")


if __name__ == "__main__":
Expand Down
64 changes: 0 additions & 64 deletions src/azmpcli/_patch.py
Original file line number Diff line number Diff line change
@@ -1,67 +1,3 @@
#### Monkey Patching Bugs in SDK ####
def _download_initial_monkey(self, scope, metric=None, custom_headers=None, raw=False, **operation_config):
import uuid # Monkey: import for monkey
from msrest.pipeline import ClientRawResponse # Monkey: import for monkey
from azure.mgmt.consumption import models # Monkey: import for monkey

# Construct URL
url = self.download.metadata['url']
path_format_arguments = {
'scope': self._serialize.url("scope", scope, 'str', skip_quote=True)
}
url = self._client.format_url(url, **path_format_arguments)

# Construct parameters
query_parameters = {}
query_parameters['api-version'] = self._serialize.query("self.api_version", self.api_version, 'str')
if metric is not None:
query_parameters['metric'] = self._serialize.query("metric", metric, 'str')

# Construct headers
header_parameters = {}
header_parameters['Accept'] = 'application/json'
if self.config.generate_client_request_id:
header_parameters['x-ms-client-request-id'] = str(uuid.uuid1())
if custom_headers:
header_parameters.update(custom_headers)
if self.config.accept_language is not None:
header_parameters['accept-language'] = self._serialize.header("self.config.accept_language", self.config.accept_language, 'str')

# Construct and send request
request = self._client.get(url, query_parameters, header_parameters) # Monkey: change POST to GET
response = self._client.send(request, stream=False, **operation_config)
response.request.method = 'POST' # Monkey: report that we did a POST

if response.status_code not in [200, 202]:
raise models.ErrorResponseException(self._deserialize, response)

deserialized = None
header_dict = {}

if response.status_code == 200:
deserialized = self._deserialize('UsageDetailsDownloadResponse', response)
header_dict = {
'Location': 'str',
'Retry-After': 'str',
'Azure-AsyncOperation': 'str',
}

if raw:
client_raw_response = ClientRawResponse(deserialized, response)
client_raw_response.add_headers(header_dict)
return client_raw_response

return deserialized


from azure.mgmt.consumption.operations.usage_details_operations import UsageDetailsOperations
from msrestazure.polling import arm_polling

UsageDetailsOperations._download_initial = _download_initial_monkey # Monkey: apply above function
arm_polling.FINISHED = frozenset(['succeeded', 'canceled', 'failed', 'completed']) # Monkey: detect completed as valid status
arm_polling.SUCCEEDED = frozenset(['succeeded', 'completed']) # Monkey: detect completed as valid status
#### /Monkey Patching Bugs in SDK ####

#### Monkey Patching Bugs in SDK ####
def get_token(self, *scopes, **kwargs): # pylint:disable=unused-argument
from azure.core.credentials import AccessToken # Monkey: import for monkey
Expand Down

0 comments on commit d55fb2a

Please sign in to comment.