← Back to branch summary

~azzar1/unity/add-show-desktop-key

« back to all changes in this revision

Viewing changes to lib/common/db.py

Added printing of full traceback when an error is caught.

Show diffs side-by-side

added added

removed removed

Lines of Context:
lib/common/db.py
1
 
# IVLE - Informatics Virtual Learning Environment
2
 
# Copyright (C) 2007-2008 The University of Melbourne
3
 
#
4
 
# This program is free software; you can redistribute it and/or modify
5
 
# it under the terms of the GNU General Public License as published by
6
 
# the Free Software Foundation; either version 2 of the License, or
7
 
# (at your option) any later version.
8
 
#
9
 
# This program is distributed in the hope that it will be useful,
10
 
# but WITHOUT ANY WARRANTY; without even the implied warranty of
11
 
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
 
# GNU General Public License for more details.
13
 
#
14
 
# You should have received a copy of the GNU General Public License
15
 
# along with this program; if not, write to the Free Software
16
 
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
17
 
 
18
 
# Module: Database
19
 
# Author: Matt Giuca
20
 
# Date:   15/2/2008
21
 
 
22
 
# Code to talk to the PostgreSQL database.
23
 
# (This is the Data Access Layer).
24
 
# All DB code should be in this module to ensure portability if we want to
25
 
# change the DB implementation.
26
 
# This means no SQL strings should be outside of this module. Add functions
27
 
# here to perform the activities needed, and place the SQL code for those
28
 
# activities within.
29
 
 
30
 
# CAUTION to editors of this module.
31
 
# All string inputs must be sanitized by calling _escape before being
32
 
# formatted into an SQL query string.
33
 
 
34
 
import pg
35
 
import conf
36
 
import md5
37
 
import copy
38
 
 
39
 
def _escape(val):
40
 
    """Wrapper around pg.escape_string. Prepares the Python value for use in
41
 
    SQL. Returns a string, which may be safely placed verbatim into an SQL
42
 
    query.
43
 
    Handles the following types:
44
 
    * str: Escapes the string, and also quotes it.
45
 
    * int/long/float: Just converts to an unquoted string.
46
 
    * bool: Returns as "TRUE" or "FALSE", unquoted.
47
 
    * NoneType: Returns "NULL", unquoted.
48
 
    Raises a DBException if val has an unsupported type.
49
 
    """
50
 
    # "E'" is postgres's way of making "escape" strings.
51
 
    # Such strings allow backslashes to escape things. Since escape_string
52
 
    # converts a single backslash into two backslashes, it needs to be fed
53
 
    # into E mode.
54
 
    # Ref: http://www.postgresql.org/docs/8.2/static/sql-syntax-lexical.html
55
 
    # WARNING: PostgreSQL-specific code
56
 
    if val is None:
57
 
        return "NULL"
58
 
    elif isinstance(val, str):
59
 
        return "E'" + pg.escape_string(val) + "'"
60
 
    elif isinstance(val, bool):
61
 
        return "TRUE" if val else "FALSE"
62
 
    elif isinstance(val, int) or isinstance(val, long) \
63
 
        or isinstance(val, float):
64
 
        return str(val)
65
 
    else:
66
 
        raise DBException("Attempt to insert an unsupported type "
67
 
            "into the database")
68
 
 
69
 
def _passhash(password):
70
 
    return md5.md5(password).hexdigest()
71
 
 
72
 
class DBException(Exception):
73
 
    """A DBException is for bad conditions in the database or bad input to
74
 
    these methods. If Postgres throws an exception it does not get rebadged.
75
 
    This is only for additional exceptions."""
76
 
    pass
77
 
 
78
 
class DB:
79
 
    """An IVLE database object. This object provides an interface to
80
 
    interacting with the IVLE database without using any external SQL.
81
 
 
82
 
    Most methods of this class have an optional dry argument. If true, they
83
 
    will return the SQL query string and NOT actually execute it. (For
84
 
    debugging purposes).
85
 
 
86
 
    Methods may throw db.DBException, or any of the pg exceptions as well.
87
 
    (In general, be prepared to catch exceptions!)
88
 
    """
89
 
    def __init__(self):
90
 
        """Connects to the database and creates a DB object.
91
 
        Takes no parameters - gets all the DB info from the configuration."""
