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 2247 aaronmk
            'for name: '+str(name)+'; value: '+strings.ustr(value), cause)
46 2240 aaronmk
        self.name = name
47
        self.value = value
48
49 2306 aaronmk
class ConstraintException(DbException):
50
    def __init__(self, name, cols, cause=None):
51 2307 aaronmk
        DbException.__init__(self, 'Violated constraint '+name+ ' on columns: '
52
            +(', '.join(cols)), cause)
53 2306 aaronmk
        self.name = name
54 468 aaronmk
        self.cols = cols
55 11 aaronmk
56 2143 aaronmk
class NameException(DbException): pass
57
58 2306 aaronmk
class DuplicateKeyException(ConstraintException): pass
59 13 aaronmk
60 2306 aaronmk
class NullValueException(ConstraintException): pass
61 13 aaronmk
62 2240 aaronmk
class FunctionValueException(ExceptionWithNameValue): pass
63 2239 aaronmk
64 2143 aaronmk
class DuplicateTableException(ExceptionWithName): pass
65
66 2188 aaronmk
class DuplicateFunctionException(ExceptionWithName): pass
67
68 89 aaronmk
class EmptyRowException(DbException): pass
69
70 865 aaronmk
##### Warnings
71
72
class DbWarning(UserWarning): pass
73
74 1930 aaronmk
##### Result retrieval
75
76
def col_names(cur): return (col[0] for col in cur.description)
77
78
def rows(cur): return iter(lambda: cur.fetchone(), None)
79
80
def consume_rows(cur):
81
    '''Used to fetch all rows so result will be cached'''
82
    iters.consume_iter(rows(cur))
83
84
def next_row(cur): return rows(cur).next()
85
86
def row(cur):
87
    row_ = next_row(cur)
88
    consume_rows(cur)
89
    return row_
90
91
def next_value(cur): return next_row(cur)[0]
92
93
def value(cur): return row(cur)[0]
94
95
def values(cur): return iters.func_iter(lambda: next_value(cur))
96
97
def value_or_none(cur):
98
    try: return value(cur)
99
    except StopIteration: return None
100
101 2101 aaronmk
##### Input validation
102
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 2247 aaronmk
        msg = exc.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 2306 aaronmk
            else: raise DuplicateKeyException(constraint, 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 2306 aaronmk
        if match: raise NullValueException('', [match.group(1)], e)
368 2240 aaronmk
369 2247 aaronmk
        match = re.search(r'invalid input (?:syntax|value)\b.*: "(.+)"\n'
370 2242 aaronmk
            r'(?:(?s).*)\bfunction "(\w+)".*\bat assignment', msg)
371 2240 aaronmk
        if match:
372 2247 aaronmk
            raise FunctionValueException(match.group(2),
373
                strings.to_unicode(match.group(1)), e)
374 2240 aaronmk
375 2148 aaronmk
        match = re.search(r'relation "(\w+)" already exists', msg)
376
        if match: raise DuplicateTableException(match.group(1), e)
377 2240 aaronmk
378 2188 aaronmk
        match = re.search(r'function "(\w+)" already exists', msg)
379
        if match: raise DuplicateFunctionException(match.group(1), e)
380 2240 aaronmk
381 2148 aaronmk
        raise # no specific exception raised
382 830 aaronmk
383 832 aaronmk
##### Basic queries
384
385 2153 aaronmk
def next_version(name):
386
    '''Prepends the version # so it won't be removed if the name is truncated'''
387 2163 aaronmk
    version = 1 # first existing name was version 0
388 2153 aaronmk
    match = re.match(r'^v(\d+)_(.*)$', name)
389
    if match:
390
        version = int(match.group(1))+1
391
        name = match.group(2)
392
    return 'v'+str(version)+'_'+name
393
394 2151 aaronmk
def run_query_into(db, query, params, into_ref=None, *args, **kw_args):
395 2085 aaronmk
    '''Outputs a query to a temp table.
396
    For params, see run_query().
397
    '''
398 2151 aaronmk
    if into_ref == None: return run_query(db, query, params, *args, **kw_args)
399 2085 aaronmk
    else: # place rows in temp table
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 2319 aaronmk
                into = sql_gen.as_Table(into_ref[0]).to_str(db)
406
                create_query += ' TABLE '+into+' AS '+query
407 2194 aaronmk
408
                return run_query(db, create_query, params, *args, **kw_args)
409 2153 aaronmk
                    # CREATE TABLE AS sets rowcount to # rows in query
410
            except DuplicateTableException, e:
411
                into_ref[0] = next_version(into_ref[0])
412
                # try again with next version of name
413 2085 aaronmk
414 2120 aaronmk
order_by_pkey = object() # tells mk_select() to order by the pkey
415
416 2199 aaronmk
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
417
418 2233 aaronmk
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
419 2293 aaronmk
    start=None, order_by=order_by_pkey, default_table=None):
