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 2210 aaronmk
##### Queries
121
122
class Query:
123
    def __init__(self, str='', values=None):
124
        if str == None: str = []
125
126
        self.str = str
127
        self.values = values
128
129 1869 aaronmk
##### Database connections
130 1849 aaronmk
131 2097 aaronmk
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
132 1926 aaronmk
133 1869 aaronmk
db_engines = {
134
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
135
    'PostgreSQL': ('psycopg2', {}),
136
}
137
138
DatabaseErrors_set = set([DbException])
139
DatabaseErrors = tuple(DatabaseErrors_set)
140
141
def _add_module(module):
142
    DatabaseErrors_set.add(module.DatabaseError)
143
    global DatabaseErrors
144
    DatabaseErrors = tuple(DatabaseErrors_set)
145
146
def db_config_str(db_config):
147
    return db_config['engine']+' database '+db_config['database']
148
149 1909 aaronmk
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
150 1894 aaronmk
151 1901 aaronmk
log_debug_none = lambda msg: None
152
153 1849 aaronmk
class DbConn:
154 2190 aaronmk
    def __init__(self, db_config, serializable=True, autocommit=False,
155
        caching=True, log_debug=log_debug_none):
156 1869 aaronmk
        self.db_config = db_config
157
        self.serializable = serializable
158 2190 aaronmk
        self.autocommit = autocommit
159
        self.caching = caching
160 1901 aaronmk
        self.log_debug = log_debug
161 2193 aaronmk
        self.debug = log_debug != log_debug_none
162 1869 aaronmk
163
        self.__db = None
164 1889 aaronmk
        self.query_results = {}
165 2139 aaronmk
        self._savepoint = 0
166 1869 aaronmk
167
    def __getattr__(self, name):
168
        if name == '__dict__': raise Exception('getting __dict__')
169
        if name == 'db': return self._db()
170
        else: raise AttributeError()
171
172
    def __getstate__(self):
173
        state = copy.copy(self.__dict__) # shallow copy
174 1915 aaronmk
        state['log_debug'] = None # don't pickle the debug callback
175 1869 aaronmk
        state['_DbConn__db'] = None # don't pickle the connection
176
        return state
177
178 2165 aaronmk
    def connected(self): return self.__db != None
179
180 1869 aaronmk
    def _db(self):
181
        if self.__db == None:
182
            # Process db_config
183
            db_config = self.db_config.copy() # don't modify input!
184 2097 aaronmk
            schemas = db_config.pop('schemas', None)
185 1869 aaronmk
            module_name, mappings = db_engines[db_config.pop('engine')]
186
            module = __import__(module_name)
187
            _add_module(module)
188
            for orig, new in mappings.iteritems():
189
                try: util.rename_key(db_config, orig, new)
190
                except KeyError: pass
191
192
            # Connect
193
            self.__db = module.connect(**db_config)
194
195
            # Configure connection
196 2190 aaronmk
            if self.serializable and not self.autocommit:
197
                self.db.set_session(isolation_level='SERIALIZABLE')
198 2101 aaronmk
            if schemas != None:
199
                schemas_ = ''.join((esc_name(self, s)+', '
200
                    for s in schemas.split(',')))
201
                run_raw_query(self, "SELECT set_config('search_path', \
202
%s || current_setting('search_path'), false)", [schemas_])
203 1869 aaronmk
204
        return self.__db
205 1889 aaronmk
206 1891 aaronmk
    class DbCursor(Proxy):
207 1927 aaronmk
        def __init__(self, outer):
208 1891 aaronmk
            Proxy.__init__(self, outer.db.cursor())
209 2191 aaronmk
            self.outer = outer
210 1927 aaronmk
            self.query_results = outer.query_results
211 1894 aaronmk
            self.query_lookup = None
212 1891 aaronmk
            self.result = []
213 1889 aaronmk
214 1894 aaronmk
        def execute(self, query, params=None):
215 1930 aaronmk
            self._is_insert = query.upper().find('INSERT') >= 0
216 1894 aaronmk
            self.query_lookup = _query_lookup(query, params)
217 2148 aaronmk
            try:
