Project

General

Profile

1 11 aaronmk
# Database access
2
3 1869 aaronmk
import copy
4 2127 aaronmk
import operator
5 11 aaronmk
import re
6 865 aaronmk
import warnings
7 11 aaronmk
8 300 aaronmk
import exc
9 1909 aaronmk
import dicts
10 1893 aaronmk
import iters
11 1960 aaronmk
import lists
12 1889 aaronmk
from Proxy import Proxy
13 1872 aaronmk
import rand
14 862 aaronmk
import strings
15 131 aaronmk
import util
16 11 aaronmk
17 832 aaronmk
##### Exceptions
18
19 2170 aaronmk
def get_cur_query(cur, input_query=None, input_params=None):
20 2168 aaronmk
    raw_query = None
21
    if hasattr(cur, 'query'): raw_query = cur.query
22
    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
23 2170 aaronmk
24
    if raw_query != None: return raw_query
25
    else: return repr(input_query)+' % '+repr(input_params)
26 14 aaronmk
27 2170 aaronmk
def _add_cursor_info(e, *args, **kw_args):
28
    '''For params, see get_cur_query()'''
29
    exc.add_msg(e, 'query: '+str(get_cur_query(*args, **kw_args)))
30 135 aaronmk
31 300 aaronmk
class DbException(exc.ExceptionWithCause):
32 14 aaronmk
    def __init__(self, msg, cause=None, cur=None):
33 2145 aaronmk
        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)
34 14 aaronmk
        if cur != None: _add_cursor_info(self, cur)
35
36 2143 aaronmk
class ExceptionWithName(DbException):
37
    def __init__(self, name, cause=None):
38 2145 aaronmk
        DbException.__init__(self, 'for name: '+str(name), cause)
39 2143 aaronmk
        self.name = name
40 360 aaronmk
41 468 aaronmk
class ExceptionWithColumns(DbException):
42
    def __init__(self, cols, cause=None):
43 2145 aaronmk
        DbException.__init__(self, 'for columns: '+(', '.join(cols)), cause)
44 468 aaronmk
        self.cols = cols
45 11 aaronmk
46 2143 aaronmk
class NameException(DbException): pass
47
48 468 aaronmk
class DuplicateKeyException(ExceptionWithColumns): pass
49 13 aaronmk
50 468 aaronmk
class NullValueException(ExceptionWithColumns): pass
51 13 aaronmk
52 2143 aaronmk
class DuplicateTableException(ExceptionWithName): pass
53
54 2188 aaronmk
class DuplicateFunctionException(ExceptionWithName): pass
55
56 89 aaronmk
class EmptyRowException(DbException): pass
57
58 865 aaronmk
##### Warnings
59
60
class DbWarning(UserWarning): pass
61
62 1930 aaronmk
##### Result retrieval
63
64
def col_names(cur): return (col[0] for col in cur.description)
65
66
def rows(cur): return iter(lambda: cur.fetchone(), None)
67
68
def consume_rows(cur):
69
    '''Used to fetch all rows so result will be cached'''
70
    iters.consume_iter(rows(cur))
71
72
def next_row(cur): return rows(cur).next()
73
74
def row(cur):
75
    row_ = next_row(cur)
76
    consume_rows(cur)
77
    return row_
78
79
def next_value(cur): return next_row(cur)[0]
80
81
def value(cur): return row(cur)[0]
82
83
def values(cur): return iters.func_iter(lambda: next_value(cur))
84
85
def value_or_none(cur):
86
    try: return value(cur)
87
    except StopIteration: return None
88
89 2101 aaronmk
##### Input validation
90
91
def clean_name(name): return re.sub(r'\W', r'', name)
92
93
def check_name(name):
94
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
95
        +'" may contain only alphanumeric characters and _')
96
97
def esc_name_by_module(module, name, ignore_case=False):
98
    if module == 'psycopg2':
99
        if ignore_case:
100
            # Don't enclose in quotes because this disables case-insensitivity
101
            check_name(name)
102
            return name
103
        else: quote = '"'
104
    elif module == 'MySQLdb': quote = '`'
105
    else: raise NotImplementedError("Can't escape name for "+module+' database')
106
    return quote + name.replace(quote, '') + quote
107
108
def esc_name_by_engine(engine, name, **kw_args):
109
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
110
111
def esc_name(db, name, **kw_args):
112
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
113
114
def qual_name(db, schema, table):
115
    def esc_name_(name): return esc_name(db, name)
116
    table = esc_name_(table)
117
    if schema != None: return esc_name_(schema)+'.'+table
118
    else: return table