420 1981 aaronmk
    '''
421 2121 aaronmk
    @param tables The single table to select from, or a list of tables to join
422 2280 aaronmk
        together, with tables after the first being sql_gen.Join objects
423 1981 aaronmk
    @param fields Use None to select all fields in the table
424 2218 aaronmk
    @param conds WHERE conditions: dict(compare_left_side=compare_right_side):
425
        * compare_left_side: Code|str (for col name)
426
        * compare_right_side: ValueCond|literal value
427 2199 aaronmk
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
428
        use all columns
429 2054 aaronmk
    @return tuple(query, params)
430 1981 aaronmk
    '''
431 2315 aaronmk
    # Parse tables param
432 2121 aaronmk
    if not lists.is_seq(tables): tables = [tables]
433 2141 aaronmk
    tables = list(tables) # don't modify input! (list() copies input)
434 2315 aaronmk
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
435 2121 aaronmk
436 2315 aaronmk
    # Parse other params
437 1135 aaronmk
    if conds == None: conds = {}
438 135 aaronmk
    assert limit == None or type(limit) == int
439 865 aaronmk
    assert start == None or type(start) == int
440 2315 aaronmk
    if order_by is order_by_pkey:
441
        if distinct_on != []: order_by = None
442
        else: order_by = pkey(db, table0, recover=True)
443 865 aaronmk
444 2315 aaronmk
    query = 'SELECT'
445 2056 aaronmk
446 2315 aaronmk
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
447 2056 aaronmk
448 2200 aaronmk
    # DISTINCT ON columns
449 2233 aaronmk
    if distinct_on != []:
450 2200 aaronmk
        query += ' DISTINCT'
451 2254 aaronmk
        if distinct_on is not distinct_on_all:
452 2200 aaronmk
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
453
454
    # Columns
455
    query += ' '
456 1135 aaronmk
    if fields == None: query += '*'
457 2056 aaronmk
    else: query += ', '.join(map(parse_col, fields))
458 2200 aaronmk
459
    # Main table
460 2271 aaronmk
    query += ' FROM '+table0.to_str(db)
461 865 aaronmk
462 2122 aaronmk
    # Add joins
463 2271 aaronmk
    left_table = table0
464 2263 aaronmk
    for join_ in tables:
465
        table = join_.table
466 2238 aaronmk
467 2343 aaronmk
        # Parse special values
468
        if join_.type_ is sql_gen.filter_out: # filter no match
469
            conds[sql_gen.Col(table_not_null_col(db, table), table)] = None
470
471 2277 aaronmk
        query += ' '+join_.to_str(db, left_table)
472 2122 aaronmk
473
        left_table = table
474
475 865 aaronmk
    missing = True
476 89 aaronmk
    if conds != {}:
477 2284 aaronmk
        query += ' WHERE '+(' AND '.join((sql_gen.as_ValueCond(r).to_str(db, l)
478
            for l, r in conds.iteritems())))
