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