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