119
120 1869 aaronmk
##### Database connections
121 1849 aaronmk
122 2097 aaronmk
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
123 1926 aaronmk
124 1869 aaronmk
db_engines = {
125
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
126
    'PostgreSQL': ('psycopg2', {}),
127
}
128
129
DatabaseErrors_set = set([DbException])
130
DatabaseErrors = tuple(DatabaseErrors_set)
131
132
def _add_module(module):
133
    DatabaseErrors_set.add(module.DatabaseError)
134
    global DatabaseErrors
135
    DatabaseErrors = tuple(DatabaseErrors_set)
136
137
def db_config_str(db_config):
138
    return db_config['engine']+' database '+db_config['database']
139
140 1909 aaronmk
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
141 1894 aaronmk
142 1901 aaronmk
log_debug_none = lambda msg: None
143
144 1849 aaronmk
class DbConn:
145 2047 aaronmk
    def __init__(self, db_config, serializable=True, log_debug=log_debug_none,
146 2050 aaronmk
        caching=True):
147 1869 aaronmk
        self.db_config = db_config
148
        self.serializable = serializable
149 1901 aaronmk
        self.log_debug = log_debug
150 2047 aaronmk
        self.caching = caching
151 1869 aaronmk
152
        self.__db = None
153 1889 aaronmk
        self.query_results = {}
154 2139 aaronmk
        self._savepoint = 0
155 1869 aaronmk
156
    def __getattr__(self, name):
157
        if name == '__dict__': raise Exception('getting __dict__')
158
        if name == 'db': return self._db()
159
        else: raise AttributeError()
160
161
    def __getstate__(self):
162
        state = copy.copy(self.__dict__) # shallow copy
163 1915 aaronmk
        state['log_debug'] = None # don't pickle the debug callback
164 1869 aaronmk
        state['_DbConn__db'] = None # don't pickle the connection
165
        return state
166
167 2165 aaronmk
    def connected(self): return self.__db != None
168
169 1869 aaronmk
    def _db(self):
170
        if self.__db == None:
171
            # Process db_config
172
            db_config = self.db_config.copy() # don't modify input!
173 2097 aaronmk
            schemas = db_config.pop('schemas', None)
174 1869 aaronmk
            module_name, mappings = db_engines[db_config.pop('engine')]
175
            module = __import__(module_name)
176
            _add_module(module)
177
            for orig, new in mappings.iteritems():
178
                try: util.rename_key(db_config, orig, new)
179
                except KeyError: pass
180
181
            # Connect
182
            self.__db = module.connect(**db_config)
183
184
            # Configure connection
