Helper code to examine corrupt zope BTree

btree.py

#!/usr/bin/env python3

# -*- coding: utf-8 -*-

import ZODB
import dm.historical
import BTrees.OOBTree


def compare_buckets(lhs: BTrees.OOBTree.BTree, rhs: BTrees.OOBTree.BTree):
    i = 0
    lhs_bucket = lhs._firstbucket
    rhs_bucket = rhs._firstbucket

    while lhs_bucket is not None:
        i += 1
        if not compare_bucket(lhs_bucket, rhs_bucket):
            print('bucket {} differs'.format(i))
            diff_lists(lhs_bucket.keys(), rhs_bucket.keys())
        lhs_bucket = lhs_bucket._next
        rhs_bucket = rhs_bucket._next

    print('There were a total of {} buckets'.format(i))


def compare_bucket(lhs_bucket, rhs_bucket):
    print("lhs_bucket: {}, rhs_bucket:{}".format(lhs_bucket, rhs_bucket))

    if lhs_bucket is None and rhs_bucket is None:
        return True
    if lhs_bucket is None:
        return False
    if rhs_bucket is None:
        return False
    lhs_keys = lhs_bucket.keys()
    rhs_keys = rhs_bucket.keys()
    return lhs_keys == rhs_keys


def check_btree(
        btree: BTrees.OOBTree.BTree,
) -> None:

    # Ensure that there is no empty bucket in the middle of the tree
    bucket = btree._firstbucket
    prev = None
    i = 0
    while bucket is not None:
        i += 1
        if bucket._next is not None and not any(bucket):
            raise AssertionError(
                'BTree bucket {} is empty and has next bucket. Status: {},'
                ' estimated size: {}, serial: {} .'
                ' prev: {}, next: {}'.format(
                    i,
                    bucket._p_status,
                    bucket._p_estimated_size,
                    bucket._p_serial,
                    prev,
                    bucket._next))
        prev = bucket
        bucket = bucket._next


def check_btree_wrapper(btree, print_e=False):
    try:
        check_btree(btree)
    except AssertionError as e:
        if print_e:
            print(str(e))
        return False
    return True


def diff_lists(lhs, rhs):
    lhs_keys = {key for key in lhs}
    rhs_keys = {key for key in rhs}

    keys_only_in_lhs = lhs_keys - rhs_keys
    keys_only_in_rhs = rhs_keys - lhs_keys
    if any(keys_only_in_lhs):
        print('{} keys only in lhs ({})'.format(
            len(keys_only_in_lhs), keys_only_in_lhs))
    if any(keys_only_in_rhs):
        print('{} keys only in rhs ({})'.format(
            len(keys_only_in_rhs), keys_only_in_rhs))


def compare_btrees(lhs: BTrees.OOBTree.BTree, rhs: BTrees.OOBTree.OOBTree):

    for key in rhs:
        print(key)
    lhs_keys = {key for key in lhs}
    rhs_keys = {key for key in rhs}

    keys_only_in_lhs = lhs_keys - rhs_keys
    keys_only_in_rhs = rhs_keys - lhs_keys
    if any(keys_only_in_lhs):
        print('{} keys only in lhs ({})'.format(
            len(keys_only_in_lhs), keys_only_in_lhs))
    if any(keys_only_in_rhs):
        print('{} keys only in rhs ({})'.format(
            len(keys_only_in_rhs), keys_only_in_rhs))


def find_last_good_btree_and_compare_buckets(btree):
    history = list(dm.historical.generateHistory(btree))
    print(
        'btree len: {}, btree history len: {}'.format(
        len(btree),
        len(history)))
    last = None
    last_good = False
    for item in history:
        if last is None:
            last = item
            continue
        good = check_btree_wrapper(item['obj'])
        if not last_good and good:
            check_btree_wrapper(last['obj'], print_e=True)
            print(item)
            compare_buckets(item['obj'], last['obj'])
            break
        last = item
        last_good = good


def find_bucket_oid(btree, oid):
    i = 0
    bucket = btree._firstbucket

    while bucket is not None:
        i += 1
        if bucket._p_oid == oid:
            return bucket
        bucket = bucket._next

    return None


def watch_oid(btree):
    history = list(dm.historical.generateHistory(btree))
    for item in history:
        bucket = find_bucket_oid(item['obj'], b'\x00\x00\x00\x00\x00\x00\r\x91')
        if bucket is not None:
            print('Time: {}, len: {}, bucket: {}'.format(item['time'], len(bucket), bucket))


if __name__ == '__main__':
    path = '/home/matthchr/Desktop/agent/agent.db'
    zodb = ZODB.DB(path)
    conn = zodb.open()

    print(conn)

    btree = conn.root()['agent.task.history'].history

    # Uncomment this and comment out watch_oid if you want to find the first transaction
    # where the Btree had gone "bad" - this also shows the delta between the last good
    # and the first bad tree
    # find_last_good_btree_and_compare_buckets(btree)

    watch_oid(btree)