479 865 aaronmk
        missing = False
480 2227 aaronmk
    if order_by != None:
481 2285 aaronmk
        query += ' ORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
482 865 aaronmk
    if limit != None: query += ' LIMIT '+str(limit); missing = False
483
    if start != None:
484
        if start != 0: query += ' OFFSET '+str(start)
485
        missing = False
486
    if missing: warnings.warn(DbWarning(
487
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
488
489 2315 aaronmk
    return (query, [])
490 11 aaronmk
491 2054 aaronmk
def select(db, *args, **kw_args):
492
    '''For params, see mk_select() and run_query()'''
493
    recover = kw_args.pop('recover', None)
494
    cacheable = kw_args.pop('cacheable', True)
495
496
    query, params = mk_select(db, *args, **kw_args)
497
    return run_query(db, query, params, recover, cacheable)
498
499 2066 aaronmk
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
500 2292 aaronmk
    returning=None, embeddable=False):
501 1960 aaronmk
    '''
502
    @param returning str|None An inserted column (such as pkey) to return
503 2070 aaronmk
    @param embeddable Whether the query should be embeddable as a nested SELECT.
504 2073 aaronmk
        Warning: If you set this and cacheable=True when the query is run, the
505
        query will be fully cached, not just if it raises an exception.
506 1960 aaronmk
    '''
507 2328 aaronmk
    table = sql_gen.as_Table(table)
508 2318 aaronmk
    if cols == []: cols = None # no cols (all defaults) = unknown col names
509
    if cols != None: cols = [sql_gen.as_Col(v).to_str(db) for v in cols]
510 2063 aaronmk
    if select_query == None: select_query = 'DEFAULT VALUES'
511 2327 aaronmk
    if returning != None: returning = sql_gen.as_Col(returning, table)
512 2063 aaronmk
513
    # Build query
514 2328 aaronmk
    query = 'INSERT INTO '+table.to_str(db)
515 2318 aaronmk
    if cols != None: query += ' ('+', '.join(cols)+')'
516 2063 aaronmk
    query += ' '+select_query
517
518
    if returning != None:
519 2327 aaronmk
        returning_name = copy.copy(returning)
520
        returning_name.table = None
521
        returning_name = returning_name.to_str(db)
522
        query += ' RETURNING '+returning_name
523 2063 aaronmk
524 2070 aaronmk
    if embeddable:
525 2327 aaronmk
        assert returning != None
526
527 2070 aaronmk
        # Create function
528 2330 aaronmk
        function_name = '_'.join(['insert', table.name] + cols)
529 2327 aaronmk
        return_type = 'SETOF '+returning.to_str(db)+'%TYPE'
530 2189 aaronmk
        while True:
531
            try:
532 2327 aaronmk
                func_schema = None
533
                if not db.debug: func_schema = 'pg_temp'
534
                function = sql_gen.Table(function_name, func_schema).to_str(db)
535 2194 aaronmk
536 2189 aaronmk
                function_query = '''\
537
CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
538 2070 aaronmk
    LANGUAGE sql
539
    AS $$'''+mogrify(db, query, params)+''';$$;
540
'''
541 2189 aaronmk
                run_query(db, function_query, recover=True, cacheable=True)
542
                break # this version was successful
543
            except DuplicateFunctionException, e:
544
                function_name = next_version(function_name)
545
                # try again with next version of name
546 2070 aaronmk
547 2337 aaronmk
        # Return query that uses function
548
        func_table = sql_gen.NamedTable('f', sql_gen.CustomCode(function+'()'),
549
            [returning_name]) # AS clause requires function alias
550
        return mk_select(db, func_table, start=0, order_by=None)
551 2070 aaronmk
552 2066 aaronmk
    return (query, params)
553
554
def insert_select(db, *args, **kw_args):
555 2085 aaronmk
    '''For params, see mk_insert_select() and run_query_into()
556 2152 aaronmk
    @param into_ref List with name of temp table to place RETURNING values in
557 2072 aaronmk
    '''
558 2151 aaronmk
    into_ref = kw_args.pop('into_ref', None)
559
    if into_ref != None: kw_args['embeddable'] = True
560 2066 aaronmk
    recover = kw_args.pop('recover', None)
561
    cacheable = kw_args.pop('cacheable', True)
562
563
    query, params = mk_insert_select(db, *args, **kw_args)
564 2153 aaronmk
    return run_query_into(db, query, params, into_ref, recover=recover,
565
        cacheable=cacheable)
566 2063 aaronmk
567 2066 aaronmk
default = object() # tells insert() to use the default value for a column
568
569 2063 aaronmk
def insert(db, table, row, *args, **kw_args):
570 2085 aaronmk
    '''For params, see insert_select()'''
571 1960 aaronmk
    if lists.is_seq(row): cols = None
572
    else:
573
        cols = row.keys()
574
        row = row.values()
575
    row = list(row) # ensure that "!= []" works
576
577 1961 aaronmk
    # Check for special values
578
    labels = []
579
    values = []
580
    for value in row:
581 2254 aaronmk
        if value is default: labels.append('DEFAULT')
582 1961 aaronmk
        else:
583
            labels.append('%s')
584
            values.append(value)
585
586
    # Build query
587 2063 aaronmk
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
588
    else: query = None
589 1554 aaronmk
590 2064 aaronmk
    return insert_select(db, table, cols, query, values, *args, **kw_args)
591 11 aaronmk
592 135 aaronmk
def last_insert_id(db):
593 1849 aaronmk
    module = util.root_module(db.db)
594 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
595
    elif module == 'MySQLdb': return db.insert_id()
596
    else: return None
597 13 aaronmk
598 1968 aaronmk
def truncate(db, table, schema='public'):
599
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
600 832 aaronmk
601
##### Database structure queries
602
603 2339 aaronmk
def table_cols(db, table, recover=None):
604
    return list(col_names(select(db, table, limit=0, order_by=None,
605
        recover=recover)))
606
607 2291 aaronmk
def pkey(db, table, recover=None):
608 832 aaronmk
    '''Assumed to be first column in table'''
609 2339 aaronmk
    return table_cols(db, table, recover)[0]
610 832 aaronmk
611 2340 aaronmk
not_null_col = 'not_null'
612
613
def table_not_null_col(db, table, recover=None):
614
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
615
    if not_null_col in table_cols(db, table, recover): return not_null_col
616
    else: return pkey(db, table, recover)
617
618 853 aaronmk
def index_cols(db, table, index):
619
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
620
    automatically created. When you don't know whether something is a UNIQUE
621
    constraint or a UNIQUE index, use this function.'''
622 1909 aaronmk
    module = util.root_module(db.db)
623
    if module == 'psycopg2':
624
        return list(values(run_query(db, '''\
625 853 aaronmk
SELECT attname
626 866 aaronmk
FROM
627
(
628
        SELECT attnum, attname
629
        FROM pg_index
630
        JOIN pg_class index ON index.oid = indexrelid
631
        JOIN pg_class table_ ON table_.oid = indrelid
632
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
633
        WHERE
634
            table_.relname = %(table)s
635
            AND index.relname = %(index)s
636
    UNION
637
        SELECT attnum, attname
638
        FROM
639
        (
640
            SELECT
641
                indrelid
642
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
643
                    AS indkey
644
            FROM pg_index
645
            JOIN pg_class index ON index.oid = indexrelid
646
            JOIN pg_class table_ ON table_.oid = indrelid
647
            WHERE
648
                table_.relname = %(table)s
649
                AND index.relname = %(index)s
650
        ) s
651
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
652
) s
653 853 aaronmk
ORDER BY attnum
654
''',
655 1909 aaronmk
            {'table': table, 'index': index}, cacheable=True)))
656
    else: raise NotImplementedError("Can't list index columns for "+module+
657
        ' database')
658 853 aaronmk
659 464 aaronmk
def constraint_cols(db, table, constraint):
660 1849 aaronmk
    module = util.root_module(db.db)
661 464 aaronmk
    if module == 'psycopg2':
662
        return list(values(run_query(db, '''\
663
SELECT attname
664
FROM pg_constraint
665
JOIN pg_class ON pg_class.oid = conrelid
666
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
667
WHERE
668
    relname = %(table)s
669
    AND conname = %(constraint)s
670
ORDER BY attnum
671
''',
672
            {'table': table, 'constraint': constraint})))
673
    else: raise NotImplementedError("Can't list constraint columns for "+module+
674
        ' database')
675
676 2096 aaronmk
row_num_col = '_row_num'
677
678 2086 aaronmk
def add_row_num(db, table):
679 2117 aaronmk
    '''Adds a row number column to a table. Its name is in row_num_col. It will
680
    be the primary key.'''
681 2320 aaronmk
    table = sql_gen.as_Table(table).to_str(db)
682 2096 aaronmk
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
683 2117 aaronmk
        +' serial NOT NULL PRIMARY KEY')
684 2086 aaronmk
685 1968 aaronmk
def tables(db, schema='public', table_like='%'):
686 1849 aaronmk
    module = util.root_module(db.db)
687 1968 aaronmk
    params = {'schema': schema, 'table_like': table_like}
688 832 aaronmk
    if module == 'psycopg2':
689 1968 aaronmk
        return values(run_query(db, '''\
690
SELECT tablename
691
FROM pg_tables
692
WHERE
693
    schemaname = %(schema)s
694
    AND tablename LIKE %(table_like)s
695
ORDER BY tablename
696
''',
697
            params, cacheable=True))
698
    elif module == 'MySQLdb':
699
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
700
            cacheable=True))
701 832 aaronmk
    else: raise NotImplementedError("Can't list tables for "+module+' database')
702 830 aaronmk
703 833 aaronmk
##### Database management
704
705 1968 aaronmk
def empty_db(db, schema='public', **kw_args):
706
    '''For kw_args, see tables()'''
707
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
708 833 aaronmk
709 832 aaronmk
##### Heuristic queries
710
711 2104 aaronmk
def put(db, table, row, pkey_=None, row_ct_ref=None):
712 1554 aaronmk
    '''Recovers from errors.
713 2077 aaronmk
    Only works under PostgreSQL (uses INSERT RETURNING).
714
    '''
715 2104 aaronmk
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
716
717 471 aaronmk
    try:
718 2149 aaronmk
        cur = insert(db, table, row, pkey_, recover=True)
719 1554 aaronmk
        if row_ct_ref != None and cur.rowcount >= 0:
720
            row_ct_ref[0] += cur.rowcount
721
        return value(cur)
722 471 aaronmk
    except DuplicateKeyException, e:
723 2104 aaronmk
        return value(select(db, table, [pkey_],
724 1069 aaronmk
            util.dict_subset_right_join(row, e.cols), recover=True))
725 471 aaronmk
726 473 aaronmk
def get(db, table, row, pkey, row_ct_ref=None, create=False):
727 830 aaronmk
    '''Recovers from errors'''
728 2209 aaronmk
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
729 14 aaronmk
    except StopIteration:
730 40 aaronmk
        if not create: raise
731 471 aaronmk
        return put(db, table, row, pkey, row_ct_ref) # insert new row
732 2078 aaronmk
733 2134 aaronmk
def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
734 2279 aaronmk
    row_ct_ref=None):
735 2078 aaronmk
    '''Recovers from errors.
736
    Only works under PostgreSQL (uses INSERT RETURNING).
737 2131 aaronmk
    @param in_tables The main input table to select from, followed by a list of
738
        tables to join with it using the main input table's pkey
739 2312 aaronmk
    @param mapping dict(out_table_col=in_table_col, ...)
740
        * out_table_col: sql_gen.Col|str
741 2323 aaronmk
        * in_table_col: sql_gen.Col Wrap literal values in a sql_gen.NamedCol
742 2312 aaronmk
    @return sql_gen.Col Where the output pkeys are made available
743 2078 aaronmk
    '''
744 2329 aaronmk
    out_table = sql_gen.as_Table(out_table)
745 2312 aaronmk
    for in_table_col in mapping.itervalues():
746
        assert isinstance(in_table_col, sql_gen.Col)
747
748 2329 aaronmk
    temp_prefix = out_table.name
749 2251 aaronmk
    pkeys_ref = [temp_prefix+'_pkeys']
750 2131 aaronmk
751 2132 aaronmk
    # Join together input tables
752 2131 aaronmk
    in_tables = in_tables[:] # don't modify input!
753
    in_tables0 = in_tables.pop(0) # first table is separate
754 2279 aaronmk
    in_pkey = pkey(db, in_tables0, recover=True)
755 2285 aaronmk
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
756 2325 aaronmk
    input_joins = [in_tables0]+[sql_gen.Join(v, {in_pkey: sql_gen.join_using})
757 2278 aaronmk
        for v in in_tables]
758 2131 aaronmk
759 2279 aaronmk
    out_pkey = pkey(db, out_table, recover=True)
760 2285 aaronmk
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
761 2142 aaronmk
762 2236 aaronmk
    pkeys = [in_pkey, out_pkey]
763
    pkeys_cols = [in_pkey_col, out_pkey_col]
764
765 2201 aaronmk
    pkeys_table_exists_ref = [False]
766
    def run_query_into_pkeys(query, params):
767
        if pkeys_table_exists_ref[0]:
768 2236 aaronmk
            insert_select(db, pkeys_ref[0], pkeys, query, params)
769 2201 aaronmk
        else:
770
            run_query_into(db, query, params, into_ref=pkeys_ref)
771
            pkeys_table_exists_ref[0] = True
772
773 2208 aaronmk
    conds = {}
774 2233 aaronmk
    distinct_on = []
775 2325 aaronmk
    def mk_main_select(joins, cols):
776
        return mk_select(db, joins, cols, conds, distinct_on, limit=limit,
777
            start=start, default_table=in_tables0)
778 2132 aaronmk
779 2309 aaronmk
    def log_exc(e):
780
        db.log_debug('Caught exception: '+exc.str_(e, first_line_only=True))
781 2245 aaronmk
    def log_ignore(in_col, value):
782 2326 aaronmk
        db.log_debug('Ignoring rows with '+str(in_col)+' = '+value)
783 2245 aaronmk
784 2206 aaronmk
    # Do inserts and selects
785 2257 aaronmk
    join_cols = {}
786 2251 aaronmk
    out_pkeys_ref = [temp_prefix+'_out_pkeys']
787 2252 aaronmk
    in_pkeys_ref = [temp_prefix+'_in_pkeys']
788 2206 aaronmk
    while True:
789 2303 aaronmk
        has_joins = join_cols != {}
790
791 2305 aaronmk
        # Prepare to insert new rows
792 2325 aaronmk
        insert_joins = input_joins[:] # don't modify original!
793
        insert_args = dict(recover=True)
794 2303 aaronmk
        if has_joins:
795 2317 aaronmk
            distinct_on = [v.to_Col() for v in join_cols.values()]
796 2325 aaronmk
            insert_joins.append(sql_gen.Join(out_table, join_cols,
797
                sql_gen.filter_out))
798
        else:
799
            insert_args.update(dict(returning=out_pkey, into_ref=out_pkeys_ref))
800 2303 aaronmk
801 2305 aaronmk
        db.log_debug('Inserting new rows')
802 2206 aaronmk
        try:
803
            cur = insert_select(db, out_table, mapping.keys(),
804 2325 aaronmk
                *mk_main_select(insert_joins, mapping.values()), **insert_args)
805 2206 aaronmk
        except DuplicateKeyException, e:
806 2309 aaronmk
            log_exc(e)
807
808 2258 aaronmk
            old_join_cols = join_cols.copy()
809 2334 aaronmk
            join_cols.update(util.dict_subset(mapping, e.cols))
810 2245 aaronmk
            db.log_debug('Ignoring existing rows, comparing on '+str(join_cols))
811 2258 aaronmk
            assert join_cols != old_join_cols # avoid infinite loops
812 2309 aaronmk
813 2206 aaronmk
            # rerun loop with additional constraints
814 2230 aaronmk
        except NullValueException, e:
815 2309 aaronmk
            log_exc(e)
816
817 2230 aaronmk
            out_col, = e.cols
818
            try: in_col = mapping[out_col]
819
            except KeyError: # no mapping for missing col, so every row invalid
820 2245 aaronmk
                db.log_debug('Missing mapping for '+out_col)
821 2230 aaronmk
                run_query_into_pkeys(*mk_select(db, insert_joins,
822 2323 aaronmk
                    [in_pkey_col, sql_gen.NamedCol(out_pkey, None)], start=0))
823 2230 aaronmk
                break
824
            else:
825 2245 aaronmk
                log_ignore(in_col, 'NULL')
826 2230 aaronmk
                conds[in_col] = sql_gen.CompareCond(None, '!=')
827
                # rerun loop with additional constraints
828 2243 aaronmk
        except FunctionValueException, e:
829 2309 aaronmk
            log_exc(e)
830
831 2243 aaronmk
            assert e.name == sql_gen.unescape_table(out_table)
832
            out_col = 'value' # assume function param was named "value"
833
            in_col = mapping[out_col]
834 2245 aaronmk
            value = e.value
835
            log_ignore(in_col, value)
836
            conds[in_col] = sql_gen.CompareCond(value, '!=')
837 2243 aaronmk
            # rerun loop with additional constraints
838 2303 aaronmk
        else:
839
            if row_ct_ref != None and cur.rowcount >= 0:
840
                row_ct_ref[0] += cur.rowcount
841
842 2305 aaronmk
            if has_joins:
843 2325 aaronmk
                select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
844 2305 aaronmk
                db.log_debug('Getting output pkeys of existing/inserted rows')
845
                run_query_into_pkeys(*mk_select(db, select_joins, pkeys_cols,
846
                    start=0))
847
            else:
848 2324 aaronmk
                add_row_num(db, out_pkeys_ref[0]) # for joining with input pkeys
849
850 2305 aaronmk
                db.log_debug('Getting input pkeys for rows in insert')
851 2325 aaronmk
                run_query_into(db, *mk_main_select(input_joins, [in_pkey]),
852 2305 aaronmk
                    into_ref=in_pkeys_ref)
853
                add_row_num(db, in_pkeys_ref[0]) # for joining with output pkeys
854
855
                db.log_debug('Joining together output and input pkeys')
856
                run_query_into_pkeys(*mk_select(db, [in_pkeys_ref[0],
857
                    sql_gen.Join(out_pkeys_ref[0],
858
                    {row_num_col: sql_gen.join_using})], pkeys, start=0))
859 2303 aaronmk
860
            break # insert successful
861 2132 aaronmk
862 2268 aaronmk
    return sql_gen.Col(out_pkey, pkeys_ref[0])
863 2115 aaronmk
864
##### Data cleanup
865
866 2290 aaronmk
def cleanup_table(db, table, cols):
867 2115 aaronmk
    def esc_name_(name): return esc_name(db, name)
868
869 2290 aaronmk
    table = sql_gen.as_Table(table).to_str(db)
870 2115 aaronmk
    cols = map(esc_name_, cols)
871
872
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
873
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
874
            for col in cols))),
875
        dict(null0='', null1=r'\N'))