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 2345 aaronmk
        DbException.__init__(self, 'Violated '+name+ ' constraint on columns: '
52 2307 aaronmk
            +(', '.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 2347 aaronmk
    def esc_name(self, name): return esc_name(self, name) # calls global func
272
273 1894 aaronmk
    def run_query(self, query, params=None, cacheable=False):
274 2148 aaronmk
        '''Translates known DB errors to typed exceptions:
275
        See self.DbCursor.execute().'''
276 2167 aaronmk
        assert query != None
277
278 2047 aaronmk
        if not self.caching: cacheable = False
279 1903 aaronmk
        used_cache = False
280
        try:
281 1927 aaronmk
            # Get cursor
282
            if cacheable:
283
                query_lookup = _query_lookup(query, params)
284
                try:
285
                    cur = self.query_results[query_lookup]
286
                    used_cache = True
287
                except KeyError: cur = self.DbCursor(self)
288
            else: cur = self.db.cursor()
289
290
            # Run query
291 2148 aaronmk
            cur.execute(query, params)
292 1903 aaronmk
        finally:
293 2193 aaronmk
            if self.debug: # only compute msg if needed
294 1903 aaronmk
                if used_cache: cache_status = 'Cache hit'
295
                elif cacheable: cache_status = 'Cache miss'
296
                else: cache_status = 'Non-cacheable'
297 1927 aaronmk
                self.log_debug(cache_status+': '
298 2170 aaronmk
                    +strings.one_line(str(get_cur_query(cur, query, params))))
299 1903 aaronmk
300
        return cur
301 1914 aaronmk
302
    def is_cached(self, query, params=None):
303
        return _query_lookup(query, params) in self.query_results
304 2139 aaronmk
305
    def with_savepoint(self, func):
306 2171 aaronmk
        savepoint = 'level_'+str(self._savepoint)
307 2139 aaronmk
        self.run_query('SAVEPOINT '+savepoint)
308
        self._savepoint += 1
309
        try:
310
            try: return_val = func()
311
            finally:
312
                self._savepoint -= 1
313
                assert self._savepoint >= 0
314
        except:
315
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint)
316
            raise
317
        else:
318
            self.run_query('RELEASE SAVEPOINT '+savepoint)
319 2191 aaronmk
            self.do_autocommit()
320 2139 aaronmk
            return return_val
321 2191 aaronmk
322
    def do_autocommit(self):
323
        '''Autocommits if outside savepoint'''
324
        assert self._savepoint >= 0
325
        if self.autocommit and self._savepoint == 0:
326
            self.log_debug('Autocommiting')
327
            self.db.commit()
328 1849 aaronmk
329 1869 aaronmk
connect = DbConn
330
331 832 aaronmk
##### Querying
332
333 1894 aaronmk
def run_raw_query(db, *args, **kw_args):
334 2085 aaronmk
    '''For params, see DbConn.run_query()'''
335 1894 aaronmk
    return db.run_query(*args, **kw_args)
336 11 aaronmk
337 2068 aaronmk
def mogrify(db, query, params):
338
    module = util.root_module(db.db)
339
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
340
    else: raise NotImplementedError("Can't mogrify query for "+module+
341
        ' database')
342
343 832 aaronmk
##### Recoverable querying
344 15 aaronmk
345 2139 aaronmk
def with_savepoint(db, func): return db.with_savepoint(func)
346 11 aaronmk
347 1894 aaronmk
def run_query(db, query, params=None, recover=None, cacheable=False):
348 830 aaronmk
    if recover == None: recover = False
349
350 2148 aaronmk
    try:
351
        def run(): return run_raw_query(db, query, params, cacheable)
352
        if recover and not db.is_cached(query, params):
353
            return with_savepoint(db, run)
354
        else: return run() # don't need savepoint if cached
355
    except Exception, e:
356
        if not recover: raise # need savepoint to run index_cols()
357 2247 aaronmk
        msg = exc.str_(e)
358 2240 aaronmk
359 2148 aaronmk
        match = re.search(r'duplicate key value violates unique constraint '
360
            r'"((_?[^\W_]+)_[^"]+)"', msg)
361
        if match:
362
            constraint, table = match.groups()
363
            try: cols = index_cols(db, table, constraint)
364
            except NotImplementedError: raise e
365 2306 aaronmk
            else: raise DuplicateKeyException(constraint, cols, e)
366 2240 aaronmk
367 2148 aaronmk
        match = re.search(r'null value in column "(\w+)" violates not-null '
368 2239 aaronmk
            r'constraint', msg)
369 2345 aaronmk
        if match: raise NullValueException('NOT NULL', [match.group(1)], e)
370 2240 aaronmk
371 2247 aaronmk
        match = re.search(r'invalid input (?:syntax|value)\b.*: "(.+)"\n'
372 2242 aaronmk
            r'(?:(?s).*)\bfunction "(\w+)".*\bat assignment', msg)
373 2240 aaronmk
        if match:
374 2247 aaronmk
            raise FunctionValueException(match.group(2),
375
                strings.to_unicode(match.group(1)), e)
376 2240 aaronmk
377 2148 aaronmk
        match = re.search(r'relation "(\w+)" already exists', msg)
378
        if match: raise DuplicateTableException(match.group(1), e)
379 2240 aaronmk
380 2188 aaronmk
        match = re.search(r'function "(\w+)" already exists', msg)
381
        if match: raise DuplicateFunctionException(match.group(1), e)
382 2240 aaronmk
383 2148 aaronmk
        raise # no specific exception raised
384 830 aaronmk
385 832 aaronmk
##### Basic queries
386
387 2153 aaronmk
def next_version(name):
388
    '''Prepends the version # so it won't be removed if the name is truncated'''
389 2163 aaronmk
    version = 1 # first existing name was version 0
390 2153 aaronmk
    match = re.match(r'^v(\d+)_(.*)$', name)
391
    if match:
392
        version = int(match.group(1))+1
393
        name = match.group(2)
394
    return 'v'+str(version)+'_'+name
395
396 2151 aaronmk
def run_query_into(db, query, params, into_ref=None, *args, **kw_args):
397 2085 aaronmk
    '''Outputs a query to a temp table.
398
    For params, see run_query().
399
    '''
400 2151 aaronmk
    if into_ref == None: return run_query(db, query, params, *args, **kw_args)
401 2085 aaronmk
    else: # place rows in temp table
402 2153 aaronmk
        kw_args['recover'] = True
403
        while True:
404
            try:
405 2194 aaronmk
                create_query = 'CREATE'
406
                if not db.debug: create_query += ' TEMP'
407 2319 aaronmk
                into = sql_gen.as_Table(into_ref[0]).to_str(db)
408
                create_query += ' TABLE '+into+' AS '+query
409 2194 aaronmk
410
                return run_query(db, create_query, params, *args, **kw_args)
411 2153 aaronmk
                    # CREATE TABLE AS sets rowcount to # rows in query
412
            except DuplicateTableException, e:
413
                into_ref[0] = next_version(into_ref[0])
414
                # try again with next version of name
415 2085 aaronmk
416 2120 aaronmk
order_by_pkey = object() # tells mk_select() to order by the pkey
417
418 2199 aaronmk
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
419
420 2233 aaronmk
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
421 2293 aaronmk
    start=None, order_by=order_by_pkey, default_table=None):