185
            if self.serializable: run_raw_query(self,
186
                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')
187 2101 aaronmk
            if schemas != None:
188
                schemas_ = ''.join((esc_name(self, s)+', '
189
                    for s in schemas.split(',')))
190
                run_raw_query(self, "SELECT set_config('search_path', \
191
%s || current_setting('search_path'), false)", [schemas_])
192 1869 aaronmk
193
        return self.__db
194 1889 aaronmk
195 1891 aaronmk
    class DbCursor(Proxy):
196 1927 aaronmk
        def __init__(self, outer):
197 1891 aaronmk
            Proxy.__init__(self, outer.db.cursor())
198 1927 aaronmk
            self.query_results = outer.query_results
199 1894 aaronmk
            self.query_lookup = None
200 1891 aaronmk
            self.result = []
201 1889 aaronmk
202 1894 aaronmk
        def execute(self, query, params=None):
203 1930 aaronmk
            self._is_insert = query.upper().find('INSERT') >= 0
204 1894 aaronmk
            self.query_lookup = _query_lookup(query, params)
205 2148 aaronmk
            try:
206
                try: return_value = self.inner.execute(query, params)
207
                finally: self.query = get_cur_query(self.inner)
208 1904 aaronmk
            except Exception, e:
209 2170 aaronmk
                _add_cursor_info(e, self, query, params)
210 1904 aaronmk
                self.result = e # cache the exception as the result
211
                self._cache_result()
212
                raise
213 1930 aaronmk
            # Fetch all rows so result will be cached
214
            if self.rowcount == 0 and not self._is_insert: consume_rows(self)
215 1894 aaronmk
            return return_value
216
217 1891 aaronmk
        def fetchone(self):
218
            row = self.inner.fetchone()
219 1899 aaronmk
            if row != None: self.result.append(row)
220
            # otherwise, fetched all rows
221 1904 aaronmk
            else: self._cache_result()
222
            return row
223
224
        def _cache_result(self):
225 1906 aaronmk
            # For inserts, only cache exceptions since inserts are not
226
            # idempotent, but an invalid insert will always be invalid
227 1930 aaronmk
            if self.query_results != None and (not self._is_insert
228 1906 aaronmk
                or isinstance(self.result, Exception)):
229
230 1894 aaronmk
                assert self.query_lookup != None
231 1916 aaronmk
                self.query_results[self.query_lookup] = self.CacheCursor(
232
                    util.dict_subset(dicts.AttrsDictView(self),
233
                    ['query', 'result', 'rowcount', 'description']))
234 1906 aaronmk
235 1916 aaronmk
        class CacheCursor:
236
            def __init__(self, cached_result): self.__dict__ = cached_result
237
238 1927 aaronmk
            def execute(self, *args, **kw_args):
239 1916 aaronmk
                if isinstance(self.result, Exception): raise self.result
240
                # otherwise, result is a rows list
241
                self.iter = iter(self.result)
242
243
            def fetchone(self):
244
                try: return self.iter.next()
245
                except StopIteration: return None
246 1891 aaronmk
247 1894 aaronmk
    def run_query(self, query, params=None, cacheable=False):
248 2148 aaronmk
        '''Translates known DB errors to typed exceptions:
249
        See self.DbCursor.execute().'''
250 2167 aaronmk
        assert query != None
251
252 2047 aaronmk
        if not self.caching: cacheable = False
253 1903 aaronmk
        used_cache = False
254
        try:
255 1927 aaronmk
            # Get cursor
256
            if cacheable:
257
                query_lookup = _query_lookup(query, params)
258
                try:
259
                    cur = self.query_results[query_lookup]
260
                    used_cache = True
261
                except KeyError: cur = self.DbCursor(self)
262
            else: cur = self.db.cursor()
263
264
            # Run query
265 2148 aaronmk
            cur.execute(query, params)
266 1903 aaronmk
        finally:
267
            if self.log_debug != log_debug_none: # only compute msg if needed
268
                if used_cache: cache_status = 'Cache hit'
269
                elif cacheable: cache_status = 'Cache miss'
270
                else: cache_status = 'Non-cacheable'
271 1927 aaronmk
                self.log_debug(cache_status+': '
272 2170 aaronmk
                    +strings.one_line(str(get_cur_query(cur, query, params))))
273 1903 aaronmk
274
        return cur
275 1914 aaronmk
276
    def is_cached(self, query, params=None):
277
        return _query_lookup(query, params) in self.query_results
278 2139 aaronmk
279
    def with_savepoint(self, func):
280 2171 aaronmk
        savepoint = 'level_'+str(self._savepoint)
281 2139 aaronmk
        self.run_query('SAVEPOINT '+savepoint)
282
        self._savepoint += 1
283
        try:
284
            try: return_val = func()
285
            finally:
286
                self._savepoint -= 1
287
                assert self._savepoint >= 0
288
        except:
289
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint)
290
            raise
291
        else:
292
            self.run_query('RELEASE SAVEPOINT '+savepoint)
293
            return return_val
294 1849 aaronmk
295 1869 aaronmk
connect = DbConn
296
297 832 aaronmk
##### Querying
298
299 1894 aaronmk
def run_raw_query(db, *args, **kw_args):
300 2085 aaronmk
    '''For params, see DbConn.run_query()'''
301 1894 aaronmk
    return db.run_query(*args, **kw_args)
302 11 aaronmk
303 2068 aaronmk
def mogrify(db, query, params):
304
    module = util.root_module(db.db)
305
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
306
    else: raise NotImplementedError("Can't mogrify query for "+module+
307
        ' database')
308
309 832 aaronmk
##### Recoverable querying
310 15 aaronmk
311 2139 aaronmk
def with_savepoint(db, func): return db.with_savepoint(func)
312 11 aaronmk
313 1894 aaronmk
def run_query(db, query, params=None, recover=None, cacheable=False):
314 830 aaronmk
    if recover == None: recover = False
315
316 2148 aaronmk
    try:
317
        def run(): return run_raw_query(db, query, params, cacheable)
318
        if recover and not db.is_cached(query, params):
319
            return with_savepoint(db, run)
320
        else: return run() # don't need savepoint if cached
321
    except Exception, e:
322
        if not recover: raise # need savepoint to run index_cols()
323
        msg = str(e)
324
        match = re.search(r'duplicate key value violates unique constraint '
325
            r'"((_?[^\W_]+)_[^"]+)"', msg)
326
        if match:
327
            constraint, table = match.groups()
328
            try: cols = index_cols(db, table, constraint)
329
            except NotImplementedError: raise e
330
            else: raise DuplicateKeyException(cols, e)