218 2191 aaronmk
                try:
219
                    return_value = self.inner.execute(query, params)
220
                    self.outer.do_autocommit()
221 2148 aaronmk
                finally: self.query = get_cur_query(self.inner)
222 1904 aaronmk
            except Exception, e:
223 2170 aaronmk
                _add_cursor_info(e, self, query, params)
224 1904 aaronmk
                self.result = e # cache the exception as the result
225
                self._cache_result()
226
                raise
227 1930 aaronmk
            # Fetch all rows so result will be cached
228
            if self.rowcount == 0 and not self._is_insert: consume_rows(self)
229 1894 aaronmk
            return return_value
230
231 1891 aaronmk
        def fetchone(self):
232
            row = self.inner.fetchone()
233 1899 aaronmk
            if row != None: self.result.append(row)
234
            # otherwise, fetched all rows
235 1904 aaronmk
            else: self._cache_result()
236
            return row
237
238
        def _cache_result(self):
239 1906 aaronmk
            # For inserts, only cache exceptions since inserts are not
240
            # idempotent, but an invalid insert will always be invalid
241 1930 aaronmk
            if self.query_results != None and (not self._is_insert
242 1906 aaronmk
                or isinstance(self.result, Exception)):
243
244 1894 aaronmk
                assert self.query_lookup != None
245 1916 aaronmk
                self.query_results[self.query_lookup] = self.CacheCursor(
246
                    util.dict_subset(dicts.AttrsDictView(self),
247
                    ['query', 'result', 'rowcount', 'description']))
248 1906 aaronmk
249 1916 aaronmk
        class CacheCursor:
250
            def __init__(self, cached_result): self.__dict__ = cached_result
251
252 1927 aaronmk
            def execute(self, *args, **kw_args):
253 1916 aaronmk
                if isinstance(self.result, Exception): raise self.result
254
                # otherwise, result is a rows list
255
                self.iter = iter(self.result)
256
257
            def fetchone(self):
258
                try: return self.iter.next()
259
                except StopIteration: return None
260 1891 aaronmk
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 2210 aaronmk
class Table:
394
    def __init__(self, name, schema=None):
395
        '''
396
        @param schema str|None (for no schema)
397
        '''
398
        self.name = name
399
        self.schema = schema
400
401
    def to_str(self, db): return qual_name(db, self.schema, self.name)
402
403
class Col:
404
    def __init__(self, name, table=None):
405
        '''
406
        @param table Table|None (for no table)
407
        '''
408
        assert table == None or isinstance(table, Table)
409
410
        self.name = name
411
        self.table = table
412
413
    def to_str(self, db):
414
        str_ = ''
415
        if self.table != None: str_ += self.table.to_str(db)+'.'
416
        str_ += esc_name(db, self.name)
417
        return str_
418
419 2120 aaronmk
order_by_pkey = object() # tells mk_select() to order by the pkey
420
421 2127 aaronmk
join_using = object() # tells mk_select() to join the column with USING
422
423 2187 aaronmk
filter_out = object() # tells mk_select() to filter out rows that match the join
424 2180 aaronmk
425 2199 aaronmk
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
426
427
def mk_select(db, tables, fields=None, conds=None, distinct_on=None, limit=None,
428
    start=None, order_by=order_by_pkey, table_is_esc=False):
429 1981 aaronmk
    '''
430 2121 aaronmk
    @param tables The single table to select from, or a list of tables to join
431 2187 aaronmk
        together: [table0, (table1, joins), ...]
432
433
        joins has the format: dict(right_col=left_col, ...)
434
        * if left_col is join_using, left_col is set to right_col
435
        * if left_col is filter_out, the tables are LEFT JOINed together and the
436
          query is filtered by `right_col IS NULL` (indicating no match)
437 1981 aaronmk
    @param fields Use None to select all fields in the table
438 2199 aaronmk
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
439
        use all columns
440 1981 aaronmk
    @param table_is_esc Whether the table name has already been escaped
441 2054 aaronmk
    @return tuple(query, params)
442 1981 aaronmk
    '''
443 2060 aaronmk
    def esc_name_(name): return esc_name(db, name)