422 1981 aaronmk
    '''
423 2121 aaronmk
    @param tables The single table to select from, or a list of tables to join
424 2280 aaronmk
        together, with tables after the first being sql_gen.Join objects
425 1981 aaronmk
    @param fields Use None to select all fields in the table
426 2218 aaronmk
    @param conds WHERE conditions: dict(compare_left_side=compare_right_side):
427
        * compare_left_side: Code|str (for col name)
428
        * compare_right_side: ValueCond|literal value
429 2199 aaronmk
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
430
        use all columns
431 2054 aaronmk
    @return tuple(query, params)
432 1981 aaronmk
    '''
433 2315 aaronmk
    # Parse tables param
434 2121 aaronmk
    if not lists.is_seq(tables): tables = [tables]
435 2141 aaronmk
    tables = list(tables) # don't modify input! (list() copies input)
436 2315 aaronmk
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
437 2121 aaronmk
438 2315 aaronmk
    # Parse other params
439 1135 aaronmk
    if conds == None: conds = {}
440 135 aaronmk
    assert limit == None or type(limit) == int
441 865 aaronmk
    assert start == None or type(start) == int
442 2315 aaronmk
    if order_by is order_by_pkey:
443
        if distinct_on != []: order_by = None
444
        else: order_by = pkey(db, table0, recover=True)