331
        match = re.search(r'null value in column "(\w+)" violates not-null '
332
            'constraint', msg)
333
        if match: raise NullValueException([match.group(1)], e)
334
        match = re.search(r'relation "(\w+)" already exists', msg)
335
        if match: raise DuplicateTableException(match.group(1), e)
336 2188 aaronmk
        match = re.search(r'function "(\w+)" already exists', msg)
337
        if match: raise DuplicateFunctionException(match.group(1), e)
338 2148 aaronmk
        raise # no specific exception raised
339 830 aaronmk
340 832 aaronmk
##### Basic queries
341
342 2153 aaronmk
def next_version(name):
343
    '''Prepends the version # so it won't be removed if the name is truncated'''
344 2163 aaronmk
    version = 1 # first existing name was version 0
345 2153 aaronmk
    match = re.match(r'^v(\d+)_(.*)$', name)
346
    if match:
347
        version = int(match.group(1))+1
348
        name = match.group(2)
349
    return 'v'+str(version)+'_'+name
350
351 2151 aaronmk
def run_query_into(db, query, params, into_ref=None, *args, **kw_args):
352 2085 aaronmk
    '''Outputs a query to a temp table.
353
    For params, see run_query().
354
    '''
355 2151 aaronmk
    if into_ref == None: return run_query(db, query, params, *args, **kw_args)
356 2085 aaronmk
    else: # place rows in temp table
357 2151 aaronmk
        check_name(into_ref[0])
358 2153 aaronmk
        kw_args['recover'] = True
359
        while True:
360
            try:
361
                return run_query(db, 'CREATE TEMP TABLE '+into_ref[0]+' AS '
362
                    +query, params, *args, **kw_args)
363
                    # CREATE TABLE AS sets rowcount to # rows in query
364
            except DuplicateTableException, e:
365
                into_ref[0] = next_version(into_ref[0])
366
                # try again with next version of name
367 2085 aaronmk
368 2120 aaronmk
order_by_pkey = object() # tells mk_select() to order by the pkey
369
370 2127 aaronmk
join_using = object() # tells mk_select() to join the column with USING
371
372 2187 aaronmk
filter_out = object() # tells mk_select() to filter out rows that match the join
373 2180 aaronmk
374 2121 aaronmk
def mk_select(db, tables, fields=None, conds=None, limit=None, start=None,
375 2120 aaronmk
    order_by=order_by_pkey, table_is_esc=False):
376 1981 aaronmk
    '''
377 2121 aaronmk
    @param tables The single table to select from, or a list of tables to join
378 2187 aaronmk
        together: [table0, (table1, joins), ...]
379
380
        joins has the format: dict(right_col=left_col, ...)
381
        * if left_col is join_using, left_col is set to right_col
382
        * if left_col is filter_out, the tables are LEFT JOINed together and the
383
          query is filtered by `right_col IS NULL` (indicating no match)
384 1981 aaronmk
    @param fields Use None to select all fields in the table
385
    @param table_is_esc Whether the table name has already been escaped
386 2054 aaronmk
    @return tuple(query, params)
387 1981 aaronmk
    '''
388 2060 aaronmk
    def esc_name_(name): return esc_name(db, name)
389 2058 aaronmk
390 2121 aaronmk
    if not lists.is_seq(tables): tables = [tables]
391 2141 aaronmk
    tables = list(tables) # don't modify input! (list() copies input)
392 2121 aaronmk
    table0 = tables.pop(0) # first table is separate
393
394 1135 aaronmk
    if conds == None: conds = {}
395 135 aaronmk
    assert limit == None or type(limit) == int
396 865 aaronmk
    assert start == None or type(start) == int
397 2120 aaronmk
    if order_by == order_by_pkey:
398 2121 aaronmk
        order_by = pkey(db, table0, recover=True, table_is_esc=table_is_esc)
399
    if not table_is_esc: table0 = esc_name_(table0)
400 865 aaronmk
401 2056 aaronmk
    params = []
402
403 2161 aaronmk
    def parse_col(field, default_table=None):
404 2056 aaronmk
        '''Parses fields'''
405 2176 aaronmk
        if field == None: field = (field,) # for None values, tuple is optional
406 2157 aaronmk
        is_tuple = isinstance(field, tuple)
407
        if is_tuple and len(field) == 1: # field is literal value
408
            value, = field
409 2056 aaronmk
            sql_ = '%s'
410
            params.append(value)
411 2157 aaronmk
        elif is_tuple and len(field) == 2: # field is col with table
412
            table, col = field
413
            if not table_is_esc: table = esc_name_(table)
