Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
  • Loading branch information
Alan Christie committed Jul 29, 2024
2 parents f6986a1 + e0efa24 commit ef1f87e
Show file tree
Hide file tree
Showing 10 changed files with 317 additions and 35 deletions.
6 changes: 4 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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].
Expand All @@ -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`

---
Expand Down
2 changes: 1 addition & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
@@ -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
91 changes: 66 additions & 25 deletions tools/coins.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand Down
119 changes: 119 additions & 0 deletions tools/create-organisations-and-units.py
Original file line number Diff line number Diff line change
@@ -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)
12 changes: 10 additions & 2 deletions tools/delete-all-instances.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
prior to a major upgrade.
"""
import argparse
import sys
from typing import Dict, List, Optional, Tuple
import urllib3

Expand All @@ -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
Expand Down Expand Up @@ -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}")
Expand Down
12 changes: 10 additions & 2 deletions tools/delete-old-instances.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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}")
Expand Down
4 changes: 4 additions & 0 deletions tools/delete-test-projects.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down
Loading

0 comments on commit ef1f87e

Please sign in to comment.