L1-218: Refector and cleanup in pallet-balances-maintenance.py (#1752)

# Description

This PR cleans up unneeded functionality in the storage maintenance
script. I also moved functions across modules.

It's mostly code removal:
```
marcin@marcin-Latitude-3520 ~/g/a/s/accounts-invariants (L1-218-part1)> git ci                                                                                                                                                                                              (base) 
[L1-218-part1 3e8af98e] L1-218: Refector and cleanup in pallet-balances-maintenance.py
 5 files changed, 110 insertions(+), 772 deletions(-)
 create mode 100644 scripts/accounts-invariants/aleph_chain_version.py
 rename scripts/accounts-invariants/{common.py => chain_operations.py} (60%)
 create mode 100644 scripts/accounts-invariants/contracts.py
 create mode 100644 scripts/accounts-invariants/logger.py
```

## Type of change

Please delete options that are not relevant.

- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)

scripts/accounts-invariants/aleph_chain_version.py

@@ -0,0 +1,28 @@
+import enum
+
+
+class AlephChainVersion(enum.IntEnum):
+    VERSION_11_4 = 65,
+    VERSION_12_0 = 67,
+    VERSION_12_2 = 68,
+    VERSION_13_0 = 70,
+    VERSION_13_2 = 71,
+    VERSION_12_3 = 72,
+    VERSION_13_3 = 73,
+    VERSION_14_X = 14000000,
+
+    @classmethod
+    def from_spec_version(cls, spec_version):
+        return cls(spec_version)
+
+
+def get_aleph_chain_version(chain_connection, block_hash):
+    """
+    Retrieves spec_version from chain and returns an `AlephChainVersion` enum
+    :param chain_connection: WS handler
+    :param block_hash: Block hash to query state from
+    :return: AlephChainVersion
+    """
+    runtime_version = chain_connection.get_block_runtime_version(block_hash)
+    spec_version = runtime_version['specVersion']
+    return AlephChainVersion.from_spec_version(spec_version)

scripts/accounts-invariants/common.py → scripts/accounts-invariants/chain_operations.py

@@ -1,60 +1,12 @@
 #!/bin/python3
 
-import enum
-import logging
-import datetime
-
 import substrateinterface
 import json
+import logging
+from tqdm import tqdm
+import sys
 