414
            sql_ = table+'.'+esc_name_(col)
415 2161 aaronmk
        else:
416
            sql_ = esc_name_(field) # field is col name
417
            if default_table != None: sql_ = default_table+'.'+sql_
418 2056 aaronmk
        return sql_
419 11 aaronmk
    def cond(entry):
420 2056 aaronmk
        '''Parses conditions'''
421 13 aaronmk
        col, value = entry
422 2187 aaronmk
        cond_ = parse_col(col)+' '
423 11 aaronmk
        if value == None: cond_ += 'IS'
424
        else: cond_ += '='
425
        cond_ += ' %s'
426
        return cond_
427 2056 aaronmk
428 1135 aaronmk
    query = 'SELECT '
429
    if fields == None: query += '*'
430 2056 aaronmk
    else: query += ', '.join(map(parse_col, fields))
431 2121 aaronmk
    query += ' FROM '+table0
432 865 aaronmk
433 2122 aaronmk
    # Add joins
434
    left_table = table0
435
    for table, joins in tables:
436
        if not table_is_esc: table = esc_name_(table)
437
438 2187 aaronmk
        left_join_ref = [False]
439
440 2122 aaronmk
        def join(entry):
441 2127 aaronmk
            '''Parses non-USING joins'''
442 2124 aaronmk
            right_col, left_col = entry
443 2173 aaronmk
444 2179 aaronmk
            # Parse special values
445 2176 aaronmk
            if left_col == None: left_col = (left_col,)
446
                # for None values, tuple is optional
447 2179 aaronmk
            elif left_col == join_using: left_col = right_col
448 2187 aaronmk
            elif left_col == filter_out:
449 2180 aaronmk
                left_col = right_col
450 2187 aaronmk
                left_join_ref[0] = True
451
                conds[(table, right_col)] = None # filter query by no match
452 2179 aaronmk
453
            # Create SQL
454 2185 aaronmk
            right_col = table+'.'+esc_name_(right_col)
455
            sql_ = right_col+' '
456 2173 aaronmk
            if isinstance(left_col, tuple) and len(left_col) == 1:
457
                # col is literal value
458
                value, = left_col
459
                if value == None: sql_ += 'IS'
460
                else: sql_ += '='
461
                sql_ += ' %s'
462
                params.append(value)
463
            else: # col is name
464
                left_col = parse_col(left_col, left_table)
465
                sql_ += ('= '+left_col+' OR ('+right_col+' IS NULL AND '
466
                    +left_col+' IS NULL)')
467
468
            return sql_
469 2122 aaronmk
470 2187 aaronmk
        # Create join condition and determine join type
471 2127 aaronmk
        if reduce(operator.and_, (v == join_using for v in joins.itervalues())):
472 2179 aaronmk
            # all cols w/ USING, so can use simpler USING syntax
473 2187 aaronmk
            join_cond = 'USING ('+(', '.join(joins.iterkeys()))+')'
474
        else: join_cond = 'ON '+(' AND '.join(map(join, joins.iteritems())))
475 2127 aaronmk
476 2187 aaronmk
        # Create join
477
        if left_join_ref[0]: query += ' LEFT'
478
        query += ' JOIN '+table+' '+join_cond
479
480 2122 aaronmk
        left_table = table
481
482 865 aaronmk
    missing = True
483 89 aaronmk
    if conds != {}:
484 2122 aaronmk
        query += ' WHERE '+(' AND '.join(map(cond, conds.iteritems())))
485 2056 aaronmk
        params += conds.values()
486 865 aaronmk
        missing = False
487 2186 aaronmk
    if order_by != None: query += ' ORDER BY '+parse_col(order_by, table0)
488 865 aaronmk
    if limit != None: query += ' LIMIT '+str(limit); missing = False
489
    if start != None:
490
        if start != 0: query += ' OFFSET '+str(start)
491
        missing = False
