"""$URL: svn+ssh://svn.mems-exchange.org/repos/trunk/dulcinea/lib/database.py $
$Id: database.py 23388 2004-02-03 00:14:30Z rmasse $

Provides the ObjectDatabase class, for initializing and accessing a ZODB
containing a set of database roots.
"""
from __future__ import generators

import sys, os, atexit, struct, time
from ZODB.DB import DB
from dulcinea.errors import ConflictError

# provide a uniform location for get_transaction
try:
    from Transaction import get_transaction # ZODB4 and newer
except ImportError:
    from __builtin__ import get_transaction # old ZODB, yuck


def pack_oid (oid_i):
    """pack_oid(oid_i : long|int) -> string

    Pack an integer or long int into an 8-byte OID string, eg.
    0x0001020304050607L -> "\0\1\2\3\4\5\6\7".
    """
    return struct.pack(">LL",
                       (oid_i & 0xffffffff00000000L) >> 32,
                       (oid_i & 0x00000000ffffffffL))

def unpack_oid (oid_s):
    """unpack_oid(oid_s : string) -> long

    Unpack an 8-byte OID string to a long int, eg.
    "\0\1\2\3\4\5\6\7" -> 0x0001020304050607L
    """
    (hi, lo) = struct.unpack(">LL", oid_s)     # two 32-bit integers
    return (hi << 32) | lo                     # a single 64-bit integer

def format_oid (oid):
    """format_oid(oid : long|int|string) -> string

    Format an OID (either a packed 8-byte string or an integer) to
    human readable form, eg.
      "\0\1\2\3\4\5\6\7" -> "0001020304050607" or
      0x0001020304050607L -> "0001020304050607"
    """
    if isinstance(oid, str):
        oid = unpack_oid(oid)
    return "%016x" % oid

def parse_oid (oid):
    """parse_oid(oid : string) -> string

    Turn a human-readable OID string into a packed 8-byte string, eg.
      "0001020304050607" -> "\0\1\2\3\4\5\6\7"
    """
    return pack_oid(long(oid, 16))


