~ubuntuone-pqm-team/canonical-identity-provider/trunk

« back to all changes in this revision

Viewing changes to identityprovider/backend/base.py

  • Committer: Danny Tamez
  • Date: 2010-04-21 15:29:24 UTC
  • Revision ID: danny.tamez@canonical.com-20100421152924-lq1m92tstk2iz75a
Canonical SSO Provider (Open Source) - Initial Commit

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright 2010 Canonical Ltd.  This software is licensed under the
 
2
# GNU Affero General Public License version 3 (see the file LICENSE).
 
3
 
 
4
__all__ = [
 
5
    'DatabaseWrapper', 'DatabaseError', 'IntegrityError',
 
6
    ]
 
7
 
 
8
import re
 
9
from hashlib import sha1
 
10
 
 
11
from django.db.backends.postgresql_psycopg2.base import (
 
12
                DatabaseWrapper as PostgresDatabaseWrapper)
 
13
from django.conf import settings
 
14
from django.core.cache import cache
 
15
from django.utils.translation import ugettext as _
 
16
 
 
17
 
 
18
class DatabaseError(Exception):
 
19
    pass
 
20
 
 
21
 
 
22
class IntegrityError(DatabaseError):
 
23
    pass
 
24
 
 
25
 
 
26
class BaseCache(object):
 
27
    """This class provides a cached version of a single database table.
 
28
 
 
29
    Queries are routed through memcached first, so only cache misses
 
30
    really hit the database.  Inserts and deletes are carried out
 
31
    in memcached only, so that the database is never written to.
 
32
 
 
33
    Subclasses need to implement update if needed.
 
34
    """
 
35
    table = None
 
36
    primary_key = None
 
37
 
 
38
    def insert(self, match, params, cursor):
 
39
        """Insert the row only into memcached."""
 
40
        cols = [x.strip('",') for x in match.group('cols').split()]
 
41
        values = dict(zip(cols, params))
 
42
        key = self.cache_key(values[self.primary_key])
 
43
        cache._cache.add(key, str([params]))
 
44
 
 
45
    def select(self, match, params, cursor):
 
46
        """Check memcached first, then hit the DB if not found."""
 
47
        pkey_cond = '"%s"."%s" = %%s' % (self.table, self.primary_key)
 
48
        assert pkey_cond in match.group('cond')
 
49
        key = self.cache_key(params[0])
 
50
        cached_value = cache._cache.get(key)
 
51
        if cached_value is None:
 
52
            cursor.execute(match.group(0), params)
 
53
            cached_value = cursor.fetchmany()
 
54
        else:
 
55
            cached_value = eval(cached_value)
 
56
        return cached_value
 
57
 
 
58
    def delete(self, match, params, cursor):
 
59
        """Mark the row as deleted in memcached.
 
60
 
 
61
        Note that future queries to this row will produce no results,
 
62
        even if there's an entry in the DB for it.
 
63
        """
 
64
        delete_pattern = r'"%s" IN \(%%s\)' % self.primary_key
 
65
        assert re.match(delete_pattern, match.group('cond'))
 
66
        for dbkey in params[0]:
 
67
            key = self.cache_key(dbkey)
 
68
            cache._cache.set(key, '[]')
 
69
 
 
70
    def cache_key(self, dbkey):
 
71
        """Returns a canonical memcached key for a row in a table."""
 
72
        hash = sha1(dbkey).hexdigest()
 
73
        return 'db-%s-%s' % (self.table, hash)
 
74
 
 
75
 
 
76
class OpenIDAssociationCache(BaseCache):
 
77
    table = 'openidassociation'
 
78
    primary_key = 'handle'
 
79
 
 
80
 
 
81
class DjangoSessionCache(BaseCache):
 
82
    table = 'django_session'
 
83
    primary_key = 'session_key'
 
84
 
 
85
    def update(self, match, params, cursor):
 
86
        """django_session is always updated in the same way, so we can map
 
87
        that in to a row in the database.
 
88
        """
 
89
        pkey_cond = '"%s"."%s" = %%s ' % (self.table, self.primary_key)
 
90
        assert pkey_cond == match.group('cond')
 
91
        update_format = '"session_data" = %s, "expire_date" = %s'
 
92
        assert update_format == match.group('cols')
 