444 2058 aaronmk
445 2121 aaronmk
    if not lists.is_seq(tables): tables = [tables]
446 2141 aaronmk
    tables = list(tables) # don't modify input! (list() copies input)
447 2121 aaronmk
    table0 = tables.pop(0) # first table is separate
448
449 1135 aaronmk
    if conds == None: conds = {}
450 135 aaronmk
    assert limit == None or type(limit) == int
451 865 aaronmk
    assert start == None or type(start) == int
452 2120 aaronmk
    if order_by == order_by_pkey:
453 2121 aaronmk
        order_by = pkey(db, table0, recover=True, table_is_esc=table_is_esc)
454
    if not table_is_esc: table0 = esc_name_(table0)
455 865 aaronmk
456 2056 aaronmk
    params = []
457
458 2161 aaronmk
    def parse_col(field, default_table=None):
459 2056 aaronmk
        '''Parses fields'''
460 2176 aaronmk
        if field == None: field = (field,) # for None values, tuple is optional
461 2157 aaronmk
        is_tuple = isinstance(field, tuple)
462
        if is_tuple and len(field) == 1: # field is literal value
463
            value, = field
464 2056 aaronmk
            sql_ = '%s'
465
            params.append(value)
466 2157 aaronmk
        elif is_tuple and len(field) == 2: # field is col with table
467
            table, col = field
468
            if not table_is_esc: table = esc_name_(table)
469
            sql_ = table+'.'+esc_name_(col)
470 2161 aaronmk
        else:
471
            sql_ = esc_name_(field) # field is col name
472
            if default_table != None: sql_ = default_table+'.'+sql_
473 2056 aaronmk
        return sql_
474 11 aaronmk
    def cond(entry):
475 2056 aaronmk
        '''Parses conditions'''
476 13 aaronmk
        col, value = entry
477 2187 aaronmk
        cond_ = parse_col(col)+' '
478 11 aaronmk
        if value == None: cond_ += 'IS'
479
        else: cond_ += '='
480
        cond_ += ' %s'
481
        return cond_
482 2056 aaronmk
483 2200 aaronmk
    query = 'SELECT'
484
485
    # DISTINCT ON columns
486
    if distinct_on != None:
487
        query += ' DISTINCT'
488
        if distinct_on != distinct_on_all:
489
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
490
491
    # Columns
492
    query += ' '
493 1135 aaronmk
    if fields == None: query += '*'
494 2056 aaronmk
    else: query += ', '.join(map(parse_col, fields))
495 2200 aaronmk
496
    # Main table
497 2121 aaronmk
    query += ' FROM '+table0
498 865 aaronmk
499 2122 aaronmk
    # Add joins
500
    left_table = table0
501
    for table, joins in tables:
502
        if not table_is_esc: table = esc_name_(table)
503
504 2187 aaronmk
        left_join_ref = [False]
505
506 2122 aaronmk
        def join(entry):
507 2127 aaronmk
            '''Parses non-USING joins'''
508 2124 aaronmk
            right_col, left_col = entry
509 2173 aaronmk
510 2179 aaronmk
            # Parse special values
511 2176 aaronmk
            if left_col == None: left_col = (left_col,)
512
                # for None values, tuple is optional
513 2179 aaronmk
            elif left_col == join_using: left_col = right_col
514 2187 aaronmk
            elif left_col == filter_out:
515 2180 aaronmk
                left_col = right_col
516 2187 aaronmk
                left_join_ref[0] = True
517
                conds[(table, right_col)] = None # filter query by no match
518 2179 aaronmk
519
            # Create SQL
520 2185 aaronmk
            right_col = table+'.'+esc_name_(right_col)
521
            sql_ = right_col+' '
522 2173 aaronmk
            if isinstance(left_col, tuple) and len(left_col) == 1:
523
                # col is literal value
524
                value, = left_col
525
                if value == None: sql_ += 'IS'
526
                else: sql_ += '='
527
                sql_ += ' %s'
528
                params.append(value)
529
            else: # col is name
530
                left_col = parse_col(left_col, left_table)
531
                sql_ += ('= '+left_col+' OR ('+right_col+' IS NULL AND '
532
                    +left_col+' IS NULL)')
