diff --git a/README.md b/README.md index c0fc789..89161f1 100644 --- a/README.md +++ b/README.md @@ -13,10 +13,10 @@ set of examples so that users can create their own utilities. > Most tools need to be executed by a user with admin privileges. ## Usage -The tools utilise the Pyhon client's `Environment` module, which expects +The tools utilise the Python client's `Environment` module, which expects you to create an `Envrionments` file - a YAML file that defines the variables used to connect to the corresponding installation. The environments -file (typically `~/.squonk2/environmemnts`) allows you to creat variables +file (typically `~/.squonk2/environmemnts`) allows you to create variables for multiple installations identified by name. See the **Environment module** section of the [Squonk2 Python Client]. @@ -43,11 +43,13 @@ display the tool's help. You should find the following tools in this repository: - - `coins` +- `create-organisations-and-units` - `delete-all-instances` - `delete-old-instances` - `delete-test-projects` - `list-environments` - `load-er` +- `load-job-manifests` - `save-er` --- diff --git a/requirements.txt b/requirements.txt index 31eddff..77fba0e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ im-squonk2-client >= 3.0.2, < 4.0.0 -python-dateutil == 2.8.2 +python-dateutil == 2.9.0 rich == 12.6.0 pyyaml == 6.0.1 diff --git a/tools/coins.py b/tools/coins.py index bb3c262..effafb7 100755 --- a/tools/coins.py +++ b/tools/coins.py @@ -3,9 +3,11 @@ """ import argparse from collections import namedtuple +import decimal from decimal import Decimal import sys -from typing import Any, Dict +from typing import Any, Dict, Optional +from attr import dataclass import urllib3 from rich.pretty import pprint @@ -16,8 +18,12 @@ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) -AdjustedCoins: namedtuple = namedtuple("AdjustedCoins", - ["coins", "fc", "ac", "aac"]) +@dataclass +class AdjustedCoins: + coins: Decimal + fc: Decimal + ac: Decimal + aac: Decimal def main(c_args: argparse.Namespace) -> None: @@ -36,6 +42,9 @@ def main(c_args: argparse.Namespace) -> None: username=env.admin_user, password=env.admin_password, ) + if not token: + console.log("[bold red]ERROR[/bold red] Failed to get token") + sys.exit(1) # Get the product details. # This gives us the product's allowance, limit and overspend multipliers @@ -55,6 +64,11 @@ def main(c_args: argparse.Namespace) -> None: remaining_days: int = p_rv.msg["product"]["coins"]["remaining_days"] + # What's the 'billing prediction' in the /product response? + # We'll compare this later to ensure it matches what we find + # when we calculate the cost to the user using the product charges. + product_response_billing_prediction: Decimal = round(Decimal(p_rv.msg["product"]["coins"]["billing_prediction"]), 2) + # Get the product's charges... pc_rv: AsApiRv = AsApi.get_product_charges(token, product_id=args.product, pbp=args.pbp) if not pc_rv.success: @@ -65,16 +79,16 @@ def main(c_args: argparse.Namespace) -> None: pprint(pc_rv.msg) # Accumulate all the storage costs - # (excluding the current which will be interpreted as the "burn rate") + # (the current record wil be used to set the future the "burn rate") num_storage_charges: int = 0 burn_rate: Decimal = Decimal() total_storage_coins: Decimal = Decimal() if "items" in pc_rv.msg["storage_charges"]: for item in pc_rv.msg["storage_charges"]["items"]: + total_storage_coins += Decimal(item["coins"]) if "current_bytes" in item["additional_data"]: - burn_rate = Decimal(item["coins"]) + burn_rate = Decimal(item["burn_rate"]) else: - total_storage_coins += Decimal(item["coins"]) num_storage_charges += 1 # Accumulate all the processing costs @@ -128,31 +142,58 @@ def main(c_args: argparse.Namespace) -> None: "Coins (Adjusted)": f"{ac.fc} + {ac.aac} = {ac.coins}", } - additional_coins: Decimal = total_uncommitted_processing_coins + burn_rate * remaining_days + # We've accumulated today's storage costs (based on the current 'peak'), + # so we can only predict further storage costs if there's more than + # 1 day left until the billing day. And that 'burn rate' is based on today's + # 'current' storage, not its 'peak'. + burn_rate_contribution: Decimal = Decimal() + burn_rate_days: int = max(remaining_days - 1, 0) + if burn_rate_days > 0: + burn_rate_contribution = burn_rate * burn_rate_days + additional_coins: Decimal = total_uncommitted_processing_coins + burn_rate_contribution predicted_total_coins: Decimal = total_coins zero: Decimal = Decimal() - if remaining_days > 0: - if burn_rate > zero: - - predicted_total_coins += additional_coins - p_ac: AdjustedCoins = _calculate_adjusted_coins( - predicted_total_coins, - allowance, - allowance_multiplier) - - invoice["Prediction"] = { - "Coins (Burn Rate)": str(burn_rate), - "Coins (Additional Spend)": f"{total_uncommitted_processing_coins} + {remaining_days} x {burn_rate} = {additional_coins}", - "Coins (Total Raw)": f"{total_coins} + {additional_coins} = {predicted_total_coins}", - "Coins (Penalty Free)": str(p_ac.fc), - "Coins (In Allowance Band)": str(p_ac.ac), - "Coins (Allowance Charge)": f"{p_ac.ac} x {allowance_multiplier} = {p_ac.aac}", - "Coins (Adjusted)": f"{p_ac.fc} + {p_ac.aac} = {p_ac.coins}", - } + calculated_billing_prediction: Decimal = Decimal() + + if remaining_days > 0 and burn_rate > zero: + + predicted_total_coins += additional_coins + p_ac: AdjustedCoins = _calculate_adjusted_coins( + predicted_total_coins, + allowance, + allowance_multiplier) + + invoice["Prediction"] = { + "Coins (Burn Rate)": str(burn_rate), + "Coins (Expected Burn Rate Contribution)": f"{burn_rate_days} x {burn_rate} = {burn_rate_contribution}", + "Coins (Additional Spend)": f"{total_uncommitted_processing_coins} + {burn_rate_contribution} = {additional_coins}", + "Coins (Total Raw)": f"{total_coins} + {additional_coins} = {predicted_total_coins}", + "Coins (Penalty Free)": str(p_ac.fc), + "Coins (In Allowance Band)": str(p_ac.ac), + "Coins (Allowance Charge)": f"{p_ac.ac} x {allowance_multiplier} = {p_ac.aac}", + "Coins (Adjusted)": f"{p_ac.fc} + {p_ac.aac} = {p_ac.coins}", + } + + calculated_billing_prediction = p_ac.coins # Now just pre-tty-print the invoice pprint(invoice) + console.log(f"Calculated billing prediction is {calculated_billing_prediction}") + console.log(f"Product response billing prediction is {product_response_billing_prediction}") + + if calculated_billing_prediction == product_response_billing_prediction: + console.log(":white_check_mark: CORRECT - Predictions match") + else: + discrepancy: Decimal = abs(calculated_billing_prediction - product_response_billing_prediction) + if calculated_billing_prediction > product_response_billing_prediction: + who_is_higher: str = "Calculated" + else: + who_is_higher: str = "Product response" + console.log(":cross_mark: ERROR - Predictions do not match.") + console.log(f"There's a discrepancy of {discrepancy} and the {who_is_higher} value is higher.") + sys.exit(1) + def _calculate_adjusted_coins(total_coins: Decimal, allowance: Decimal, diff --git a/tools/create-organisations-and-units.py b/tools/create-organisations-and-units.py new file mode 100755 index 0000000..e5539fc --- /dev/null +++ b/tools/create-organisations-and-units.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python +"""Creates organisations using a YAML file to define their names and owners. +The file is simply a list of organisations that have a `name`, and `owner` +with an optional list of `units` with names and billing days (with a default of '3) +(which are created in the same way). +""" +import argparse +from pathlib import Path +import sys +from typing import Any, Dict, List +import urllib3 + +from rich.console import Console +from squonk2.auth import Auth +from squonk2.as_api import AsApi, AsApiRv +from squonk2.environment import Environment +import yaml + +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + + +def main(c_args: argparse.Namespace, filename: str) -> None: + """Main function.""" + + console = Console() + + _ = Environment.load() + env: Environment = Environment(c_args.environment) + AsApi.set_api_url(env.as_api) + + token: str = Auth.get_access_token( + keycloak_url=env.keycloak_url, + keycloak_realm=env.keycloak_realm, + keycloak_client_id=env.keycloak_as_client_id, + username=env.admin_user, + password=env. + admin_password, + ) + if not token: + print("Failed to get token") + sys.exit(1) + + # Get the current organisations (as an admin user you should see them all) + org_rv: AsApiRv = AsApi.get_organisations(token) + if not org_rv.success: + console.log(':boom: Failed to get existing organisations') + sys.exit(1) + existing_org_names: List[str] = [] + existing_orgs: Dict[str, str] = {} + for org in org_rv.msg['organisations']: + if org['name'] not in ['Default']: + existing_orgs[org['name']] = org['id'] + existing_org_names.append(org['name']) + + # Just read the list from the chosen file + file_content: str = Path(filename).read_text(encoding='utf8') + orgs: List[Dict[str, Any]] = yaml.load(file_content, Loader=yaml.FullLoader) + # Create the organisations one at a time (to handle any errors gracefully) + for org in orgs: + org_name: str = org.get('name') + if not org_name: + console.log(':boom: File has an organisation without a name') + sys.exit(1) + owner: str = org.get('owner') + if not owner: + console.log(':boom: File has an organisation without an owner') + sys.exit(1) + # Now try and create the organisation (if it's new)... + if org_name in existing_org_names: + console.log(f':white_check_mark: Skipping organisation "{org_name}" - it already exists') + else: + org_rv: AsApiRv = AsApi.create_organisation(token, org_name=org_name, org_owner=owner) + if org_rv.success: + emoji = ':white_check_mark:' + existing_orgs[org_name] = org_rv.msg['id'] + else: + emoji = ':cross_mark:' + # Log + console.log(f'{emoji} {org_name} ({owner})') + # Units? + if 'units' in org: + org_rv: AsApiRv = AsApi.get_units(token, org_id=existing_orgs[org_name]) + existing_unit_names: List[str] = [unit['name'] for unit in org_rv.msg['units']] + for unit in org['units']: + unit_name: str = unit.get('name') + if not unit_name: + console.log(':boom: File has a unit without a name') + sys.exit(1) + if unit_name in existing_unit_names: + console.log(f':white_check_mark: Skipping unit "{org_name}/{unit_name}" - it already exists') + else: + billing_day: int = unit.get('billing_day', 3) + # Now try and create the unit (if it's new)... + unit_rv: AsApiRv = AsApi.create_unit(token, org_id=existing_orgs[org_name], unit_name=unit_name, billing_day=billing_day) + emoji = ':white_check_mark:' if unit_rv.success else ':cross_mark:' + # Log + console.log(f' {emoji} {unit_name} (billing day {billing_day})') + + +if __name__ == "__main__": + + # Parse command line arguments + parser = argparse.ArgumentParser( + prog="create-organisations", + description="Creates Organisations and Units (from a YAML file). You will need admin privileges to use this tool." + ) + parser.add_argument('environment', type=str, help='The environment name') + parser.add_argument('file', type=str, help='The source file') + args: argparse.Namespace = parser.parse_args() + + filename: str = args.file + if not filename.endswith('.yaml'): + filename += '.yaml' + + # File must exist + if not Path(filename).is_file(): + parser.error(f"File '{filename}' does not exist") + + main(args, filename) diff --git a/tools/delete-all-instances.py b/tools/delete-all-instances.py index fa051fa..f2d0cbc 100755 --- a/tools/delete-all-instances.py +++ b/tools/delete-all-instances.py @@ -11,6 +11,7 @@ prior to a major upgrade. """ import argparse +import sys from typing import Dict, List, Optional, Tuple import urllib3 @@ -34,13 +35,18 @@ def main(c_args: argparse.Namespace) -> None: username=env.admin_user, password=env.admin_password, ) + if not token: + print("Failed to get token") + sys.exit(1) # The collection of instances project_instances: Dict[str, List[Tuple[str, str]]] = {} # To see everything we need to become admin... rv: DmApiRv = DmApi.set_admin_state(token, admin=True) - assert rv.success + if not rv.success: + print("Failed to set admin state") + sys.exit(1) # Iterate through projects to get instances... num_instances: int = 0 @@ -82,7 +88,9 @@ def main(c_args: argparse.Namespace) -> None: # Revert to a non-admin state # To see everything we need to become admin... rv = DmApi.set_admin_state(token, admin=False) - assert rv.success + if not rv.success: + print("Failed to unset admin state") + sys.exit(1) print(f"Found {num_instances}") print(f"Deleted {num_deleted}") diff --git a/tools/delete-old-instances.py b/tools/delete-old-instances.py index d437e93..37b95de 100755 --- a/tools/delete-old-instances.py +++ b/tools/delete-old-instances.py @@ -8,6 +8,7 @@ is that the user has admin rights. """ import argparse +import sys from datetime import datetime, timedelta from typing import List, Optional, Tuple import urllib3 @@ -33,10 +34,15 @@ def main(c_args: argparse.Namespace) -> None: username=env.admin_user, password=env.admin_password, ) + if not token: + print("Failed to get token") + sys.exit(1) # To see everything we need to become admin... rv: DmApiRv = DmApi.set_admin_state(token, admin=True) - assert rv.success + if not rv.success: + print("Failed to set admin state") + sys.exit(1) # Max age? max_stopped_age: timedelta = timedelta(hours=args.age) @@ -76,7 +82,9 @@ def main(c_args: argparse.Namespace) -> None: # Revert to a non-admin state # To see everything we need to become admin... rv = DmApi.set_admin_state(token, admin=False) - assert rv.success + if not rv.success: + print("Failed to unset admin state") + sys.exit(1) print(f"Found {len(old_instances)}") print(f"Deleted {num_deleted}") diff --git a/tools/delete-test-projects.py b/tools/delete-test-projects.py index fccd1f9..826144a 100755 --- a/tools/delete-test-projects.py +++ b/tools/delete-test-projects.py @@ -4,6 +4,7 @@ """Deletes projects in the DM created by a built-in test user. """ import argparse +import sys from typing import Any, Dict, List, Optional import urllib3 @@ -29,6 +30,9 @@ def main(c_args: argparse.Namespace) -> None: username=env.admin_user, password=env.admin_password, ) + if not token: + print("Failed to get token") + sys.exit(1) ret_val: DmApiRv = DmApi.get_available_projects(token) assert ret_val.success diff --git a/tools/load-er.py b/tools/load-er.py index 1b55d29..f6f4c1b 100755 --- a/tools/load-er.py +++ b/tools/load-er.py @@ -16,7 +16,7 @@ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) -def main(c_args: argparse.Namespace, filename: StopIteration) -> None: +def main(c_args: argparse.Namespace, filename: str) -> None: """Main function.""" console = Console() @@ -32,6 +32,9 @@ def main(c_args: argparse.Namespace, filename: StopIteration) -> None: username=env.admin_user, password=env.admin_password, ) + if not token: + print("Failed to get token") + sys.exit(1) # Just read the list from the chosen file file_content: str = Path(filename).read_text(encoding='utf8') @@ -89,8 +92,9 @@ def main(c_args: argparse.Namespace, filename: StopIteration) -> None: if not num_rates and not num_rates_failed: console.log('Loaded [bold red1]nothing[/bold red1]') - # Error states - if num_rates_failed or not num_rates and not num_rates_failed: + # Error if nothing was loaded. + # Errors or jobs without rates are not considered an error + if num_rates == 0: sys.exit(1) diff --git a/tools/load-job-manifests.py b/tools/load-job-manifests.py new file mode 100755 index 0000000..44feae7 --- /dev/null +++ b/tools/load-job-manifests.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python +"""Loads Job Manifests using a YAML file to define their origin. +The file is simply a list of manifests that have a `url`, optional `header`, +and `params` (both of which are expected to be JSON strings of keys and values). +""" +import argparse +from pathlib import Path +import sys +from typing import Any, Dict, List +import urllib3 + +from rich.console import Console +from squonk2.auth import Auth +from squonk2.dm_api import DmApi, DmApiRv +from squonk2.environment import Environment +import yaml + +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + + +def main(c_args: argparse.Namespace, filename: str) -> None: + """Main function.""" + + console = Console() + + _ = Environment.load() + env: Environment = Environment(c_args.environment) + DmApi.set_api_url(env.dm_api) + + token: str = Auth.get_access_token( + keycloak_url=env.keycloak_url, + keycloak_realm=env.keycloak_realm, + keycloak_client_id=env.keycloak_dm_client_id, + username=env.admin_user, + password=env.admin_password, + ) + if not token: + print("Failed to get token") + sys.exit(1) + + # Just read the list from the chosen file + file_content: str = Path(filename).read_text(encoding='utf8') + manifests: List[Dict[str, Any]] = yaml.load(file_content, Loader=yaml.FullLoader) + # Load the manifests one at a time (to handle any errors gracefully) + num_manifests: int = 0 + num_manifests_failed: int = 0 + for manifest in manifests: + # A manifest must have a url and optional header and params + url: str = manifest.get('url') + if not url: + console.log(':boom: File has a manifest without a URL') + sys.exit(1) + header: str = manifest.get('header') + params: str = manifest.get('params') + # Now try and set the rate... + jm_rv: DmApiRv = DmApi.put_job_manifest(token, url=url, header=header, params=params) + if jm_rv.success: + num_manifests += 1 + emoji = ':white_check_mark:' + else: + num_manifests_failed += 1 + emoji = ':cross_mark:' + # Log + console.log(f'{emoji} {url}') + + # Summary + if num_manifests: + console.log(f'Job manifests loaded {num_manifests}') + # Error states + if num_manifests_failed: + console.log(f'Job manifest failures {num_manifests_failed}') + + +if __name__ == "__main__": + + # Parse command line arguments + parser = argparse.ArgumentParser( + prog="load-job-manifests", + description="Loads Job Manifests (from a YAML file)" + ) + parser.add_argument('environment', type=str, help='The environment name') + parser.add_argument('file', type=str, help='The source file') + args: argparse.Namespace = parser.parse_args() + + filename: str = args.file + if not filename.endswith('.yaml'): + filename += '.yaml' + + # File must exist + if not Path(filename).is_file(): + parser.error(f"File '{filename}' does not exist") + + main(args, filename) diff --git a/tools/save-er.py b/tools/save-er.py index 31db474..320bd47 100755 --- a/tools/save-er.py +++ b/tools/save-er.py @@ -32,6 +32,9 @@ def main(c_args: argparse.Namespace) -> None: username=env.admin_user, password=env.admin_password, ) + if not token: + print("Failed to get token") + sys.exit(1) er_rv: DmApiRv = DmApi.get_job_exchange_rates(token) if not er_rv.success: