#!/www/python/bin/python

"""$URL: svn://svn/repos/trunk/dulcinea/lib/tools/zodb_index.py $
$Id: zodb_index.py 20912 2003-03-04 16:07:12Z dbinger $

Tool to build and query various low-level indices of a ZODB.  Currently
has three indices:
  * forward object references: maps each OID to the list of OIDs
    referenced by that object
  * reverse object reference: maps each OID to the list of OIDs
    that reference that object
  * class instance: maps the name of every class in the database to
    the list of instances of that class (list of OIDs)

Needs to be run once with the -c option to create the indices; all are
saved to a single Berkeley DB hash file (specified with -i, default
"index.db").

Once the index is created, you can query it.  You can also specify query
options (-f, -r, -C) on the creation run if you want.  Run with --help
to find out about the query options.

"""

from __future__ import nested_scopes, generators

import sys, os
from bsddb3 import db as bsddb
from ZODB.referencesf import referencesf
from dulcinea.script_util import Option, OptionParser, run, writenow
from dulcinea.database import pack_oid, unpack_oid, format_oid, parse_oid
from dulcinea.indented_file import IndentedFile


# -- Index open/creation -----------------------------------------------

def open_index (index_filename, create=0, readonly=1):
    """open_index(index_file : string,
                  create : boolean = false,
                  readonly : boolean = true)
       -> index_db : DB

    Open an index file in preparation for either querying it or creating
    the indices.  If 'create' is true, the file will be created and
    truncated (ie. any existing file called 'index_filename' will be
    clobbered).  This does *not* create the actual indices; use
    create_index() for that.
    """
    if create:
        readonly = 0
    if create:
        flags = bsddb.DB_CREATE | bsddb.DB_TRUNCATE
    elif readonly:
        flags = bsddb.DB_RDONLY
    else:
        flags = 0

    index_db = bsddb.DB()
    index_db.set_flags(bsddb.DB_DUP)
    if index_filename is None:
        index_filename = "index.db"
    index_db.open(index_filename, 
                  dbtype=bsddb.DB_HASH,
                  flags=flags, mode=0664)
    return index_db


def create_index (db, index_db):
    """create_index(index_db : DB)

    Creates the OID indices in index_db needed for later queries.  This
    is quite time-consuming, so it should only be done at a
    user's explicit instructions -- eg. when this script is run
    with -c.
    """

    def update_indices (oid_i, oid, object):

        # Find the objects referenced by this object (forward refs).
        (pickle, rev_id) = db.get_storage().load(oid, None)
        referenced_oids = []                   # list of raw OIDs
        referencesf(pickle, referenced_oids)

        # Update the forward and reverse object reference indices,
        # handling the case where object A references B twice.
        seen = {}
        for r_oid in referenced_oids:
            if seen.get(r_oid):
                continue
            index_db.put("F"+oid, r_oid)
            index_db.put("R"+r_oid, oid)
            seen[r_oid] = 1

        # Use the object to update the class index.
        klass = object.__class__
        klass_name = "%s.%s" % (klass.__module__, klass.__name__)
        index_db.put("C"+klass_name, oid)


    database = db.get_database()
    connection = db.get_connection()
    iterate(database, connection, callback=update_indices)


# -- Index query -------------------------------------------------------
# (useful for other modules as well as the "perform/report query"
# functions below)

def get_object_references (index_db, connection, oid, direction="R"):
    """get_object_references(index_db : DB, oid : string, direction : string)
       -> [(oid : string, object : instance)]

    If direction == "R" (reverse, the default), get the list of objects
    that reference 'oid' and return a list of (oid, object) pairs.

    If direction == "F" (forward), return the list of objects referenced
    by 'oid'.

    oid must be a packed 8-byte OID string.
    """
    if direction not in ("F", "R"):
        raise ValueError, "invalid direction %r" % direction

    # Now get the list of objects that refer to this OID.
    c = index_db.cursor()
    try:
        data = c.set(direction+oid)
    except bsddb.DBNotFoundError:
        return

    while data:
        r_oid = data[1]
        object = connection[r_oid]
        yield (r_oid, object)
        data = c.next_dup()

    c.close()