492
    if missing: warnings.warn(DbWarning(
493
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
494
495 2056 aaronmk
    return (query, params)
496 11 aaronmk
497 2054 aaronmk
def select(db, *args, **kw_args):
498
    '''For params, see mk_select() and run_query()'''
499
    recover = kw_args.pop('recover', None)
500
    cacheable = kw_args.pop('cacheable', True)
501
502
    query, params = mk_select(db, *args, **kw_args)
503
    return run_query(db, query, params, recover, cacheable)
504
505 2066 aaronmk
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
506 2070 aaronmk
    returning=None, embeddable=False, table_is_esc=False):
507 1960 aaronmk
    '''
508
    @param returning str|None An inserted column (such as pkey) to return
509 2070 aaronmk
    @param embeddable Whether the query should be embeddable as a nested SELECT.
510 2073 aaronmk
        Warning: If you set this and cacheable=True when the query is run, the
511
        query will be fully cached, not just if it raises an exception.
512 1960 aaronmk
    @param table_is_esc Whether the table name has already been escaped
513
    '''
514 2063 aaronmk
    if select_query == None: select_query = 'DEFAULT VALUES'
515
    if cols == []: cols = None # no cols (all defaults) = unknown col names
516 1960 aaronmk
    if not table_is_esc: check_name(table)
517 2063 aaronmk
518
    # Build query
519
    query = 'INSERT INTO '+table
520
    if cols != None:
521
        map(check_name, cols)
522
        query += ' ('+', '.join(cols)+')'
523
    query += ' '+select_query
524
525
    if returning != None:
526
        check_name(returning)
527
        query += ' RETURNING '+returning
528
529 2070 aaronmk
    if embeddable:
530
        # Create function
531 2189 aaronmk
        function_name = '_'.join(map(clean_name, ['insert', table] + cols))
532 2070 aaronmk
        return_type = 'SETOF '+table+'.'+returning+'%TYPE'
533 2189 aaronmk
        while True:
534
            try:
535
                function = 'pg_temp.'+function_name
536
                function_query = '''\
537
CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
538 2070 aaronmk
    LANGUAGE sql
539
    AS $$'''+mogrify(db, query, params)+''';$$;
540
'''
541 2189 aaronmk
                run_query(db, function_query, recover=True, cacheable=True)
542
                break # this version was successful
543
            except DuplicateFunctionException, e:
544
                function_name = next_version(function_name)
545
                # try again with next version of name
546 2070 aaronmk
547
        # Return query that uses function
548 2134 aaronmk
        return mk_select(db, function+'() AS f ('+returning+')', start=0,
549
            order_by=None, table_is_esc=True)# AS clause requires function alias
550 2070 aaronmk
551 2066 aaronmk
    return (query, params)
552
553
def insert_select(db, *args, **kw_args):
554 2085 aaronmk
    '''For params, see mk_insert_select() and run_query_into()
555 2152 aaronmk
    @param into_ref List with name of temp table to place RETURNING values in
556 2072 aaronmk
    '''
557 2151 aaronmk
    into_ref = kw_args.pop('into_ref', None)
558
    if into_ref != None: kw_args['embeddable'] = True
559 2066 aaronmk
    recover = kw_args.pop('recover', None)
560
    cacheable = kw_args.pop('cacheable', True)
561
562
    query, params = mk_insert_select(db, *args, **kw_args)
563 2153 aaronmk
    return run_query_into(db, query, params, into_ref, recover=recover,
564
        cacheable=cacheable)
565 2063 aaronmk
566 2066 aaronmk
default = object() # tells insert() to use the default value for a column
567
568 2063 aaronmk
def insert(db, table, row, *args, **kw_args):
569 2085 aaronmk
    '''For params, see insert_select()'''
570 1960 aaronmk
    if lists.is_seq(row): cols = None
571
    else:
572
        cols = row.keys()
573
        row = row.values()
574
    row = list(row) # ensure that "!= []" works
575
576 1961 aaronmk
    # Check for special values
577
    labels = []
578
    values = []
579
    for value in row:
580
        if value == default: labels.append('DEFAULT')
581
        else:
582
            labels.append('%s')
583
            values.append(value)
584
585
    # Build query
586 2063 aaronmk
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
587
    else: query = None
588 1554 aaronmk
589 2064 aaronmk
    return insert_select(db, table, cols, query, values, *args, **kw_args)
590 11 aaronmk
591 135 aaronmk
def last_insert_id(db):
592 1849 aaronmk
    module = util.root_module(db.db)
593 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
594
    elif module == 'MySQLdb': return db.insert_id()
595
    else: return None
596 13 aaronmk
597 1968 aaronmk
def truncate(db, table, schema='public'):
598
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
599 832 aaronmk
600
##### Database structure queries
601
602 2084 aaronmk
def pkey(db, table, recover=None, table_is_esc=False):
603 832 aaronmk
    '''Assumed to be first column in table'''
604 2120 aaronmk
    return col_names(select(db, table, limit=0, order_by=None, recover=recover,
605 2084 aaronmk
        table_is_esc=table_is_esc)).next()
606 832 aaronmk
607 853 aaronmk
def index_cols(db, table, index):
608
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
609
    automatically created. When you don't know whether something is a UNIQUE
610
    constraint or a UNIQUE index, use this function.'''
611
    check_name(table)
612
    check_name(index)
613 1909 aaronmk
    module = util.root_module(db.db)
614
    if module == 'psycopg2':
615
        return list(values(run_query(db, '''\
616 853 aaronmk
SELECT attname
617 866 aaronmk
FROM
618
(
619
        SELECT attnum, attname
620
        FROM pg_index
621
        JOIN pg_class index ON index.oid = indexrelid
622
        JOIN pg_class table_ ON table_.oid = indrelid
623
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
624
        WHERE
625
            table_.relname = %(table)s
626
            AND index.relname = %(index)s
627
    UNION
628
        SELECT attnum, attname
629
        FROM
630
        (
631
            SELECT
632
                indrelid
633
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
634
                    AS indkey
635
            FROM pg_index
636
            JOIN pg_class index ON index.oid = indexrelid
637
            JOIN pg_class table_ ON table_.oid = indrelid
638
            WHERE
639
                table_.relname = %(table)s
640
                AND index.relname = %(index)s
641
        ) s
642
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
643
) s
644 853 aaronmk
ORDER BY attnum
645
''',
646 1909 aaronmk
            {'table': table, 'index': index}, cacheable=True)))
647
    else: raise NotImplementedError("Can't list index columns for "+module+
648
        ' database')
649 853 aaronmk
650 464 aaronmk
def constraint_cols(db, table, constraint):
651
    check_name(table)
652
    check_name(constraint)
653 1849 aaronmk
    module = util.root_module(db.db)
654 464 aaronmk
    if module == 'psycopg2':
655
        return list(values(run_query(db, '''\
656
SELECT attname
657
FROM pg_constraint
658
JOIN pg_class ON pg_class.oid = conrelid
659
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
660
WHERE
661
    relname = %(table)s
662
    AND conname = %(constraint)s
663
ORDER BY attnum
664
''',
665
            {'table': table, 'constraint': constraint})))
666
    else: raise NotImplementedError("Can't list constraint columns for "+module+
667
        ' database')
668
669 2096 aaronmk
row_num_col = '_row_num'
670
671 2086 aaronmk
def add_row_num(db, table):
672 2117 aaronmk
    '''Adds a row number column to a table. Its name is in row_num_col. It will
673
    be the primary key.'''
674 2086 aaronmk
    check_name(table)
675 2096 aaronmk
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
676 2117 aaronmk
        +' serial NOT NULL PRIMARY KEY')
677 2086 aaronmk
678 1968 aaronmk
def tables(db, schema='public', table_like='%'):
679 1849 aaronmk
    module = util.root_module(db.db)
680 1968 aaronmk
    params = {'schema': schema, 'table_like': table_like}
681 832 aaronmk
    if module == 'psycopg2':
682 1968 aaronmk
        return values(run_query(db, '''\
683
SELECT tablename
684
FROM pg_tables
685
WHERE
686
    schemaname = %(schema)s
687
    AND tablename LIKE %(table_like)s
688
ORDER BY tablename
689
''',
690
            params, cacheable=True))
691
    elif module == 'MySQLdb':
692
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
693
            cacheable=True))
694 832 aaronmk
    else: raise NotImplementedError("Can't list tables for "+module+' database')
695 830 aaronmk
696 833 aaronmk
##### Database management
697
698 1968 aaronmk
def empty_db(db, schema='public', **kw_args):
699
    '''For kw_args, see tables()'''
700
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
701 833 aaronmk
702 832 aaronmk
##### Heuristic queries
703
704 2104 aaronmk
def put(db, table, row, pkey_=None, row_ct_ref=None):
705 1554 aaronmk
    '''Recovers from errors.
706 2077 aaronmk
    Only works under PostgreSQL (uses INSERT RETURNING).
707
    '''
708 2104 aaronmk
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
709
710 471 aaronmk
    try:
711 2149 aaronmk
        cur = insert(db, table, row, pkey_, recover=True)
712 1554 aaronmk
        if row_ct_ref != None and cur.rowcount >= 0:
713
            row_ct_ref[0] += cur.rowcount
714
        return value(cur)
715 471 aaronmk
    except DuplicateKeyException, e:
716 2104 aaronmk
        return value(select(db, table, [pkey_],
717 1069 aaronmk
            util.dict_subset_right_join(row, e.cols), recover=True))
718 471 aaronmk
719 473 aaronmk
def get(db, table, row, pkey, row_ct_ref=None, create=False):
720 830 aaronmk
    '''Recovers from errors'''
721
    try: return value(select(db, table, [pkey], row, 1, recover=True))
722 14 aaronmk
    except StopIteration:
723 40 aaronmk
        if not create: raise
724 471 aaronmk
        return put(db, table, row, pkey, row_ct_ref) # insert new row
725 2078 aaronmk
726 2134 aaronmk
def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
727
    row_ct_ref=None, table_is_esc=False):
728 2078 aaronmk
    '''Recovers from errors.
729
    Only works under PostgreSQL (uses INSERT RETURNING).
730 2131 aaronmk
    @param in_tables The main input table to select from, followed by a list of
731
        tables to join with it using the main input table's pkey
732 2133 aaronmk
    @return (table, col) Where the pkeys (from INSERT RETURNING) are made
733 2078 aaronmk
        available
734
    '''
735 2162 aaronmk
    temp_suffix = clean_name(out_table)
736 2158 aaronmk
        # suffix, not prefix, so main name won't be removed if name is truncated
737
    pkeys_ref = ['pkeys_'+temp_suffix]
738 2131 aaronmk
739 2132 aaronmk
    # Join together input tables
740 2131 aaronmk
    in_tables = in_tables[:] # don't modify input!
741
    in_tables0 = in_tables.pop(0) # first table is separate
742
    in_pkey = pkey(db, in_tables0, recover=True, table_is_esc=table_is_esc)
743 2178 aaronmk
    insert_joins = [in_tables0]+[(t, {in_pkey: join_using}) for t in in_tables]
744 2131 aaronmk
745 2142 aaronmk
    out_pkey = pkey(db, out_table, recover=True, table_is_esc=table_is_esc)
746
    pkeys_cols = [in_pkey, out_pkey]
747
748 2132 aaronmk
    def mk_select_(cols):
749 2178 aaronmk
        return mk_select(db, insert_joins, cols, limit=limit, start=start,
750 2134 aaronmk
            table_is_esc=table_is_esc)
751 2132 aaronmk
752 2158 aaronmk
    out_pkeys_ref = ['out_pkeys_'+temp_suffix]
753 2181 aaronmk
    def insert_(pkeys_table_exists=False):
754 2148 aaronmk
        '''Inserts and capture output pkeys.'''
755 2132 aaronmk
        cur = insert_select(db, out_table, mapping.keys(),
756 2133 aaronmk
            *mk_select_(mapping.values()), returning=out_pkey,
757 2155 aaronmk
            into_ref=out_pkeys_ref, recover=True, table_is_esc=table_is_esc)
758 2078 aaronmk
        if row_ct_ref != None and cur.rowcount >= 0:
759
            row_ct_ref[0] += cur.rowcount
760 2155 aaronmk
            add_row_num(db, out_pkeys_ref[0]) # for joining it with input pkeys
761 2086 aaronmk
762 2132 aaronmk
        # Get input pkeys corresponding to rows in insert
763 2158 aaronmk
        in_pkeys_ref = ['in_pkeys_'+temp_suffix]
764 2155 aaronmk
        run_query_into(db, *mk_select_([in_pkey]), into_ref=in_pkeys_ref)
765
        add_row_num(db, in_pkeys_ref[0]) # for joining it with output pkeys
766 2086 aaronmk
767 2155 aaronmk
        # Join together output and input pkeys
768 2181 aaronmk
        select_query = mk_select(db, [in_pkeys_ref[0],
769
            (out_pkeys_ref[0], {row_num_col: join_using})], pkeys_cols, start=0)
770
        if pkeys_table_exists:
771
            insert_select(db, pkeys_ref[0], pkeys_cols, *select_query)
772
        else: run_query_into(db, *select_query, into_ref=pkeys_ref)
773 2132 aaronmk
774 2142 aaronmk
    # Do inserts and selects
775 2148 aaronmk
    try: insert_()
776 2142 aaronmk
    except DuplicateKeyException, e:
777
        join_cols = util.dict_subset_right_join(mapping, e.cols)
778 2178 aaronmk
        select_joins = insert_joins + [(out_table, join_cols)]
779
780
        # Get pkeys of already existing rows
781
        run_query_into(db, *mk_select(db, select_joins, pkeys_cols, start=0,
782 2154 aaronmk
            table_is_esc=table_is_esc), into_ref=pkeys_ref, recover=True)
783 2142 aaronmk
784 2154 aaronmk
    return (pkeys_ref[0], out_pkey)
785 2115 aaronmk
786
##### Data cleanup
787
788
def cleanup_table(db, table, cols, table_is_esc=False):
789
    def esc_name_(name): return esc_name(db, name)
790
791
    if not table_is_esc: check_name(table)
792
    cols = map(esc_name_, cols)
793
794
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
795
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
796
            for col in cols))),
797
        dict(null0='', null1=r'\N'))