Merge branch 'main' of https://github.com/InformaticsMatters/squonk2-python-cl-tools

Author: Alan Christie

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`
 
 ---

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

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,

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)

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}")

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}")

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