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