class ObjectDatabase:

    """Class used to initialize, access, and close a ZODB containing a set of
    root objects.
    
    NOTE: Care must be taken to ensure that only one instance exists for each
    ZODB database.

    Instance attributes:
      _dbspec : string
        the connection specification for this database
      _current_root_accessors : [str]
      storage : ZODB.BaseStorageBaseStorage
      database : ZODB.DB.DB
      connection : ZODB.Connection.Connection
      root : any
    """

    def __init__ (self):
        """ObjectDatabase()

        Creates a new ObjectDatabase object.
        """
        self._dbspec = None
        self.storage = None
        self.database = None
        self.connection = None
        self.root = None
        self._current_root_accessors = []
        atexit.register(self.close) # XXX is this really necessary?

    def open(self, dbspec):
        assert self.root is None, 'database already open'
        self._dbspec = dbspec
        if dbspec.find(":") != -1:
            try:
                (storage_type, args) = dbspec.split(":")
            except ValueError:
                raise ValueError, ("dbspec may have at most one "
                                   "colon: '%s' invalid" % dbspec)
            args = tuple(args.split(","))
        else:
            storage_type = dbspec
            args = ()

        if storage_type == "file":
            (filename,) = args
            (self.database, self.connection) = self._open_file(filename)

        elif storage_type == "dbm":
            (filename,) = args
            (self.database, self.connection) = self._open_dbm(filename)

        elif storage_type == "gdbm":
            (filename,) = args
            (self.database, self.connection) = self._open_gdbm(filename)

        elif storage_type == "berkeley":
            (filename,) = args
            (self.database, self.connection) = self._open_berkeley(filename)

        elif storage_type == "client":
            try:
                host, port = args
                port = int(port)
            except ValueError:
                raise ValueError("invalid dbspec: must be "
                                 "'client:<host>,<port>'" % dbspec)
            (self.database, self.connection) = self._open_client((host, port))

        elif storage_type == "empty":
            # leave database and connection as None
            self.root = {}

        else:
            raise ValueError, "invalid storage type '%s'" % storage_type

        if self.connection is not None:
            self.root = self.connection.root()

        if self.database is not None:
            # increase the maximum number of objects cached
            self.database.setCacheSize(10000)

        self.update_root_accessors()


    def close(self):
        "Close the database connection"
        self._dbspec = None
        if self.connection is not None:
            self.connection.close()
            self.connection = None
        if self.database is not None:
            self.database.close()
            self.database = None
        if self.storage is not None:
            self.storage.close()
            self.storage = None
        self.root = None

    def __repr__ (self):
        return '<%s: %s>' % (self.__class__.__name__, self._dbspec)

    # Accessors for the ZODB objects
    def get_storage (self):    return self.storage
    def get_connection (self): return self.connection
    def get_database (self):  return self.database
    def get_root (self):       return self.root

    def get_cache_size (self):
        """() -> current_size:int, nominal_size:int
        Return the number of objects currently in the ZODB cache,
        and the size of the cache.
        """
        return (self.database.cacheDetailSize()[0]['size'],
                self.database.getCacheSize())

    def pack(self):
        self.database.pack(time.time())

    def set_cache_size (self, size):
        """(int)
        Set the maximum cache size.
        """
        if self.database is not None:
            self.database.setCacheSize(size)

    def _open_file (self, filename):
        """Open 'filename' as a FileStorage, and then open a ZODB database
        around that.  Return a (database, connection) tuple containing the
        DB instance and the Connection instance that results from calling
        its 'open()' method."""

        from ZODB.FileStorage import FileStorage
        created = (not os.path.isfile(filename))

        self.storage = FileStorage(filename)
        if created:
            os.chmod(filename, 0664)            # ensure group writeable
            os.chmod(filename + ".lock", 0664)  # gag! retch!
            os.chmod(filename + ".tmp", 0664)   # ...and puuuuke!
        db = DB(self.storage)
        return (db, db.open())

    def _open_dbm (self, filename):
        """Open 'filename' as a dbmStorage, and then open a ZODB database
        around that.  Return a (database, connection) tuple containing the
        DB instance and the Connection instance that results from calling
        its 'open()' method."""

        from ZODB.dbmStorage import anydbmStorage

        self.storage = anydbmStorage(filename, flag='c')
        db = DB(self.storage)
        return (db, db.open())

    def _open_gdbm (self, filename):
        """Open 'filename' as a gdbmStorage, and then open a ZODB database
        around that.  Return a (database, connection) tuple containing the
        DB instance and the Connection instance that results from calling
        its 'open()' method."""

        from ZODB.dbmStorage import gdbmStorage

        self.storage = gdbmStorage(filename, flag='c')
        db = DB(self.storage)
        return (db, db.open())

    def _open_berkeley (self, filename):
        raise NotImplementedError, \
              "BerkeleyStorage interface not implemented yet"

    def _open_client (self, location):
        """Open 'location' (a (hostname, port_number) tuple) as a ZEO
        ClientStorage, and then open a ZODB database around that.  Return
        a (database, connection) tuple."""

        host, port = location
        if host == "":
            # If the specified hostname is the empty string, then
            # 'localhost' is used.
            location = ('localhost', port)

        from ZEO.ClientStorage import ClientStorage
        self.storage = ClientStorage(location, var='/var/tmp',
                                     wait_for_server_on_startup=0)
        db = DB(self.storage)
        return (db, db.open())

    def get_root_object (self, name):
        assert self.root is not None, "database not initialized"
        if not self.root.has_key(name):
            raise ValueError, "invalid database root '%s'" % name
        return self.root[name]

    def _add_root_accessor(self, name):
        if hasattr(self, name):
            raise ConflictError, 'attribute %s already exists' % name
        setattr(self, name, self.root[name])
        self._current_root_accessors.append(name)

    def init_root (self, name, modulename, klassname):
        __import__(modulename)
        module = sys.modules[modulename]
        klass = getattr(module, klassname)
        assert name not in self.root.keys(), name
        self.root[name] = klass()
        self._add_root_accessor(name)

    def update_root_accessors (self):
        for name in self._current_root_accessors:
            if hasattr(self, name):
                delattr(self, name)
        for name in self.root.keys():
            self._add_root_accessor(name)

    def get_root_names (self):
        return self.root.keys()

    def iterate (self, verbose=-1, status_every=0x800):
        """
        Iterate over all OIDs in a ZODB database, yielding every object
        found.

        If verbose is true or it is -1 and stdout is a tty then iterate()
        prints a status line to stdout every status_every OIDs (default
        0x800, or 2048).

        Note that this method currently only works for FileStorage
        connections.
        """
        if not hasattr(self.storage, "_index"):
            raise SystemExit, "must use FileStorage when using iterate()"

        if verbose and not sys.stdout.isatty():
            # no status reports when not a TTY
            status_every = 0

        def write_status (oid, count, newline=0):
            sys.stdout.write("\rOID: %s (objects seen: %d)" %
                             (format_oid(oid), count))
            if newline:
                sys.stdout.write("\n")
            sys.stdout.flush()

        count = 0
        for oid in self.storage._index.keys():
            count += 1
            if status_every and (count % status_every) == 0:
                write_status(oid, count)
            yield self.connection[oid]

        if status_every:
            write_status(oid, count, newline=1)