92
 
        self.db = pg.connect(dbname=conf.db_dbname, host=conf.db_host,
93
 
                port=conf.db_port, user=conf.db_user, passwd=conf.db_password)
94
 
        self.open = True
95
 
 
96
 
    def __del__(self):
97
 
        if self.open:
98
 
            self.db.close()
99
 
 
100
 
    # GENERIC DB FUNCTIONS #
101
 
 
102
 
    @staticmethod
103
 
    def check_dict(dict, tablefields, disallowed=frozenset([]), must=False):
104
 
        """Checks that a dict does not contain keys that are not fields
105
 
        of the specified table.
106
 
        dict: A mapping from string keys to values; the keys are checked to
107
 
            see that they correspond to login table fields.
108
 
        tablefields: Collection of strings for field names in the table.
109
 
            Only these fields will be allowed.
110
 
        disallowed: Optional collection of strings for field names that are
111
 
            not allowed.
112
 
        must: If True, the dict MUST contain all fields in tablefields.
113
 
            If False, it may contain any subset of the fields.
114
 
        Returns True if the dict is valid, False otherwise.
115
 
        """
116
 
        allowed = frozenset(tablefields) - frozenset(disallowed)
117
 
        dictkeys = frozenset(dict.keys())
118
 
        if must:
119
 
            return allowed == dictkeys
120
 
        else:
121
 
            return allowed.issuperset(dictkeys)
122
 
 
123
 
    def insert(self, dict, tablename, tablefields, disallowed=frozenset([]),
124
 
        dry=False):
125
 
        """Inserts a new row in a table, using data from a supplied
126
 
        dictionary (which will be checked by check_dict).
127
 
        dict: Dictionary mapping column names to values. The values may be
128
 
            any of the following types:
129
 
            str, int, long, float, NoneType.
130
 
        tablename: String, name of the table to insert into. Will NOT be
131
 
            escaped - must be a valid identifier.
132
 
        tablefields, disallowed: see check_dict.
133
 
        dry: Returns the SQL query as a string, and does not execute it.
134
 
        Raises a DBException if the dictionary contains invalid fields.
135
 
        """
136
 
        if not DB.check_dict(dict, tablefields, disallowed):
137
 
            raise DBException("Supplied dictionary contains invalid fields.")
138
 
        # Build two lists concurrently: field names and values, as SQL strings
139
 
        fieldnames = []
140
 
        values = []
141
 
        for k,v in dict.items():
142
 
            fieldnames.append(k)
143
 
            values.append(_escape(v))
144
 
        if len(fieldnames) == 0: return
145
 
        fieldnames = ', '.join(fieldnames)
146
 
        values = ', '.join(values)
147
 
        query = ("INSERT INTO %s (%s) VALUES (%s);"
148
 
            % (tablename, fieldnames, values))
149
 
        if dry: return query
150
 
        self.db.query(query)
151
 
 
152
 
    def update(self, primarydict, updatedict, tablename, tablefields,
153
 
        primary_keys, disallowed_update=frozenset([]), dry=False):
154
 
        """Updates a row in a table, matching against primarydict to find the
155
 
        row, and using the data in updatedict (which will be checked by
156
 
        check_dict).
157
 
        primarydict: Dict mapping column names to values. The keys should be
158
 
            the table's primary key. Only rows which match this dict's values
159
 
            will be updated.
160
 
        updatedict: Dict mapping column names to values. The columns will be
161
 
            updated with the given values for the matched rows.
162
 
        tablename, tablefields, disallowed_update: See insert.
163
 
        primary_keys: Collection of strings which together form the primary
164
 
            key for this table. primarydict must contain all of these as keys,
165
 
            and only these keys.
166
 
        """
167
 
        if (not (DB.check_dict(primarydict, primary_keys, must=True)
168
 
            and DB.check_dict(updatedict, tablefields, disallowed_update))):
169
 
            raise DBException("Supplied dictionary contains invalid or "
170
 
                " missing fields.")
171
 
        # Make a list of SQL fragments of the form "field = 'new value'"
172
 
        # These fragments are ALREADY-ESCAPED
173
 
        setlist = []
174
 
        for k,v in updatedict.items():
175
 
            setlist.append("%s = %s" % (k, _escape(v)))
176
 
        wherelist = []