93
        key = self.cache_key(params[2])
 
94
        new_value = [params[2], params[0], params[1]]
 
95
        cache._cache.set(key, str([new_value]))
 
96
 
 
97
 
 
98
class AuthUserCache(BaseCache):
 
99
    table = 'auth_user'
 
100
    primary_key = 'auth_user_pkey'
 
101
 
 
102
    def select(self, match, params, cursor):
 
103
        """Skip memcached completely."""
 
104
        cursor.execute(match.group(0), params)
 
105
        cached_value = cursor.fetchmany()
 
106
        return cached_value
 
107
 
 
108
    def update(self, match, params, cursor):
 
109
        """Does nothing.
 
110
        During readonly mode auth_user will only be updated to set
 
111
        last_login, so we ignore all updates.
 
112
        """
 
113
        pass
 
114
 
 
115
cached_tables = {
 
116
    'django_session': DjangoSessionCache(),
 
117
    'openidassociation': OpenIDAssociationCache(),
 
118
    'auth_user': AuthUserCache(),
 
119
    }
 
120
 
 
121
 
 
122
class CursorReadOnlyWrapper(object):
 
123
    delete_pattern = r'^DELETE FROM "(?P<table>.*)" WHERE (?P<cond>.*)$'
 
124
    insert_pattern = r'^INSERT INTO "(?P<table>.*)" \((?P<cols>.*)\) VALUES \((?P<values>.*)\)$'
 
125
    select_pattern = r'^SELECT (?P<cols>.*) FROM "(?P<table>.*)" WHERE (?P<cond>.*)$'
 
126
    update_pattern = r'^UPDATE "(?P<table>.*)" SET (?P<cols>.*) WHERE (?P<cond>.*)$'
 
127
 
 
128
    def __init__(self, cursor):
 
129
        self.cursor = cursor
 
130
        self.cache = None
 
131
 
 
132
    def execute_cached(self, command, sql, params):
 
133
        """Attempt to carry out a command against memcache.
 
134
 
 
135
        Return True if the command is successfully carried out.
 
136
        """
 
137
        pattern = getattr(self, command + '_pattern', None)
 
138
        if pattern is None:
 
139
            return False
 
140
        match = re.match(pattern, sql)
 
141
        if match is not None:
 
142
            table = match.group('table')
 
143
            if table in cached_tables:
 
144
                self.cache = cached_tables[table]
 
145
                method = getattr(self.cache, command)
 
146
                self._values = method(match, params, self.cursor)
 
147
                return True
 
148
        return False
 
149
 
 
150
    def execute(self, sql, params=()):
 
151
        command = sql.split(' ', 1)[0].lower()
 
152
        executed = self.execute_cached(command, sql, params)
 
153
        if executed:
 
154
            return
 
155
        if command in ['select', 'savepoint']:
 
156
            return self.cursor.execute(sql, params)
 
157
        else:
 
158
            raise DatabaseError(_('Attempted to %(command)s while in '
 
159
                            'read-only mode: \'%(sql)s\' %% (%(params)s)') %
 
160
                            {'command': command, 'sql': sql, 'params': params})
 
161
 
 
162
    def fetchmany(self, chunk):
 
163
        if self.cache is not None:
 
164
            if len(self._values) == 0:
 
165
                raise StopIteration()
 
166
            values = self._values[:chunk]
 
167
            self._values = self._values[chunk:]
 
168
            return values
 
169
        else:
 
170
            return self.cursor.fetchmany(chunk)
 
171
 
 
172
    def __getattr__(self, attr):
 
173
        if attr in self.__dict__:
 
174
            return self.__dict__[attr]
 
175
        else:
 
176
            return getattr(self.cursor, attr)
 
177
 
 
178
 
 
179
class DatabaseWrapper(PostgresDatabaseWrapper):
 
180
    def _cursor(self, *args):
 
181
        cursor = super(DatabaseWrapper, self)._cursor(*args)
 
182
        if getattr(settings, 'READ_ONLY_MODE', False):
 
183
            cursor = CursorReadOnlyWrapper(cursor)
 
184
        return cursor
 
185
 
 
186
    def cursor(self):
 
187
        from django.conf import settings
 
188
        cursor = self._cursor(settings)
 
189
        return self.make_debug_cursor(cursor)