From 8e7749d512279582fe69133c0ba5536dae240302 Mon Sep 17 00:00:00 2001
From: Dan Helfman <dan@luminotes.com>
Date: Wed, 12 Sep 2007 21:28:28 +0000
Subject: [PATCH] controller.Database.size() function for getting the
 approximate storage size of an object.

---
 controller/Database.py           | 22 ++++++++++++++++++++++
 controller/test/Test_database.py | 26 ++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/controller/Database.py b/controller/Database.py
index 1ee9391..04de62a 100644
--- a/controller/Database.py
+++ b/controller/Database.py
@@ -193,6 +193,28 @@ class Database( object ):
 
     self.__db.sync()
 
+  def size( self, object_id, revision = None ):
+    """
+    Load the object corresponding to the given object id from the database, and return the size of
+    its pickled data in bytes. If a revision is provided, a specific revision of the object will be
+    loaded.
+
+    @type object_id: unicode
+    @param object_id: id of the object whose size should be returned
+    @type revision: int or NoneType
+    @param revision: revision of the object to load (optional)
+    """
+    if revision is not None:
+      object_id = Persistent.make_revision_id( object_id, revision )
+
+    object_id = unicode( object_id ).encode( "utf8" )
+
+    pickled = self.__db.get( object_id )
+    if pickled is None or pickled == "":
+      return None
+
+    return len( pickled )
+
   @staticmethod
   def generate_id():
     int_id = random.getrandbits( Database.ID_BITS )
diff --git a/controller/test/Test_database.py b/controller/test/Test_database.py
index 08923ea..1cb7e28 100644
--- a/controller/test/Test_database.py
+++ b/controller/test/Test_database.py
@@ -269,6 +269,32 @@ class Test_database( object ):
     self.scheduler.add( g )
     self.scheduler.wait_for( g )
 
+  def test_size( self ):
+    def gen():
+      basic_obj = Some_object( object_id = "5", value = 1 )
+      original_revision = basic_obj.revision
+
+      self.database.save( basic_obj, self.scheduler.thread )
+      yield Scheduler.SLEEP
+      if self.clear_cache: self.database.clear_cache()
+
+      size = self.database.size( basic_obj.object_id )
+
+      from cPickle import Pickler
+      from StringIO import StringIO
+      buffer = StringIO()
+      pickler = Pickler( buffer, protocol = -1 )
+      pickler.dump( basic_obj )
+      expected_size = len( buffer.getvalue() )
+
+      # as long as the size is close to the expected size, that's fine
+      assert abs( size - expected_size ) < 10
+
+    g = gen()
+    self.scheduler.add( g )
+    self.scheduler.wait_for( g )
+
+
   def test_next_id( self ):
     def gen():
       self.database.next_id( self.scheduler.thread )