~loggerhead-team/loggerhead/trunk-rich

#

# Copyright (C) 2006  Robey Pointer <robey@lag.net>

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 51 Franklin Street, Suite 500, Boston, MA 02110-1335  USA

#

"""

a cache for chewed-up 'file change' data structures, which are basically just

a different way of storing a revision delta.  the cache improves lookup times

10x over bazaar's xml revision structure, though, so currently still worth

doing.

once a revision is committed in bazaar, it never changes, so once we have

cached a change, it's good forever.

"""

import marshal

import os

import pickle

import tempfile

import zlib

from sqlite3 import dbapi2

# We take an optimistic approach to concurrency here: we might do work twice

# in the case of races, but not crash or corrupt data.

def safe_init_db(filename, init_sql):

    # To avoid races around creating the database, we create the db in

    # a temporary file and rename it into the ultimate location.

    fd, temp_path = tempfile.mkstemp(dir=os.path.dirname(filename))

    os.close(fd)

    con = dbapi2.connect(temp_path)

    cur = con.cursor()

    cur.execute(init_sql)

    con.commit()

    con.close()

    os.rename(temp_path, filename)

class FakeShelf(object):

    def __init__(self, filename):

        create_table = not os.path.exists(filename)

        if create_table:

            safe_init_db(

                filename, "create table RevisionData "

                "(revid binary primary key, data binary)")

        self.connection = dbapi2.connect(filename)

        self.cursor = self.connection.cursor()

    def _create_table(self, filename):

        con = dbapi2.connect(filename)

        cur = con.cursor()

        cur.execute(

            "create table RevisionData "

            "(revid binary primary key, data binary)")

        con.commit()

        con.close()

    def _serialize(self, obj):

        return dbapi2.Binary(pickle.dumps(obj, protocol=2))

    def _unserialize(self, data):

        return pickle.loads(str(data))

    def get(self, revid):

        self.cursor.execute(

            "select data from revisiondata where revid = ?", (revid, ))

        filechange = self.cursor.fetchone()

        if filechange is None:

            return None

        else:

            return self._unserialize(filechange[0])

    def add(self, revid, object):

        try:

            self.cursor.execute(

                "insert into revisiondata (revid, data) values (?, ?)",

                (revid, self._serialize(object)))

            self.connection.commit()

        except dbapi2.IntegrityError:

            # If another thread or process attempted to set the same key, we

            # assume it set it to the same value and carry on with our day.

            pass

class RevInfoDiskCache(object):

    """Like `RevInfoMemoryCache` but backed in a sqlite DB."""

    def __init__(self, cache_path):

        if not os.path.exists(cache_path):

            os.mkdir(cache_path)

        filename = os.path.join(cache_path, 'revinfo.sql')

        create_table = not os.path.exists(filename)

        if create_table:

            safe_init_db(

                filename, "create table Data "

                "(key binary primary key, revid binary, data binary)")

        self.connection = dbapi2.connect(filename)

        self.cursor = self.connection.cursor()

    def get(self, key, revid):

        if not isinstance(key, bytes):

            raise TypeError(key)

        if not isinstance(revid, bytes):

            raise TypeError(revid)

        self.cursor.execute(

            "select revid, data from data where key = ?", (dbapi2.Binary(key),))

        row = self.cursor.fetchone()

        if row is None:

            return None

        elif str(row[0]) != revid:

            return None

        else:

            try:

                return marshal.loads(zlib.decompress(row[1]))

            except (EOFError, ValueError, TypeError):

                return None

    def set(self, key, revid, data):

        if not isinstance(key, bytes):

            raise TypeError(key)

        if not isinstance(revid, bytes):

            raise TypeError(revid)

        try:

            self.cursor.execute(

                'delete from data where key = ?', (dbapi2.Binary(key), ))

            blob = zlib.compress(marshal.dumps(data, 2))

            self.cursor.execute(

                "insert into data (key, revid, data) values (?, ?, ?)",

                list(map(dbapi2.Binary, [key, revid, blob])))

            self.connection.commit()

        except dbapi2.IntegrityError:

            # If another thread or process attempted to set the same key, we

            # don't care too much -- it's only a cache after all!

            pass