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