177
 
        for k,v in primarydict.items():
178
 
            wherelist.append("%s = %s" % (k, _escape(v)))
179
 
        if len(setlist) == 0 or len(wherelist) == 0:
180
 
            return
181
 
        # Join the fragments into a comma-separated string
182
 
        setstring = ', '.join(setlist)
183
 
        wherestring = ' AND '.join(wherelist)
184
 
        # Build the whole query as an UPDATE statement
185
 
        query = ("UPDATE %s SET %s WHERE %s;"
186
 
            % (tablename, setstring, wherestring))
187
 
        if dry: return query
188
 
        self.db.query(query)
189
 
 
190
 
    def delete(self, primarydict, tablename, primary_keys, dry=False):
191
 
        """Deletes a row in the table, matching against primarydict to find
192
 
        the row.
193
 
        primarydict, tablename, primary_keys: See update.
194
 
        """
195
 
        if not DB.check_dict(primarydict, primary_keys, must=True):
196
 
            raise DBException("Supplied dictionary contains invalid or "
197
 
                " missing fields.")
198
 
        wherelist = []
199
 
        for k,v in primarydict.items():
200
 
            wherelist.append("%s = %s" % (k, _escape(v)))
201
 
        if len(wherelist) == 0:
202
 
            return
203
 
        wherestring = ' AND '.join(wherelist)
204
 
        query = ("DELETE FROM %s WHERE %s;" % (tablename, wherestring))
205
 
        if dry: return query
206
 
        self.db.query(query)
207
 
 
208
 
    def get_single(self, primarydict, tablename, getfields, primary_keys,
209
 
        error_notfound="No rows found", dry=False):
210
 
        """Retrieves a single row from a table, returning it as a dictionary
211
 
        mapping field names to values. Matches against primarydict to find the
212
 
        row.
213
 
        primarydict, tablename, primary_keys: See update/delete.
214
 
        getfields: Collection of strings; the field names which will be
215
 
            returned as keys in the dictionary.
216
 
        error_notfound: Error message if 0 rows match.
217
 
        Raises a DBException if 0 rows match, with error_notfound as the msg.
218
 
        Raises an AssertError if >1 rows match (this should not happen if
219
 
            primary_keys is indeed the primary key).
220
 
        """
221
 
        if not DB.check_dict(primarydict, primary_keys, must=True):
222
 
            raise DBException("Supplied dictionary contains invalid or "
223
 
                " missing fields.")
224
 
        wherelist = []
225
 
        for k,v in primarydict.items():
226
 
            wherelist.append("%s = %s" % (k, _escape(v)))
227
 
        if len(getfields) == 0 or len(wherelist) == 0:
228
 
            return
229
 
        # Join the fragments into a comma-separated string
230
 
        getstring = ', '.join(getfields)
231
 
        wherestring = ' AND '.join(wherelist)
232
 
        # Build the whole query as an SELECT statement
233
 
        query = ("SELECT %s FROM %s WHERE %s;"
234
 
            % (getstring, tablename, wherestring))
235
 
        if dry: return query
236
 
        result = self.db.query(query)
237
 
        # Expecting exactly one
238
 
        if result.ntuples() != 1:
239
 
            # It should not be possible for ntuples to be greater than 1
240
 
            assert (result.ntuples() < 1)
241
 
            raise DBException(error_notfound)
242
 
        # Return as a dictionary
243
 
        return result.dictresult()[0]
244
 
 
245
 
    def get_all(self, tablename, getfields, dry=False):
246
 
        """Retrieves all rows from a table, returning it as a list of
247
 
        dictionaries mapping field names to values.
248
 
        tablename, getfields: See get_single.
249
 
        """
250
 
        if len(getfields) == 0:
251
 
            return
252
 
        getstring = ', '.join(getfields)
253
 
        query = ("SELECT %s FROM %s;" % (getstring, tablename))
254
 
        if dry: return query
255
 
        return self.db.query(query).dictresult()
256
 
 
257
 
    # USER MANAGEMENT FUNCTIONS #
258
 
 
259
 
    login_primary = frozenset(["login"])
260
 
    login_fields = frozenset([
261
 
        "login", "passhash", "state", "unixid", "email", "nick", "fullname",
262
 
        "rolenm", "studentid", "acct_exp", "pass_exp"
263
 
    ])