-def get_global_logger():
-    log_formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
-    root_logger = logging.getLogger()
-    root_logger.setLevel('DEBUG')
-
-    time_now = datetime.datetime.now().strftime("%d-%m-%Y_%H:%M:%S")
-    file_handler = logging.FileHandler(f"pallet-balances-maintenance-{time_now}.log")
-    file_handler.setFormatter(log_formatter)
-    file_handler.setLevel(logging.DEBUG)
-    root_logger.addHandler(file_handler)
-
-    console_handler = logging.StreamHandler()
-    console_handler.setFormatter(log_formatter)
-    console_handler.setLevel(logging.INFO)
-    root_logger.addHandler(console_handler)
-
-    return logging
-
-
-log = get_global_logger()
-
-
-class ChainMajorVersion(enum.Enum):
-    PRE_12_MAJOR_VERSION = 65,
-    AT_LEAST_12_MAJOR_VERSION = 68,
-    AT_LEAST_13_2_VERSION = 71,
-
-    @classmethod
-    def from_spec_version(cls, spec_version):
-        if spec_version <= 65:
-            return cls(ChainMajorVersion.PRE_12_MAJOR_VERSION)
-        elif 68 <= spec_version < 71 or spec_version == 72:
-            return cls(ChainMajorVersion.AT_LEAST_12_MAJOR_VERSION)
-        elif spec_version >= 71:
-            return cls(ChainMajorVersion.AT_LEAST_13_2_VERSION)
-
-
-def get_chain_major_version(chain_connection, block_hash):
-    """
-    Retrieves spec_version from chain and returns an enum whether this is pre 12 version or at least 12 version
-    :param chain_connection: WS handler
-    :param block_hash: Block hash to query state from
-    :return: ChainMajorVersion
-    """
-    runtime_version = chain_connection.get_block_runtime_version(block_hash)
-    spec_version = runtime_version['specVersion']
-    major_version = ChainMajorVersion.from_spec_version(spec_version)
-    return major_version
+log = logging.getLogger()
 
 
 def filter_accounts(chain_connection,
@@ -82,15 +34,16 @@ def filter_accounts(chain_connection,
     total_accounts_count = 0
     total_issuance = 0
 
-    for (i, (account_id, info)) in enumerate(account_query):
+    for (i, (account_id, info)) in tqdm(iterable=enumerate(account_query),
+                                        desc="Accounts checked",
+                                        unit="",
+                                        file=sys.stdout):
         total_accounts_count += 1
         free = info['data']['free'].value
         reserved = info['data']['reserved'].value
         total_issuance += free + reserved
         if check_accounts_predicate(info, chain_major_version, ed):
             accounts_that_do_meet_predicate.append([account_id.value, info.serialize()])
-        if i % 5000 == 0 and i > 0:
-            log.info(f"Checked {i} accounts")
 
     log.info(
         f"Total accounts that match given predicate {check_accounts_predicate_name} is {len(accounts_that_do_meet_predicate)}")
@@ -111,44 +64,6 @@ def format_balance(chain_connection, amount):
     return f"{amount} {token}"
 
 
-def batch_transfer(chain_connection,
-                   input_args,
-                   accounts,
-                   amount,
-                   sender_keypair):
-    """
-    Send Balance.Transfer calls in a batch
-    :param chain_connection: WS connection handler
-    :param input_args: script input arguments returned from argparse
-    :param accounts: transfer beneficents
-    :param amount: amount to be transferred
-    :param sender_keypair: keypair of sender account
-    :return: None. Can raise exception in case of SubstrateRequestException thrown
-    """
-    for (i, account_ids_chunk) in enumerate(chunks(accounts, input_args.transfer_calls_in_batch)):
-        balance_calls = list(map(lambda account: chain_connection.compose_call(
-            call_module='Balances',
-            call_function='transfer',
-            call_params={
-                'dest': account,
-                'value': amount,
-            }), account_ids_chunk))
-        batch_call = chain_connection.compose_call(
-            call_module='Utility',
-            call_function='batch',
-            call_params={
-                'calls': balance_calls
-            }
-        )
-
-        extrinsic = chain_connection.create_signed_extrinsic(call=batch_call, keypair=sender_keypair)
-        log.info(f"About to send {len(balance_calls)} transfers, each with {format_balance(chain_connection, amount)} "
-                 f"from {sender_keypair.ss58_address} to below accounts: "
-                 f"{account_ids_chunk}")
-
-        submit_extrinsic(chain_connection, extrinsic, len(balance_calls), args.dry_run)
-
-
 def submit_extrinsic(chain_connection,
                      extrinsic,
                      expected_number_of_events,

scripts/accounts-invariants/contracts.py

@@ -0,0 +1,44 @@
+import logging
+from tqdm import tqdm
+import sys
+
+log = logging.getLogger()
+
+
+def query_contract_and_code_owners_accounts(chain_connection, block_hash):
+    """
+    Returns contract accounts and code owners.
+    """
+    code_owners = set()
+    contract_accounts = set()
+
+    log.info(f"Querying code owners.")
+    code_info_of_query = chain_connection.query_map(module='Contracts',
+                                                    storage_function='CodeInfoOf',
+                                                    page_size=500,
+                                                    block_hash=block_hash)
+
+    for (i, (account_id, info)) in tqdm(iterable=enumerate(code_info_of_query),
+                                        desc="CodeInfoOfs checked",
+                                        unit="",
+                                        file=sys.stdout):
+        code_owners.add(info.serialize()['owner'])
+
+    log.info(f"Total code owners is {len(code_owners)}")
+    log.debug(f"Code owners: {code_owners}")
+
+    log.info(f"Querying contract accounts.")
+    contract_info_of_query = chain_connection.query_map(module='Contracts',
+                                                        storage_function='ContractInfoOf',
+                                                        page_size=1000,
+                                                        block_hash=block_hash)
+    for (i, (account_id, info)) in tqdm(iterable=enumerate(contract_info_of_query),
+                                        desc="ContractInfoOfs checked",
+                                        unit="",
+                                        file=sys.stdout):
+        contract_accounts.add(account_id.value)
+
+    log.info(f"Total contracts count is {len(contract_accounts)}")
+    log.debug(f"Contract accounts: {contract_accounts}")
+
+    return code_owners, contract_accounts

scripts/accounts-invariants/logger.py

@@ -0,0 +1,19 @@
+import logging
+import datetime
+
+
+def setup_global_logger():
+    log_formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
+    root_logger = logging.getLogger()
+    root_logger.setLevel('DEBUG')
+
+    time_now = datetime.datetime.now().strftime("%d-%m-%Y_%H:%M:%S")
+    file_handler = logging.FileHandler(f"pallet-balances-maintenance-{time_now}.log")
+    file_handler.setFormatter(log_formatter)
+    file_handler.setLevel(logging.DEBUG)
+    root_logger.addHandler(file_handler)
+
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(log_formatter)
+    console_handler.setLevel(logging.INFO)
+    root_logger.addHandler(console_handler)