diff --git a/database_storage/database_storage.py b/database_storage/database_storage.py index e70b4cb..cbc1286 100644 --- a/database_storage/database_storage.py +++ b/database_storage/database_storage.py @@ -128,7 +128,7 @@ def _open(self, name, mode='rb'): """ assert mode == 'rb', "DatabaseStorage open mode must be 'rb'." - query = 'SELECT %(data_column)s FROM %(table)s ' + \ + query = 'SELECT %(data_column)s, %(size_column)s FROM %(table)s ' + \ 'WHERE %(name_column)s = %%s' query %= self.__dict__ cursor = connection.cursor() @@ -140,6 +140,7 @@ def _open(self, name, mode='rb'): inMemFile = StringIO.StringIO(base64.b64decode(row[0])) inMemFile.name = name inMemFile.mode = mode + inMemFile.size = row[1] return File(inMemFile) diff --git a/database_storage/test/test_database_storage.py b/database_storage/test/test_database_storage.py index 453bb55..623a072 100644 --- a/database_storage/test/test_database_storage.py +++ b/database_storage/test/test_database_storage.py @@ -68,6 +68,13 @@ def test_save_and_retrieve(self): storage._save(name, self._wrap_content(content)) self.assertEqual(storage._open(name).read(), content) + def test_open_with_size(self): + name = 'foo.bin' + content = self._get_binary_data(10 * 1024) + storage = DatabaseStorage(DEFAULT_OPTIONS) + storage._save(name, self._wrap_content(content)) + self.assertEqual(storage._open(name).size, 10 * 1024) + def test_overwrite(self): "Overwrite a file with new content, verify it updates correctly." name = 'foo.bin'