445 865 aaronmk
446 2315 aaronmk
    query = 'SELECT'
447 2056 aaronmk
448 2315 aaronmk
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
449 2056 aaronmk
450 2200 aaronmk
    # DISTINCT ON columns
451 2233 aaronmk
    if distinct_on != []:
452 2200 aaronmk
        query += ' DISTINCT'
453 2254 aaronmk
        if distinct_on is not distinct_on_all:
454 2200 aaronmk
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
455
456
    # Columns
457
    query += ' '
458 1135 aaronmk
    if fields == None: query += '*'
459 2056 aaronmk
    else: query += ', '.join(map(parse_col, fields))
460 2200 aaronmk
461
    # Main table
462 2271 aaronmk
    query += ' FROM '+table0.to_str(db)
463 865 aaronmk
464 2122 aaronmk
    # Add joins
465 2271 aaronmk
    left_table = table0
466 2263 aaronmk
    for join_ in tables:
467
        table = join_.table
468 2238 aaronmk
469 2343 aaronmk
        # Parse special values
470
        if join_.type_ is sql_gen.filter_out: # filter no match
471
            conds[sql_gen.Col(table_not_null_col(db, table), table)] = None
472
473 2277 aaronmk
        query += ' '+join_.to_str(db, left_table)
474 2122 aaronmk
475
        left_table = table
476
477 865 aaronmk
    missing = True
478 89 aaronmk
    if conds != {}:
479 2284 aaronmk
        query += ' WHERE '+(' AND '.join((sql_gen.as_ValueCond(r).to_str(db, l)
480
            for l, r in conds.iteritems())))
481 865 aaronmk
        missing = False
482 2227 aaronmk
    if order_by != None:
483 2285 aaronmk
        query += ' ORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
484 865 aaronmk
    if limit != None: query += ' LIMIT '+str(limit); missing = False
485
    if start != None:
486
        if start != 0: query += ' OFFSET '+str(start)
487
        missing = False
