From 1c9134d87fad052484c2e70aece440f1311c5721 Mon Sep 17 00:00:00 2001
From: Florian Stecker <oss@florianstecker.net>
Date: Sat, 11 Oct 2025 14:54:08 -0400
Subject: [PATCH] add CLI

---
 .gitignore |   1 +
 iavlread   | 117 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 .gitignore
 create mode 100755 iavlread

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..bee8a64
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+__pycache__
diff --git a/iavlread b/iavlread
new file mode 100755
index 0000000..d2e2139
--- /dev/null
+++ b/iavlread
@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+import argparse
+import plyvel
+import iavltree
+import json
+
+def decode_protobuf(subformats: dict, format_prefix: str, data: bytes):
+    result = []
+    for (k,v) in iavltree.parse_struct(data):
+        idx = f'{format_prefix}.{k}'
+        if idx in subformats:
+            f = subformats[idx]
+            if f == 'str':
+                decoded_value = v.decode('utf-8')
+            elif f == 'int':
+                decoded_value = int(v)
+            elif f == 'float':
+                decoded_value = float(v)
+            elif f == 'proto':
+                decoded_value = decode_protobuf(subformats, idx, v)
+            elif f == 'protodict':
+                decoded_value = dict(decode_protobuf(subformats, idx, v))
+            else:
+                decoded_value = v
+        else:
+            decoded_value = v
+        result.append((k, decoded_value))
+    return result
+
+def decode_output(format: str, data: bytes) -> str:
+    if format == 'str':
+        return data.decode('utf-8')
+    elif format == 'int':
+        return int(data)
+    elif format == 'float':
+        return float(data)
+    elif format.startswith('protodict'):
+        subformats = {'.' + id: subformat for x in format.split(',')[1:] for id, subformat in (x.split('='),)}
+        return dict(decode_protobuf(subformats, '', data))
+    elif format.startswith('proto'):
+        subformats = {'.' + id: subformat for x in format.split(',')[1:] for id, subformat in (x.split('='),)}
+        return decode_protobuf(subformats, '', data)
+    else:
+        return data
+
+parser = argparse.ArgumentParser(description="Read the IAVL tree in a cosmos snapshot")
+
+parser.add_argument('-d', '--database', help='Path to database (application.db folder)')
+parser.add_argument('-H', '--height', type=int, help='Block height')
+parser.add_argument('-k', '--keyformat', help='Key format for maps (e.g. Qss)')
+parser.add_argument('-v', '--valueformat', help='Value format')
+
+subparsers = parser.add_subparsers(required=True, dest='cmd')
+p_max_height = subparsers.add_parser('max_height', help = 'Get the max block height in the snapshot')
+p_get = subparsers.add_parser('get', help = 'Retrieve a single item')
+p_get.add_argument('prefix', help = 'Prefix (e.g. "s/k:emissions/")')
+p_get.add_argument('key', nargs='+', help = 'Key parts')
+p_count = subparsers.add_parser('count', help = 'Count number of items with a prefix')
+p_count.add_argument('prefix', help = 'Prefix (e.g. "s/k:emissions/")')
+p_count.add_argument('key', nargs='*', help = 'Key parts')
+p_iterate = subparsers.add_parser('iterate', help = 'Iterate over items with some prefix')
+p_iterate.add_argument('prefix', help = 'Prefix (e.g. "s/k:emissions/")')
+p_iterate.add_argument('key', nargs='*', help = 'Key parts')
+p_iterate = subparsers.add_parser('iterate_keys', help = 'Iterate over items with some prefix, output keys only')
+p_iterate.add_argument('prefix', help = 'Prefix (e.g. "s/k:emissions/")')
+p_iterate.add_argument('key', nargs='*', help = 'Key parts')
+p_iterate = subparsers.add_parser('iterate_values', help = 'Iterate over items with some prefix, output values only')
+p_iterate.add_argument('prefix', help = 'Prefix (e.g. "s/k:emissions/")')
+p_iterate.add_argument('key', nargs='*', help = 'Key parts')
+
+args = parser.parse_args()
+
+dbpath = args.database if args.database is not None else 'data/application.db'
+keyformat = args.keyformat if args.keyformat is not None else ''
+valueformat = args.valueformat if args.valueformat is not None else 'b'
+
+if args.key is None or len(args.key) == 0:
+    key = None
+else:
+    if len(args.key) > len(keyformat) + 1:
+        raise Exception('Too many key elements for keyformat')
+    key = [int(args.key[0])]
+    for f, k in zip(keyformat, args.key[1:]):
+        if f in ['i', 'I', 'q', 'Q']:
+            key.append(int(k))
+        else:
+            key.append(k)
+
+with plyvel.DB(dbpath) as db:
+    if args.height is None or args.cmd == 'max_height':
+        height = iavltree.max_height(db)
+    else:
+        height = args.height
+
+    if args.cmd == 'max_height':
+        print(height)
+    elif args.cmd == 'get':
+        result = iavltree.walk_disk(db, args.prefix, height, keyformat, key)
+
+        print(decode_output(valueformat, result))
+    elif args.cmd == 'count':
+        result = iavltree.count(db, args.prefix, height, keyformat, key = key)
+
+        print(result)
+    elif args.cmd == 'iterate' or args.cmd == 'iterate_keys' or args.cmd == 'iterate_values':
+        it = iavltree.iterate(db, args.prefix, height, keyformat, key = key)
+
+        try:
+            for k, v in it:
+                if args.cmd == 'iterate_keys':
+                    print(k)
+                elif args.cmd == 'iterate_values':
+                    print(decode_output(valueformat,v))
+                else:
+                    print((k, decode_output(valueformat, v)))
+        except BrokenPipeError:
+            pass