533
534
            return sql_
535 2122 aaronmk
536 2187 aaronmk
        # Create join condition and determine join type
537 2127 aaronmk
        if reduce(operator.and_, (v == join_using for v in joins.itervalues())):
538 2179 aaronmk
            # all cols w/ USING, so can use simpler USING syntax
539 2187 aaronmk
            join_cond = 'USING ('+(', '.join(joins.iterkeys()))+')'
540
        else: join_cond = 'ON '+(' AND '.join(map(join, joins.iteritems())))
541 2127 aaronmk
542 2187 aaronmk
        # Create join
543
        if left_join_ref[0]: query += ' LEFT'
544
        query += ' JOIN '+table+' '+join_cond
545
546 2122 aaronmk
        left_table = table
547
548 865 aaronmk
    missing = True
549 89 aaronmk
    if conds != {}:
550 2122 aaronmk
        query += ' WHERE '+(' AND '.join(map(cond, conds.iteritems())))
551 2056 aaronmk
        params += conds.values()
552 865 aaronmk
        missing = False
553 2186 aaronmk
    if order_by != None: query += ' ORDER BY '+parse_col(order_by, table0)
554 865 aaronmk
    if limit != None: query += ' LIMIT '+str(limit); missing = False
555
    if start != None:
556
        if start != 0: query += ' OFFSET '+str(start)
557
        missing = False
