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