autoimport/autoimport/tmpdb.py

70 lines
1.8 KiB
Python

"""
This module provides a way to construct the temporary database
used by ``autoimport``.
The database is used to transfer the data between the ``autoimport``
modules: ``select``, ``order`` and ``write``.
``autoimport`` always uses a sqlite3 database as an interface but
the database can be stored in memory (fast) or on the disk
(for huge amounts of images).
"""
import sqlite3
import tempfile
import abc
from .db import AbstractDatabase, initialize_database
def _open_db_mem():
return (sqlite3.connect(":memory:"), None)
def _open_db_disk():
file = tempfile.NamedTemporaryFile()
db = sqlite3.connect(file.name)
return (db, file)
class AbstractTemporaryDatabase(AbstractDatabase):
pass
class MemoryTemporaryDatabase(AbstractTemporaryDatabase):
def __init__(self):
AbstractTemporaryDatabase.__init__(self)
self._db,_ = _open_db_mem()
def close(self):
self._db.close()
class DiskTemporaryDatabase(AbstractTemporaryDatabase):
def __init__(self):
AbstractTemporaryDatabase.__init__(self)
db, file = _open_db_disk()
self._db = db
self._file = file
def close(self):
self._db.close()
self._file.close()
def get_temporary_db(type_):
"""
Return an open ``TemporaryDatabase`` with already set up tables.
``type_`` is either ``"mem"`` for the in-memory implementation or
``"disk"`` for the on-disk implementation.
"""
implementations = {"mem": MemoryTemporaryDatabase,
"disk": DiskTemporaryDatabase}
if(not type_ in implementations):
raise ValueError("unsuppored implementation: {}".format(type_))
impl = implementations[type_]
instance = impl()
initialize_database(instance)
return instance