558
    if missing: warnings.warn(DbWarning(
559
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
560
561 2056 aaronmk
    return (query, params)
562 11 aaronmk
563 2054 aaronmk
def select(db, *args, **kw_args):
564
    '''For params, see mk_select() and run_query()'''
565
    recover = kw_args.pop('recover', None)
566
    cacheable = kw_args.pop('cacheable', True)
567
568
    query, params = mk_select(db, *args, **kw_args)
569
    return run_query(db, query, params, recover, cacheable)
570
571 2066 aaronmk
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
572 2070 aaronmk
    returning=None, embeddable=False, table_is_esc=False):
573 1960 aaronmk
    '''
574
    @param returning str|None An inserted column (such as pkey) to return
575 2070 aaronmk
    @param embeddable Whether the query should be embeddable as a nested SELECT.
576 2073 aaronmk
        Warning: If you set this and cacheable=True when the query is run, the
577
        query will be fully cached, not just if it raises an exception.
578 1960 aaronmk
    @param table_is_esc Whether the table name has already been escaped
579
    '''
580 2063 aaronmk
    if select_query == None: select_query = 'DEFAULT VALUES'
581
    if cols == []: cols = None # no cols (all defaults) = unknown col names
582 1960 aaronmk
    if not table_is_esc: check_name(table)
583 2063 aaronmk
584
    # Build query
585
    query = 'INSERT INTO '+table
586
    if cols != None:
587
        map(check_name, cols)
588
        query += ' ('+', '.join(cols)+')'
589
    query += ' '+select_query
590
591
    if returning != None:
592
        check_name(returning)
593
        query += ' RETURNING '+returning
594
595 2070 aaronmk
    if embeddable:
596
        # Create function
597 2189 aaronmk
        function_name = '_'.join(map(clean_name, ['insert', table] + cols))
598 2070 aaronmk
        return_type = 'SETOF '+table+'.'+returning+'%TYPE'
599 2189 aaronmk
        while True:
600
            try:
601 2194 aaronmk
                function = function_name
602
                if not db.debug: function = 'pg_temp.'+function
603
604 2189 aaronmk
                function_query = '''\
605
CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
606 2070 aaronmk
    LANGUAGE sql
607
    AS $$'''+mogrify(db, query, params)+''';$$;
608
'''
609 2189 aaronmk
                run_query(db, function_query, recover=True, cacheable=True)
610
                break # this version was successful
611
            except DuplicateFunctionException, e:
612
                function_name = next_version(function_name)
613
                # try again with next version of name
614 2070 aaronmk
615
        # Return query that uses function
616 2134 aaronmk
        return mk_select(db, function+'() AS f ('+returning+')', start=0,
617
            order_by=None, table_is_esc=True)# AS clause requires function alias
618 2070 aaronmk
619 2066 aaronmk
    return (query, params)
620
621
def insert_select(db, *args, **kw_args):
622 2085 aaronmk
    '''For params, see mk_insert_select() and run_query_into()
623 2152 aaronmk
    @param into_ref List with name of temp table to place RETURNING values in
624 2072 aaronmk
    '''
625 2151 aaronmk
    into_ref = kw_args.pop('into_ref', None)
626
    if into_ref != None: kw_args['embeddable'] = True
627 2066 aaronmk
    recover = kw_args.pop('recover', None)
628
    cacheable = kw_args.pop('cacheable', True)
629
630
    query, params = mk_insert_select(db, *args, **kw_args)
631 2153 aaronmk
    return run_query_into(db, query, params, into_ref, recover=recover,
632
        cacheable=cacheable)
633 2063 aaronmk
634 2066 aaronmk
default = object() # tells insert() to use the default value for a column
635
636 2063 aaronmk
def insert(db, table, row, *args, **kw_args):
637 2085 aaronmk
    '''For params, see insert_select()'''
638 1960 aaronmk
    if lists.is_seq(row): cols = None
639
    else:
640
        cols = row.keys()
641
        row = row.values()
642
    row = list(row) # ensure that "!= []" works
643
644 1961 aaronmk
    # Check for special values
645
    labels = []
646
    values = []
647
    for value in row:
648
        if value == default: labels.append('DEFAULT')
649
        else:
650
            labels.append('%s')
651
            values.append(value)
652
653
    # Build query
654 2063 aaronmk
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
655
    else: query = None
656 1554 aaronmk
657 2064 aaronmk
    return insert_select(db, table, cols, query, values, *args, **kw_args)
658 11 aaronmk
659 135 aaronmk
def last_insert_id(db):
660 1849 aaronmk
    module = util.root_module(db.db)
661 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
662
    elif module == 'MySQLdb': return db.insert_id()
663
    else: return None
664 13 aaronmk
665 1968 aaronmk
def truncate(db, table, schema='public'):
666
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
667 832 aaronmk
668
##### Database structure queries
669
670 2084 aaronmk
def pkey(db, table, recover=None, table_is_esc=False):
671 832 aaronmk
    '''Assumed to be first column in table'''
672 2120 aaronmk
    return col_names(select(db, table, limit=0, order_by=None, recover=recover,
673 2084 aaronmk
        table_is_esc=table_is_esc)).next()
674 832 aaronmk
675 853 aaronmk
def index_cols(db, table, index):
676
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
677
    automatically created. When you don't know whether something is a UNIQUE
678
    constraint or a UNIQUE index, use this function.'''
679
    check_name(table)
680
    check_name(index)
681 1909 aaronmk
    module = util.root_module(db.db)
682
    if module == 'psycopg2':
683
        return list(values(run_query(db, '''\
684 853 aaronmk
SELECT attname
685 866 aaronmk
FROM
686
(
687
        SELECT attnum, attname
688
        FROM pg_index
689
        JOIN pg_class index ON index.oid = indexrelid
690
        JOIN pg_class table_ ON table_.oid = indrelid
691
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
692
        WHERE
693
            table_.relname = %(table)s
694
            AND index.relname = %(index)s
695
    UNION
696
        SELECT attnum, attname
697
        FROM
698
        (
699
            SELECT
700
                indrelid
701
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
702
                    AS indkey
703
            FROM pg_index
704
            JOIN pg_class index ON index.oid = indexrelid
705
            JOIN pg_class table_ ON table_.oid = indrelid
706
            WHERE
707
                table_.relname = %(table)s
708
                AND index.relname = %(index)s
709
        ) s
710
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
711
) s
712 853 aaronmk
ORDER BY attnum
713
''',
714 1909 aaronmk
            {'table': table, 'index': index}, cacheable=True)))
715
    else: raise NotImplementedError("Can't list index columns for "+module+
716
        ' database')
717 853 aaronmk
718 464 aaronmk
def constraint_cols(db, table, constraint):
719
    check_name(table)
720
    check_name(constraint)
721 1849 aaronmk
    module = util.root_module(db.db)
722 464 aaronmk
    if module == 'psycopg2':
723
        return list(values(run_query(db, '''\
724
SELECT attname
725
FROM pg_constraint
726
JOIN pg_class ON pg_class.oid = conrelid
727
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
728
WHERE
729
    relname = %(table)s
730
    AND conname = %(constraint)s
731
ORDER BY attnum
732
''',
733
            {'table': table, 'constraint': constraint})))
734
    else: raise NotImplementedError("Can't list constraint columns for "+module+
735
        ' database')
736
737 2096 aaronmk
row_num_col = '_row_num'
738
739 2086 aaronmk
def add_row_num(db, table):
740 2117 aaronmk
    '''Adds a row number column to a table. Its name is in row_num_col. It will
741
    be the primary key.'''
742 2086 aaronmk
    check_name(table)
743 2096 aaronmk
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
744 2117 aaronmk
        +' serial NOT NULL PRIMARY KEY')
745 2086 aaronmk
746 1968 aaronmk
def tables(db, schema='public', table_like='%'):
747 1849 aaronmk
    module = util.root_module(db.db)
748 1968 aaronmk
    params = {'schema': schema, 'table_like': table_like}
749 832 aaronmk
    if module == 'psycopg2':
750 1968 aaronmk
        return values(run_query(db, '''\
751
SELECT tablename
752
FROM pg_tables
753
WHERE
754
    schemaname = %(schema)s
755
    AND tablename LIKE %(table_like)s
756
ORDER BY tablename
757
''',
758
            params, cacheable=True))
759
    elif module == 'MySQLdb':
760
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
761
            cacheable=True))
762 832 aaronmk
    else: raise NotImplementedError("Can't list tables for "+module+' database')
763 830 aaronmk
764 833 aaronmk
##### Database management
765
766 1968 aaronmk
def empty_db(db, schema='public', **kw_args):
767
    '''For kw_args, see tables()'''
768
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
769 833 aaronmk
770 832 aaronmk
##### Heuristic queries
771
772 2104 aaronmk
def put(db, table, row, pkey_=None, row_ct_ref=None):
773 1554 aaronmk
    '''Recovers from errors.
774 2077 aaronmk
    Only works under PostgreSQL (uses INSERT RETURNING).
775
    '''
776 2104 aaronmk
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
777
778 471 aaronmk
    try:
779 2149 aaronmk
        cur = insert(db, table, row, pkey_, recover=True)
780 1554 aaronmk
        if row_ct_ref != None and cur.rowcount >= 0:
781
            row_ct_ref[0] += cur.rowcount
782
        return value(cur)
783 471 aaronmk
    except DuplicateKeyException, e:
784 2104 aaronmk
        return value(select(db, table, [pkey_],
785 1069 aaronmk
            util.dict_subset_right_join(row, e.cols), recover=True))
786 471 aaronmk
787 473 aaronmk
def get(db, table, row, pkey, row_ct_ref=None, create=False):
788 830 aaronmk
    '''Recovers from errors'''
789 2209 aaronmk
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
790 14 aaronmk
    except StopIteration:
791 40 aaronmk
        if not create: raise
792 471 aaronmk
        return put(db, table, row, pkey, row_ct_ref) # insert new row
793 2078 aaronmk
794 2134 aaronmk
def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
795
    row_ct_ref=None, table_is_esc=False):
796 2078 aaronmk
    '''Recovers from errors.
797
    Only works under PostgreSQL (uses INSERT RETURNING).
798 2131 aaronmk
    @param in_tables The main input table to select from, followed by a list of
799
        tables to join with it using the main input table's pkey
800 2133 aaronmk
    @return (table, col) Where the pkeys (from INSERT RETURNING) are made
801 2078 aaronmk
        available
802
    '''
803 2162 aaronmk
    temp_suffix = clean_name(out_table)
804 2158 aaronmk
        # suffix, not prefix, so main name won't be removed if name is truncated
805
    pkeys_ref = ['pkeys_'+temp_suffix]
806 2131 aaronmk
807 2132 aaronmk
    # Join together input tables
808 2131 aaronmk
    in_tables = in_tables[:] # don't modify input!
809
    in_tables0 = in_tables.pop(0) # first table is separate
810
    in_pkey = pkey(db, in_tables0, recover=True, table_is_esc=table_is_esc)
811 2178 aaronmk
    insert_joins = [in_tables0]+[(t, {in_pkey: join_using}) for t in in_tables]
812 2131 aaronmk
813 2142 aaronmk
    out_pkey = pkey(db, out_table, recover=True, table_is_esc=table_is_esc)
814
    pkeys_cols = [in_pkey, out_pkey]
815
816 2201 aaronmk
    pkeys_table_exists_ref = [False]
817
    def run_query_into_pkeys(query, params):
818
        if pkeys_table_exists_ref[0]:
819
            insert_select(db, pkeys_ref[0], pkeys_cols, query, params)
820
        else:
821
            run_query_into(db, query, params, into_ref=pkeys_ref)
822
            pkeys_table_exists_ref[0] = True
823
824 2208 aaronmk
    conds = {}
825 2207 aaronmk
    distinct_on = None
826
    def mk_main_select(cols):
827 2208 aaronmk
        return mk_select(db, insert_joins, cols, conds, distinct_on,
828 2205 aaronmk
            order_by=None, limit=limit, start=start, table_is_esc=table_is_esc)
829 2132 aaronmk
830 2206 aaronmk
    # Do inserts and selects
831 2158 aaronmk
    out_pkeys_ref = ['out_pkeys_'+temp_suffix]
832 2206 aaronmk
    while True:
833
        try:
834
            cur = insert_select(db, out_table, mapping.keys(),
835 2207 aaronmk
                *mk_main_select(mapping.values()), returning=out_pkey,
836 2206 aaronmk
                into_ref=out_pkeys_ref, recover=True, table_is_esc=table_is_esc)
837
            if row_ct_ref != None and cur.rowcount >= 0:
838
                row_ct_ref[0] += cur.rowcount
839
                add_row_num(db, out_pkeys_ref[0]) # for joining with input pkeys
840
841
            # Get input pkeys corresponding to rows in insert
842
            in_pkeys_ref = ['in_pkeys_'+temp_suffix]
843 2207 aaronmk
            run_query_into(db, *mk_main_select([in_pkey]),
844
                into_ref=in_pkeys_ref)
845 2206 aaronmk
            add_row_num(db, in_pkeys_ref[0]) # for joining with output pkeys
846
847
            # Join together output and input pkeys
848
            run_query_into_pkeys(*mk_select(db, [in_pkeys_ref[0],
849
                (out_pkeys_ref[0], {row_num_col: join_using})], pkeys_cols,
850
                start=0))
851
852
            break # insert successful
853
        except DuplicateKeyException, e:
854
            join_cols = util.dict_subset_right_join(mapping, e.cols)
855
            select_joins = insert_joins + [(out_table, join_cols)]
856
857
            # Get pkeys of already existing rows
858
            run_query_into_pkeys(*mk_select(db, select_joins, pkeys_cols,
859
                order_by=None, start=0, table_is_esc=table_is_esc))
860
861
            # Save existing pkeys in temp table for joining on
862
            existing_pkeys_ref = ['existing_pkeys_'+temp_suffix]
863
            run_query_into(db, *mk_select(db, pkeys_ref[0], [in_pkey],
864
                order_by=None, start=0, table_is_esc=True),
865
                into_ref=existing_pkeys_ref)
866
                # need table_is_esc=True to make table name case-insensitive
867
868
            # rerun loop with additional constraints
869
            break # but until NullValueExceptions are handled, end loop here
870 2132 aaronmk
871 2154 aaronmk
    return (pkeys_ref[0], out_pkey)
872 2115 aaronmk
873
##### Data cleanup
874
875
def cleanup_table(db, table, cols, table_is_esc=False):
876
    def esc_name_(name): return esc_name(db, name)
877
878
    if not table_is_esc: check_name(table)
879
    cols = map(esc_name_, cols)
880
881
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
882
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
883
            for col in cols))),
884
        dict(null0='', null1=r'\N'))