path: root/trunk/dimbola/db.py

# Copyright (C) 2009  Lars Wirzenius <liw@liw.fi>
# 
# 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 3 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, see <http://www.gnu.org/licenses/>.


import logging
import os
import sqlite3

import dimbola


class Database(object):

    '''Interface to the database.
    
    This provides a thin layer of abstraction to the Sqlite3 database
    file we use to store data. It makes testing those parts that access
    the database easier: all actual SQL queries are in this class, so
    it is easy to mock up when testing upper level.
    
    '''
    
    SUPPORTED = 2
    
    def __init__(self, filename):
        self.conn = sqlite3.connect(filename)
        self.transaction = None
        
    def init_db(self):
        '''Create tables in the database.'''

        logging.debug('Creating database tables.')
        
        self.begin_transaction()
        
        self.execute('''create table if not exists dbmeta (version integer)''')
        version = self.query_first('select version from dbmeta')
        if version is None:
            version = 1
        if version > self.SUPPORTED:
            raise Exception('Database is version %s, we support up to %s' %
                            (version, self.SUPPORTED))
        if version < self.SUPPORTED:
            self.execute('delete from dbmeta')
            self.execute('insert into dbmeta (version) values (?)',
                         (self.SUPPORTED,))
        
        self.execute('''create table if not exists photos
                        (photoid integer primary key autoincrement,
                         folderid integer,
                         basename string,
                         rating integer,
                         rotate integer)''')
        if version < 2:
            self.execute('alter table photos add column sha1 string')
        
        self.execute('''create table if not exists folders
                        (folderid integer primary key autoincrement,
                         foldername string unique)''')
                      
        self.execute('''create table if not exists thumbnails
                        (photoid integer unique,
                         thumbnail blob)''')
                      
        self.execute('''create table if not exists previews
                        (photoid integer unique,
                         preview blob)''')
        
        self.execute('''create table if not exists tagnames
                        (tagid integer primary key,
                         tagname string,
                         tagparent integer)''')
                      
        self.execute('''create table if not exists tags
                        (tagid integer,
                         photoid integer)''')
                      
        self.execute('''create table if not exists exif
                        (exifid integer,
                         photoid integer,
                         exifvalue string)''')
                      
        self.execute('''create table if not exists exifnames
                        (exifid integer primary key,
                         exifname string unique)''')
                      
        self.end_transaction()

    def begin_transaction(self):
        '''Start a transaction.
        
        Any multiple-query thing should happen within a transaction.
        
        '''
        
        assert self.transaction is None
        self.transaction = self.conn.cursor()
        
    def end_transaction(self):
        '''End the current transaction.'''
        
        assert self.transaction is not None
        self.transaction = None
        self.conn.commit()
        
    def __enter__(self):
        self.begin_transaction()
        
    def __exit__(self, exc_type, exc_value, exc_traceback):
        if exc_type is None:
            self.end_transaction()
        else:
            self.conn.rollback()
            self.transaction = None

    def execute(self, sql, args=None):
        '''Execute some SQL within the current transaction.'''
        logging.debug('Executing SQL: "%s" with arguments %s' % (sql, args))
        if args is None:
            self.transaction.execute(sql)
        else:
            self.transaction.execute(sql, args)

    def query(self, sql, args=None):
        '''Query the database within the current transaction.
        
        Generate result rows, one by one.
        
        '''
        
        self.execute(sql, args)
        for row in self.transaction:
            yield row

    def query_first(self, sql, args=None):
        '''Like query, but return first column of first match, or None.'''
        for row in self.query(sql, args):
            return row[0]
        return None

    def find_folder(self, foldername):
        '''Return folderid corresponding to a foldername, or None.'''
        sql = 'select folderid from folders where foldername = ?'
        return self.query_first(sql, (foldername,))

    def get_folder_name(self, folderid):
        '''Return foldername that corresponds to folder id.'''
        sql = 'select foldername from folders where folderid = ?'
        return self.query_first(sql, (folderid,))

    def add_folder(self, foldername):
        '''Add a new folder to the database, return its id.'''
        self.execute('insert into folders (foldername) values (?)',
                     (foldername,))
        folderid = self.transaction.lastrowid
        logging.debug('Added folder %s as %s' % (foldername, folderid))
        return folderid

    def find_photoids(self):
        '''Return list of ids of all photos in database.'''
        return [row[0] for row in self.query('select photoid from photos')]

    def find_photoids_in_folder(self, folderid):
        '''Return ids of photos in a folder.'''
        sql = 'select photoid from photos where folderid = ?'
        return [row[0] for row in self.query(sql, (folderid,))]

    def add_photo(self, folderid, basename, rating, rotate):
        '''Add a new photo to the database, return its id.'''
        self.execute('''insert into photos 
                        (folderid, basename, rating, rotate)
                        values (:folderid, :basename, :rating, :rotate)''',
                     { 'folderid': folderid,
                       'basename': basename,
                       'rating': rating,
                       'rotate': rotate })
        photoid = self.transaction.lastrowid
        logging.debug('Added photo %s in folder %s as %s' % 
                      (basename, folderid, photoid))
        return photoid
        
    def remove_photo(self, photoid):
        '''Remove photo from database.
        
        This also removes all tags and other metadata related to the
        photo.
        
        '''
        
        self.execute('delete from photos where photoid = ?', (photoid,))
        self.execute('delete from thumbnails where photoid = ?', (photoid,))
        self.execute('delete from previews where photoid = ?', (photoid,))
        self.execute('delete from tags where photoid = ?', (photoid,))
        self.execute('delete from exif where photoid = ?', (photoid,))
        
    def set_sha1(self, photoid, sha1):
        '''Set the sha1 checksum for a photo.'''
        self.execute('update photos set sha1 = ? where photoid = ?',
                     (sha1, photoid))

    def get_sha1(self, photoid):
        '''Return the sha1 checksum for a photo.'''
        return self.query_first('select sha1 from photos where photoid = ?',
                                (photoid,))
    
    def find_photos_without_checksum(self):
        '''Return list of photos without checksums.'''
        sql = 'select photoid from photos where sha1 isnull'
        return [row[0] for row in self.query(sql)]
    
    def get_basic_photo_metadata(self, photoid):
        '''Return the basic metadata about a photo.
        
        Return folderid, basename, rating, rotate.
        
        '''
        
        sql = ('select folderid, basename, rating, rotate from photos '
               'where photoid = ?')
        for row in self.query(sql, (photoid,)):
            return row
        return None, None, None, None

    def get_photo_pathname(self, photoid):
        '''Return the full pathname of a photo.'''
        folderid, basename, c, d = self.get_basic_photo_metadata(photoid)
        foldername = self.get_folder_name(folderid)
        return os.path.join(foldername, basename)

    def find_exifname(self, exifname):
        '''Return exifid corresponding to exifname, or None.'''
        sql = 'select exifid from exifnames where exifname = ?'
        return self.query_first(sql, (exifname,))

    def add_exifname(self, exifname):
        '''Add a new exifname to the database, return its id.'''
        self.execute('insert into exifnames (exifname) values (?)',
                     (exifname,))
        exifid = self.transaction.lastrowid
        logging.debug('Added exifname %s as %s' % (exifname, exifid))
        return exifid

    def add_exif(self, photoid, exifid, exifvalue):
        '''Add a new exif header for a photo.'''
        self.execute('''insert into exif (photoid, exifid, exifvalue)
                        values (:photoid, :exifid, :exifvalue)''',
                     { 'photoid': photoid,
                       'exifid': exifid,
                       'exifvalue': exifvalue })

    def get_exif(self, photoid, exifname):
        '''Return a given exif header for a given photo.'''
        exifid = self.find_exifname(exifname)
        sql = 'select exifvalue from exif where photoid = ? and exifid = ?'
        value = self.query_first(sql, (photoid, exifid))
        if value is not None:
            value = str(value)
        return value

    def add_thumbnail(self, photoid, thumbnail):
        '''Add a new thumbnail for a photo.'''
        self.execute('''insert into thumbnails (photoid, thumbnail)
                        values (:photoid, :thumbnail)''',
                     { 'photoid': photoid,
                       'thumbnail': buffer(thumbnail) })

    def get_thumbnail(self, photoid):
        '''Return the thumbnail for a given photo, or None.'''
        sql = 'select thumbnail from thumbnails where photoid = ?'
        return self.query_first(sql, (photoid,))

    def add_preview(self, photoid, preview):
        '''Add a new preview for a photo.'''
        self.execute('''insert into previews (photoid, preview)
                        values (:photoid, :preview)''',
                     { 'photoid': photoid,
                       'preview': buffer(preview) })

    def get_preview(self, photoid):
        '''Return the preview for a given photo, or None.'''
        sql = 'select preview from previews where photoid = ?'
        return self.query_first(sql, (photoid,))
        
    def get_tagnames(self):
        '''Return generator to go through all tagid, name, parentid tuples.'''
        sql = 'select tagid, tagname, tagparent from tagnames'
        for tagid, tagname, tagparent in self.query(sql):
            yield tagid, tagname, tagparent

    def get_tagname(self, tagid):
        '''Return name corresponding to a tagid.'''
        sql = 'select tagname from tagnames where tagid = ?'
        return self.query_first(sql, (tagid,))

    def add_tagname(self, tagname):
        '''Add a new tagname to the database, return its id.'''
        self.execute('insert into tagnames (tagname) values (?)',
                     (tagname,))
        tagid = self.transaction.lastrowid
        logging.debug('Added tagname %s as %s' % (tagname, tagid))
        return tagid

    def remove_tagname(self, tagid):
        '''Remove a tag name from the database.
        
        This also removes the tag from the photos that have it.
        
        '''
        self.execute('delete from tags where tagid = ?', (tagid,))
        self.execute('delete from tagnames where tagid = ?', (tagid,))

    def change_tagname(self, tagid, tagname):
        '''Change the name of a tag.'''
        self.execute('update tagnames set tagname = ? where tagid = ?', 
                     (tagname, tagid))
    
    def set_tagparent(self, tagid, parentid):
        '''Set parent of a tag.'''
        self.execute('update tagnames set tagparent = ? where tagid = ?',
                     (parentid, tagid))

    def get_tagchildren(self, tagid):
        '''Return ids of all child tags of a given tag.'''
        sql = 'select tagid from tagnames where tagparent = ?'
        for childid in self.query(sql, (tagid,)):
            yield childid[0]
    
    def get_tagids(self, photoid):
        '''Return list of tagids that apply to a photo.'''
        sql = 'select tagid from tags where photoid = ?'
        tagids = []
        for row in self.query(sql, (photoid,)):
            tagids.append(row[0])
        return tagids

    def add_tagid(self, photoid, tagid):
        '''Add a tagid for a photo.'''
        sql = 'insert into tags (photoid, tagid) values (?, ?)'
        self.execute(sql, (photoid, tagid))

    def remove_tagid(self, photoid, tagid):
        '''Remove a tagid for a photo.'''
        sql = 'delete from tags where photoid = ? and tagid = ?'
        self.execute(sql, (photoid, tagid))
        
    def set_rating(self, photoid, rating):
        '''Set rating for a photo.'''
        sql = 'update photos set rating = ? where photoid = ?'
        self.execute(sql, (rating, photoid))

    def set_rotate(self, photoid, rotate):
        '''Set rotation angle for a photo.'''
        sql = 'update photos set rotate = ? where photoid = ?'
        self.execute(sql, (rotate, photoid))