def get_class_instances (index_db, connection, klass_name):
    """get_class_instances(index_db : DB,
                           connection : Connection,
                           klass_name : string)
       -> [(oid : string, object : instance)]

    Get the list of all instances of class 'klass_name'.  Returns a list
    of (oid, object) pairs.
    """
    c = index_db.cursor()
    try:
        data = c.set("C"+klass_name)
    except bsddb.DBNotFoundError:
        return

    while data:
        oid = data[1]
        object = connection[oid]
        yield (oid, object)
        data = c.next_dup()


# -- Perform query and report results ----------------------------------
# (only used when run as a script)

def report_references (oid, index_db, connection, query_type):
    print "%s: %r" % (format_oid(oid), connection[oid])
    references = get_object_references(index_db, connection, oid,
                                       direction=query_type)
    for (r_oid, object) in references:
        object_repr = repr(object)
        l = len(object_repr)
        if l > 60:
            object_repr = object_repr[:50] + " ... " + object_repr[l-5:]
        print "  %s: %s" % (format_oid(r_oid), object_repr)


def follow_references (oid, index_db, connection, query_type,
                       outfile=None, seen=None):
    if outfile is None:
        outfile = IndentedFile(sys.stdout, indent=0, base_indent=2)
    if seen is None:
        seen = {}

    outfile.write("%s: %r" % (format_oid(oid), connection[oid]))
    if seen.get(oid):
        outfile.write(" (already seen)\n")
        return
    else:
        outfile.write("\n")
        seen[oid] = 1

    references = get_object_references(index_db, connection, oid,
                                       direction=query_type)
    if not references:                  # recursion bottomed out
        return

    for (r_oid, object) in references:
        outfile.indent_more()
        follow_references(r_oid, index_db, connection, query_type,
                          outfile=outfile, seen=seen)
        outfile.indent_less()

def reference_query (db, index_db, connection, oids, options):

    for oid in oids:
        # First convert human-readable string to packed binary OID
        oid = parse_oid(oid)

        if options.recurse:
            # And print this object, everything that references it,
            # everything that references each of these, etc.  (Or, if
            # options.reverse is false, print this object, everything
            # referenced by it, etc.)
            follow_references(oid, index_db, db.get_connection(),
                              options.query_type)
        else:
            # Just print this object and its references (forwards
            # or backwards)
            report_references(oid, index_db, db.get_connection(),
                              options.query_type)


def class_query (index_db, connection, klass_names):
    for klass_name in klass_names:
        instances = get_class_instances(index_db, connection, klass_name)
        if instances:
            print "%s:" % klass_name
            for (oid, object) in instances:
                print "  %s: %r" % (format_oid(oid), object)
        else:
            print "%s: no objects in database" % klass_name


# -- Main program ------------------------------------------------------

def main (db, default_dbspec):
    usage = "usage: %prog [options] [oid ...]"

    option_list = [
        Option("-i", "--index-file", type="string", dest="index_filename",
               metavar="FILE",
               help="specify file where the indices are stored"),
        Option("-c", "--create-index",
               action="store_true", dest="create_index",
               help="create the index file before proceeding"),
        Option("-d", "--database",
               type="string", dest="database",
               help="database specification (e.g. file:db.fs)"),
        Option("-r", "--reverse",
               action="store_const", dest="query_type", const="R",
               help="do reverse lookups, ie. given an OID x, report all "
                    "OIDs that reference x [default]"),
        Option("-f", "--forward",
               action="store_const", dest="query_type", const="F",
               help="do forward lookups, ie. report all OIDS referenced by x"),
        Option("-C", "--class",
               action="store_const", dest="query_type", const="C",
               help="report all instances of class CLASS"),
        Option("-R", "--recurse", action="store_true", dest="recurse",
               help="report OID references recursively"),
        ]
    parser = OptionParser(usage, option_list)
    parser.set_defaults(query_type="R")
    (options, args) = parser.parse_args()
    if not (args or options.create_index):
        parser.error("either -c or a list of OIDs must be supplied")

    # Open the index file (creating it if necessary)
    try:
        index_db = open_index(options.index_filename,
                              create=options.create_index,
                              readonly=not options.create_index)
    except bsddb.DBNoSuchFileError:
        parser.error("index file does not exist -- run with -c to create it")

    dbspec = options.database or default_dbspec
    db.open(dbspec)

    if options.create_index:
        create_index(db, index_db)

    if options.query_type in ("R", "F"): # object ref queries
        reference_query(db, index_db, db.get_connection(), args, options)

    elif options.query_type == "C":
        class_query(index_db, db.get_connection(), args)

    index_db.close()