264
 
    # Do not return passhash when reading from the DB
265
 
    login_getfields = login_fields - frozenset(["passhash"])
266
 
 
267
 
    def create_user(self, dry=False, **kwargs):
268
 
        """Creates a user login entry in the database.
269
 
        All user fields are to be passed as args. The argument names
270
 
        are the field names of the "login" table of the DB schema.
271
 
        However, instead of supplying a "passhash", you must supply a
272
 
        "password" argument, which will be hashed internally.
273
 
        Also "state" must not given explicitly; it is implicitly set to
274
 
        "no_agreement".
275
 
        Raises an exception if the user already exists, or the dict contains
276
 
        invalid keys or is missing required keys.
277
 
        """
278
 
        if 'passhash' in kwargs:
279
 
            raise DBException("Supplied arguments include passhash (invalid).")
280
 
        # Make a copy of the dict. Change password to passhash (hashing it),
281
 
        # and set 'state' to "no_agreement".
282
 
        kwargs = copy.copy(kwargs)
283
 
        kwargs['passhash'] = _passhash(kwargs['password'])
284
 
        del kwargs['password']
285
 
        kwargs['state'] = "no_agreement"
286
 
        # Execute the query.
287
 
        return self.insert(kwargs, "login", self.login_fields, dry=dry)
288
 
 
289
 
    def update_user(self, login, dry=False, **kwargs):
290
 
        """Updates fields of a particular user. login is the name of the user
291
 
        to update. The dict contains the fields which will be modified, and
292
 
        their new values. If any value is omitted from the dict, it does not
293
 
        get modified. login and studentid may not be modified.
294
 
        Passhash may be modified by supplying a "password" field, in
295
 
        cleartext, not a hashed password.
296
 
 
297
 
        Note that no checking is done. It is expected this function is called
298
 
        by a trusted source. In particular, it allows the password to be
299
 
        changed without knowing the old password. The caller should check
300
 
        that the user knows the existing password before calling this function
301
 
        with a new one.
302
 
        """
303
 
        if 'passhash' in kwargs:
304
 
            raise DBException("Supplied arguments include passhash (invalid).")
305
 
        if "password" in kwargs:
306
 
            kwargs = copy.copy(kwargs)
307
 
            kwargs['passhash'] = _passhash(kwargs['password'])
308
 
            del kwargs['password']
309
 
        return self.update({"login": login}, kwargs, "login",
310
 
            self.login_fields, self.login_primary, ["login", "studentid"],
311
 
            dry=dry)
312
 
 
313
 
    def get_user(self, login, dry=False):
314
 
        """Given a login, returns a dictionary of the user's DB fields,
315
 
        excluding the passhash field.
316
 
 
317
 
        Raises a DBException if the login is not found in the DB.
318
 
        """
319
 
        return self.get_single({"login": login}, "login",
320
 
            self.login_getfields, self.login_primary,
321
 
            error_notfound="get_user: No user with that login name", dry=dry)
322
 
 
323
 
    def get_users(self, dry=False):
324
 
        """Returns a list of all users. The list elements are a dictionary of
325
 
        the user's DB fields, excluding the passhash field.
326
 
        """
327
 
        return self.get_all("login", self.login_getfields, dry=dry)
328
 
 
329
 
    def user_authenticate(self, login, password, dry=False):
330
 
        """Performs a password authentication on a user. Returns True if
331
 
        "passhash" is the correct passhash for the given login, False
332
 
        otherwise.
333
 
        Also returns False if the login does not exist (so if you want to
334
 
        differentiate these cases, use get_user and catch an exception).
335
 
        """
336
 
        query = ("SELECT login FROM login "
337
 
            "WHERE login = '%s' AND passhash = %s;"
338
 
            % (login, _escape(_passhash(password))))
339
 
        if dry: return query
340
 
        result = self.db.query(query)
341
 
        # If one row was returned, succeed.
342
 
        # Otherwise, fail to authenticate.
343
 
        return result.ntuples() == 1
344
 
 
345
 
    def close(self):
346
 
        """Close the DB connection. Do not call any other functions after
347
 
        this. (The behaviour of doing so is undefined).
348
 
        """
349
 
        self.db.close()
350
 
        self.open = False