488
    if missing: warnings.warn(DbWarning(
489
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
490
491 2315 aaronmk
    return (query, [])
492 11 aaronmk
493 2054 aaronmk
def select(db, *args, **kw_args):
494
    '''For params, see mk_select() and run_query()'''
495
    recover = kw_args.pop('recover', None)
496
    cacheable = kw_args.pop('cacheable', True)
497
498
    query, params = mk_select(db, *args, **kw_args)
499
    return run_query(db, query, params, recover, cacheable)
500
501 2066 aaronmk
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
502 2292 aaronmk
    returning=None, embeddable=False):
503 1960 aaronmk
    '''
504
    @param returning str|None An inserted column (such as pkey) to return
505 2070 aaronmk
    @param embeddable Whether the query should be embeddable as a nested SELECT.
506 2073 aaronmk
        Warning: If you set this and cacheable=True when the query is run, the
507
        query will be fully cached, not just if it raises an exception.
508 1960 aaronmk
    '''
509 2328 aaronmk
    table = sql_gen.as_Table(table)
510 2318 aaronmk
    if cols == []: cols = None # no cols (all defaults) = unknown col names
511
    if cols != None: cols = [sql_gen.as_Col(v).to_str(db) for v in cols]
512 2063 aaronmk
    if select_query == None: select_query = 'DEFAULT VALUES'
513 2327 aaronmk
    if returning != None: returning = sql_gen.as_Col(returning, table)
514 2063 aaronmk
515
    # Build query
516 2328 aaronmk
    query = 'INSERT INTO '+table.to_str(db)
517 2318 aaronmk
    if cols != None: query += ' ('+', '.join(cols)+')'
518 2063 aaronmk
    query += ' '+select_query
519
520
    if returning != None:
521 2327 aaronmk
        returning_name = copy.copy(returning)
522
        returning_name.table = None
523
        returning_name = returning_name.to_str(db)
524
        query += ' RETURNING '+returning_name
525 2063 aaronmk
526 2070 aaronmk
    if embeddable:
527 2327 aaronmk
        assert returning != None
528
529 2070 aaronmk
        # Create function
530 2330 aaronmk
        function_name = '_'.join(['insert', table.name] + cols)
531 2327 aaronmk
        return_type = 'SETOF '+returning.to_str(db)+'%TYPE'
532 2189 aaronmk
        while True:
533
            try:
534 2327 aaronmk
                func_schema = None
535
                if not db.debug: func_schema = 'pg_temp'
536
                function = sql_gen.Table(function_name, func_schema).to_str(db)
537 2194 aaronmk
538 2189 aaronmk
                function_query = '''\
539
CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
540 2070 aaronmk
    LANGUAGE sql
541
    AS $$'''+mogrify(db, query, params)+''';$$;
542
'''
543 2189 aaronmk
                run_query(db, function_query, recover=True, cacheable=True)
544
                break # this version was successful
545
            except DuplicateFunctionException, e:
546
                function_name = next_version(function_name)
547
                # try again with next version of name
548 2070 aaronmk
549 2337 aaronmk
        # Return query that uses function
550
        func_table = sql_gen.NamedTable('f', sql_gen.CustomCode(function+'()'),
551
            [returning_name]) # AS clause requires function alias
552
        return mk_select(db, func_table, start=0, order_by=None)
553 2070 aaronmk
554 2066 aaronmk
    return (query, params)
555
556
def insert_select(db, *args, **kw_args):
557 2085 aaronmk
    '''For params, see mk_insert_select() and run_query_into()
558 2152 aaronmk
    @param into_ref List with name of temp table to place RETURNING values in
559 2072 aaronmk
    '''
560 2151 aaronmk
    into_ref = kw_args.pop('into_ref', None)
561
    if into_ref != None: kw_args['embeddable'] = True
562 2066 aaronmk
    recover = kw_args.pop('recover', None)
563
    cacheable = kw_args.pop('cacheable', True)
564
565
    query, params = mk_insert_select(db, *args, **kw_args)
566 2153 aaronmk
    return run_query_into(db, query, params, into_ref, recover=recover,
567
        cacheable=cacheable)
568 2063 aaronmk
569 2066 aaronmk
default = object() # tells insert() to use the default value for a column
570
571 2063 aaronmk
def insert(db, table, row, *args, **kw_args):
572 2085 aaronmk
    '''For params, see insert_select()'''
573 1960 aaronmk
    if lists.is_seq(row): cols = None
574
    else:
575
        cols = row.keys()
576
        row = row.values()
577
    row = list(row) # ensure that "!= []" works
578
579 1961 aaronmk
    # Check for special values
580
    labels = []
581
    values = []
582
    for value in row:
583 2254 aaronmk
        if value is default: labels.append('DEFAULT')
584 1961 aaronmk
        else:
585
            labels.append('%s')
586
            values.append(value)
587
588
    # Build query
589 2063 aaronmk
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
590
    else: query = None
591 1554 aaronmk
592 2064 aaronmk
    return insert_select(db, table, cols, query, values, *args, **kw_args)
593 11 aaronmk
594 135 aaronmk
def last_insert_id(db):
595 1849 aaronmk
    module = util.root_module(db.db)
596 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
597
    elif module == 'MySQLdb': return db.insert_id()
598
    else: return None
599 13 aaronmk
600 1968 aaronmk
def truncate(db, table, schema='public'):
601
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
602 832 aaronmk
603
##### Database structure queries
604
605 2339 aaronmk
def table_cols(db, table, recover=None):
606
    return list(col_names(select(db, table, limit=0, order_by=None,
607
        recover=recover)))
608
609 2291 aaronmk
def pkey(db, table, recover=None):
610 832 aaronmk
    '''Assumed to be first column in table'''
611 2339 aaronmk
    return table_cols(db, table, recover)[0]
612 832 aaronmk
613 2340 aaronmk
not_null_col = 'not_null'
614
615
def table_not_null_col(db, table, recover=None):
616
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
617
    if not_null_col in table_cols(db, table, recover): return not_null_col
618
    else: return pkey(db, table, recover)
619
620 853 aaronmk
def index_cols(db, table, index):
621
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
622
    automatically created. When you don't know whether something is a UNIQUE
623
    constraint or a UNIQUE index, use this function.'''
624 1909 aaronmk
    module = util.root_module(db.db)
625
    if module == 'psycopg2':
626
        return list(values(run_query(db, '''\
627 853 aaronmk
SELECT attname
628 866 aaronmk
FROM
629
(
630
        SELECT attnum, attname
631
        FROM pg_index
632
        JOIN pg_class index ON index.oid = indexrelid
633
        JOIN pg_class table_ ON table_.oid = indrelid
634
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
635
        WHERE
636
            table_.relname = %(table)s
637
            AND index.relname = %(index)s
638
    UNION
639
        SELECT attnum, attname
640
        FROM
641
        (
642
            SELECT
643
                indrelid
644
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
645
                    AS indkey
646
            FROM pg_index
647
            JOIN pg_class index ON index.oid = indexrelid
648
            JOIN pg_class table_ ON table_.oid = indrelid
649
            WHERE
650
                table_.relname = %(table)s
651
                AND index.relname = %(index)s
652
        ) s
653
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
654
) s
655 853 aaronmk
ORDER BY attnum
656
''',
657 1909 aaronmk
            {'table': table, 'index': index}, cacheable=True)))
658
    else: raise NotImplementedError("Can't list index columns for "+module+
659
        ' database')
660 853 aaronmk
661 464 aaronmk
def constraint_cols(db, table, constraint):
662 1849 aaronmk
    module = util.root_module(db.db)
663 464 aaronmk
    if module == 'psycopg2':
664
        return list(values(run_query(db, '''\
665
SELECT attname
666
FROM pg_constraint
667
JOIN pg_class ON pg_class.oid = conrelid
668
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
669
WHERE
670
    relname = %(table)s
671
    AND conname = %(constraint)s
672
ORDER BY attnum
673
''',
674
            {'table': table, 'constraint': constraint})))
675
    else: raise NotImplementedError("Can't list constraint columns for "+module+
676
        ' database')
677
678 2096 aaronmk
row_num_col = '_row_num'
679
680 2086 aaronmk
def add_row_num(db, table):
681 2117 aaronmk
    '''Adds a row number column to a table. Its name is in row_num_col. It will
682
    be the primary key.'''
683 2320 aaronmk
    table = sql_gen.as_Table(table).to_str(db)
684 2096 aaronmk
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
685 2117 aaronmk
        +' serial NOT NULL PRIMARY KEY')
686 2086 aaronmk
687 1968 aaronmk
def tables(db, schema='public', table_like='%'):
688 1849 aaronmk
    module = util.root_module(db.db)
689 1968 aaronmk
    params = {'schema': schema, 'table_like': table_like}
690 832 aaronmk
    if module == 'psycopg2':
691 1968 aaronmk
        return values(run_query(db, '''\
692
SELECT tablename
693
FROM pg_tables
694
WHERE
695
    schemaname = %(schema)s
696
    AND tablename LIKE %(table_like)s
697
ORDER BY tablename
698
''',
699
            params, cacheable=True))
700
    elif module == 'MySQLdb':
701
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
702
            cacheable=True))
703 832 aaronmk
    else: raise NotImplementedError("Can't list tables for "+module+' database')
704 830 aaronmk
705 833 aaronmk
##### Database management
706
707 1968 aaronmk
def empty_db(db, schema='public', **kw_args):
708
    '''For kw_args, see tables()'''
709
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
710 833 aaronmk
711 832 aaronmk
##### Heuristic queries
712
713 2104 aaronmk
def put(db, table, row, pkey_=None, row_ct_ref=None):
714 1554 aaronmk
    '''Recovers from errors.
715 2077 aaronmk
    Only works under PostgreSQL (uses INSERT RETURNING).
716
    '''
717 2104 aaronmk
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
718
719 471 aaronmk
    try:
720 2149 aaronmk
        cur = insert(db, table, row, pkey_, recover=True)
721 1554 aaronmk
        if row_ct_ref != None and cur.rowcount >= 0:
722
            row_ct_ref[0] += cur.rowcount
723
        return value(cur)
724 471 aaronmk
    except DuplicateKeyException, e:
725 2104 aaronmk
        return value(select(db, table, [pkey_],
726 1069 aaronmk
            util.dict_subset_right_join(row, e.cols), recover=True))
727 471 aaronmk
728 473 aaronmk
def get(db, table, row, pkey, row_ct_ref=None, create=False):
729 830 aaronmk
    '''Recovers from errors'''
730 2209 aaronmk
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
731 14 aaronmk
    except StopIteration:
732 40 aaronmk
        if not create: raise
733 471 aaronmk
        return put(db, table, row, pkey, row_ct_ref) # insert new row
734 2078 aaronmk
735 2134 aaronmk
def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
736 2279 aaronmk
    row_ct_ref=None):
737 2078 aaronmk
    '''Recovers from errors.
738
    Only works under PostgreSQL (uses INSERT RETURNING).
739 2131 aaronmk
    @param in_tables The main input table to select from, followed by a list of
740
        tables to join with it using the main input table's pkey
741 2312 aaronmk
    @param mapping dict(out_table_col=in_table_col, ...)
742
        * out_table_col: sql_gen.Col|str
743 2323 aaronmk
        * in_table_col: sql_gen.Col Wrap literal values in a sql_gen.NamedCol
744 2312 aaronmk
    @return sql_gen.Col Where the output pkeys are made available
745 2078 aaronmk
    '''
746 2329 aaronmk
    out_table = sql_gen.as_Table(out_table)
747 2312 aaronmk
    for in_table_col in mapping.itervalues():
748
        assert isinstance(in_table_col, sql_gen.Col)
749
750 2329 aaronmk
    temp_prefix = out_table.name
751 2251 aaronmk
    pkeys_ref = [temp_prefix+'_pkeys']
752 2131 aaronmk
753 2132 aaronmk
    # Join together input tables
754 2131 aaronmk
    in_tables = in_tables[:] # don't modify input!
755
    in_tables0 = in_tables.pop(0) # first table is separate
756 2279 aaronmk
    in_pkey = pkey(db, in_tables0, recover=True)
757 2285 aaronmk
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
758 2352 aaronmk
    input_joins = [in_tables0]+[sql_gen.Join(v, {in_pkey: sql_gen.join_same})
759 2278 aaronmk
        for v in in_tables]
760 2131 aaronmk
761 2279 aaronmk
    out_pkey = pkey(db, out_table, recover=True)
762 2285 aaronmk
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
763 2142 aaronmk
764 2236 aaronmk
    pkeys = [in_pkey, out_pkey]
765
    pkeys_cols = [in_pkey_col, out_pkey_col]
766
767 2201 aaronmk
    pkeys_table_exists_ref = [False]
768
    def run_query_into_pkeys(query, params):
769
        if pkeys_table_exists_ref[0]:
770 2236 aaronmk
            insert_select(db, pkeys_ref[0], pkeys, query, params)
771 2201 aaronmk
        else:
772
            run_query_into(db, query, params, into_ref=pkeys_ref)
773
            pkeys_table_exists_ref[0] = True
774
775 2208 aaronmk
    conds = {}
776 2233 aaronmk
    distinct_on = []
777 2325 aaronmk
    def mk_main_select(joins, cols):
778
        return mk_select(db, joins, cols, conds, distinct_on, limit=limit,
779
            start=start, default_table=in_tables0)
780 2132 aaronmk
781 2309 aaronmk
    def log_exc(e):
782
        db.log_debug('Caught exception: '+exc.str_(e, first_line_only=True))
783 2245 aaronmk
    def log_ignore(in_col, value):
784 2326 aaronmk
        db.log_debug('Ignoring rows with '+str(in_col)+' = '+value)
785 2245 aaronmk
786 2206 aaronmk
    # Do inserts and selects
787 2257 aaronmk
    join_cols = {}
788 2251 aaronmk
    out_pkeys_ref = [temp_prefix+'_out_pkeys']
789 2252 aaronmk
    in_pkeys_ref = [temp_prefix+'_in_pkeys']
790 2206 aaronmk
    while True:
791 2303 aaronmk
        has_joins = join_cols != {}
792
793 2305 aaronmk
        # Prepare to insert new rows
794 2325 aaronmk
        insert_joins = input_joins[:] # don't modify original!
795
        insert_args = dict(recover=True)
796 2303 aaronmk
        if has_joins:
797 2317 aaronmk
            distinct_on = [v.to_Col() for v in join_cols.values()]
798 2325 aaronmk
            insert_joins.append(sql_gen.Join(out_table, join_cols,
799
                sql_gen.filter_out))
800
        else:
801
            insert_args.update(dict(returning=out_pkey, into_ref=out_pkeys_ref))
802 2303 aaronmk
803 2305 aaronmk
        db.log_debug('Inserting new rows')
804 2206 aaronmk
        try:
805
            cur = insert_select(db, out_table, mapping.keys(),
806 2325 aaronmk
                *mk_main_select(insert_joins, mapping.values()), **insert_args)
807 2357 aaronmk
            break # insert successful
808 2206 aaronmk
        except DuplicateKeyException, e:
809 2309 aaronmk
            log_exc(e)
810
811 2258 aaronmk
            old_join_cols = join_cols.copy()
812 2334 aaronmk
            join_cols.update(util.dict_subset(mapping, e.cols))
813 2245 aaronmk
            db.log_debug('Ignoring existing rows, comparing on '+str(join_cols))
814 2258 aaronmk
            assert join_cols != old_join_cols # avoid infinite loops
815 2309 aaronmk
816 2206 aaronmk
            # rerun loop with additional constraints
817 2230 aaronmk
        except NullValueException, e:
818 2309 aaronmk
            log_exc(e)
819
820 2230 aaronmk
            out_col, = e.cols
821
            try: in_col = mapping[out_col]
822 2356 aaronmk
            except KeyError:
823
                db.log_debug('Missing mapping for NOT NULL '+out_col)
824
                limit = 0 # just create an empty pkeys table
825 2230 aaronmk
            else:
826 2245 aaronmk
                log_ignore(in_col, 'NULL')
827 2230 aaronmk
                conds[in_col] = sql_gen.CompareCond(None, '!=')
828 2357 aaronmk
829
            # rerun loop with additional constraints
830 2243 aaronmk
        except FunctionValueException, e:
831 2309 aaronmk
            log_exc(e)
832
833 2344 aaronmk
            assert e.name == out_table.name
834 2243 aaronmk
            out_col = 'value' # assume function param was named "value"
835
            in_col = mapping[out_col]
836 2245 aaronmk
            value = e.value
837
            log_ignore(in_col, value)
838
            conds[in_col] = sql_gen.CompareCond(value, '!=')
839 2357 aaronmk
840 2243 aaronmk
            # rerun loop with additional constraints
841 2132 aaronmk
842 2357 aaronmk
    if row_ct_ref != None and cur.rowcount >= 0:
843
        row_ct_ref[0] += cur.rowcount
844
845
    if has_joins:
846
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
847
        db.log_debug('Getting output pkeys of existing/inserted rows')
848
        run_query_into_pkeys(*mk_select(db, select_joins, pkeys_cols,
849
            start=0))
850
    else:
851
        add_row_num(db, out_pkeys_ref[0]) # for joining with input pkeys
852
853
        db.log_debug('Getting input pkeys for rows in insert')
854
        run_query_into(db, *mk_main_select(input_joins, [in_pkey]),
855
            into_ref=in_pkeys_ref)
856
        add_row_num(db, in_pkeys_ref[0]) # for joining with output pkeys
857
858
        db.log_debug('Joining together output and input pkeys')
859
        pkey_joins = [in_pkeys_ref[0], sql_gen.Join(out_pkeys_ref[0],
860
            {row_num_col: sql_gen.join_same_not_null})]
861
        run_query_into_pkeys(*mk_select(db, pkey_joins, pkeys, start=0))
862
863
    db.log_debug("Setting missing rows' pkeys to NULL")
864
    missing_rows_joins = input_joins+[sql_gen.Join(pkeys_ref[0],
865
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
866
        # must use join_same_not_null or query will take forever
867
    run_query_into_pkeys(*mk_select(db, missing_rows_joins,
868
        [in_pkey_col, sql_gen.NamedCol(out_pkey, None)], start=0))
869
870 2268 aaronmk
    return sql_gen.Col(out_pkey, pkeys_ref[0])
871 2115 aaronmk
872
##### Data cleanup
873
874 2290 aaronmk
def cleanup_table(db, table, cols):
875 2115 aaronmk
    def esc_name_(name): return esc_name(db, name)
876
877 2290 aaronmk
    table = sql_gen.as_Table(table).to_str(db)
878 2115 aaronmk
    cols = map(esc_name_, cols)
879
880
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
881
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
882
            for col in cols))),
883
        dict(null0='', null1=r'\N'))