Project

General

Profile

1
# Database access
2

    
3
import copy
4
import operator
5
import re
6
import warnings
7

    
8
import exc
9
import dicts
10
import iters
11
import lists
12
from Proxy import Proxy
13
import rand
14
import sql_gen
15
import strings
16
import util
17

    
18
##### Exceptions
19

    
20
def get_cur_query(cur, input_query=None):
21
    raw_query = None
22
    if hasattr(cur, 'query'): raw_query = cur.query
23
    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
24
    
25
    if raw_query != None: return raw_query
26
    else: return '[input] '+strings.ustr(input_query)
27

    
28
def _add_cursor_info(e, *args, **kw_args):
29
    '''For params, see get_cur_query()'''
30
    exc.add_msg(e, 'query: '+strings.ustr(get_cur_query(*args, **kw_args)))
31

    
32
class DbException(exc.ExceptionWithCause):
33
    def __init__(self, msg, cause=None, cur=None):
34
        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)
35
        if cur != None: _add_cursor_info(self, cur)
36

    
37
class ExceptionWithName(DbException):
38
    def __init__(self, name, cause=None):
39
        DbException.__init__(self, 'for name: '+strings.as_tt(str(name)), cause)
40
        self.name = name
41

    
42
class ExceptionWithNameValue(DbException):
43
    def __init__(self, name, value, cause=None):
44
        DbException.__init__(self, 'for name: '+strings.as_tt(str(name))
45
            +'; value: '+strings.as_tt(repr(value)), cause)
46
        self.name = name
47
        self.value = value
48

    
49
class ExceptionWithNameType(DbException):
50
    def __init__(self, type_, name, cause=None):
51
        DbException.__init__(self, 'for type: '+strings.as_tt(str(type_))
52
            +'; name: '+strings.as_tt(name), cause)
53
        self.type = type_
54
        self.name = name
55

    
56
class ConstraintException(DbException):
57
    def __init__(self, name, cols, cause=None):
58
        DbException.__init__(self, 'Violated '+strings.as_tt(name)
59
            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)
60
        self.name = name
61
        self.cols = cols
62

    
63
class MissingCastException(DbException):
64
    def __init__(self, type_, col, cause=None):
65
        DbException.__init__(self, 'Missing cast to type '+strings.as_tt(type_)
66
            +' on column: '+strings.as_tt(col), cause)
67
        self.type = type_
68
        self.col = col
69

    
70
class NameException(DbException): pass
71

    
72
class DuplicateKeyException(ConstraintException): pass
73

    
74
class NullValueException(ConstraintException): pass
75

    
76
class FunctionValueException(ExceptionWithNameValue): pass
77

    
78
class DuplicateException(ExceptionWithNameType): pass
79

    
80
class EmptyRowException(DbException): pass
81

    
82
##### Warnings
83

    
84
class DbWarning(UserWarning): pass
85

    
86
##### Result retrieval
87

    
88
def col_names(cur): return (col[0] for col in cur.description)
89

    
90
def rows(cur): return iter(lambda: cur.fetchone(), None)
91

    
92
def consume_rows(cur):
93
    '''Used to fetch all rows so result will be cached'''
94
    iters.consume_iter(rows(cur))
95

    
96
def next_row(cur): return rows(cur).next()
97

    
98
def row(cur):
99
    row_ = next_row(cur)
100
    consume_rows(cur)
101
    return row_
102

    
103
def next_value(cur): return next_row(cur)[0]
104

    
105
def value(cur): return row(cur)[0]
106

    
107
def values(cur): return iters.func_iter(lambda: next_value(cur))
108

    
109
def value_or_none(cur):
110
    try: return value(cur)
111
    except StopIteration: return None
112

    
113
##### Escaping
114

    
115
def esc_name_by_module(module, name):
116
    if module == 'psycopg2' or module == None: quote = '"'
117
    elif module == 'MySQLdb': quote = '`'
118
    else: raise NotImplementedError("Can't escape name for "+module+' database')
119
    return sql_gen.esc_name(name, quote)
120

    
121
def esc_name_by_engine(engine, name, **kw_args):
122
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
123

    
124
def esc_name(db, name, **kw_args):
125
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
126

    
127
def qual_name(db, schema, table):
128
    def esc_name_(name): return esc_name(db, name)
129
    table = esc_name_(table)
130
    if schema != None: return esc_name_(schema)+'.'+table
131
    else: return table
132

    
133
##### Database connections
134

    
135
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
136

    
137
db_engines = {
138
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
139
    'PostgreSQL': ('psycopg2', {}),
140
}
141

    
142
DatabaseErrors_set = set([DbException])
143
DatabaseErrors = tuple(DatabaseErrors_set)
144

    
145
def _add_module(module):
146
    DatabaseErrors_set.add(module.DatabaseError)
147
    global DatabaseErrors
148
    DatabaseErrors = tuple(DatabaseErrors_set)
149

    
150
def db_config_str(db_config):
151
    return db_config['engine']+' database '+db_config['database']
152

    
153
log_debug_none = lambda msg, level=2: None
154

    
155
class DbConn:
156
    def __init__(self, db_config, autocommit=True, caching=True,
157
        log_debug=log_debug_none, debug_temp=False):
158
        '''
159
        @param debug_temp Whether temporary objects should instead be permanent.
160
            This assists in debugging the internal objects used by the program.
161
        '''
162
        self.db_config = db_config
163
        self.autocommit = autocommit
164
        self.caching = caching
165
        self.log_debug = log_debug
166
        self.debug = log_debug != log_debug_none
167
        self.debug_temp = debug_temp
168
        
169
        self.__db = None
170
        self.query_results = {}
171
        self._savepoint = 0
172
        self._notices_seen = set()
173
    
174
    def __getattr__(self, name):
175
        if name == '__dict__': raise Exception('getting __dict__')
176
        if name == 'db': return self._db()
177
        else: raise AttributeError()
178
    
179
    def __getstate__(self):
180
        state = copy.copy(self.__dict__) # shallow copy
181
        state['log_debug'] = None # don't pickle the debug callback
182
        state['_DbConn__db'] = None # don't pickle the connection
183
        return state
184
    
185
    def connected(self): return self.__db != None
186
    
187
    def _db(self):
188
        if self.__db == None:
189
            # Process db_config
190
            db_config = self.db_config.copy() # don't modify input!
191
            schemas = db_config.pop('schemas', None)
192
            module_name, mappings = db_engines[db_config.pop('engine')]
193
            module = __import__(module_name)
194
            _add_module(module)
195
            for orig, new in mappings.iteritems():
196
                try: util.rename_key(db_config, orig, new)
197
                except KeyError: pass
198
            
199
            # Connect
200
            self.__db = module.connect(**db_config)
201
            
202
            # Configure connection
203
            if hasattr(self.db, 'set_isolation_level'):
204
                import psycopg2.extensions
205
                self.db.set_isolation_level(
206
                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
207
            if schemas != None:
208
                search_path = [self.esc_name(s) for s in schemas.split(',')]
209
                search_path.append(value(run_query(self, 'SHOW search_path',
210
                    log_level=4)))
211
                run_query(self, 'SET search_path TO '+(','.join(search_path)),
212
                    log_level=3)
213
        
214
        return self.__db
215
    
216
    class DbCursor(Proxy):
217
        def __init__(self, outer):
218
            Proxy.__init__(self, outer.db.cursor())
219
            self.outer = outer
220
            self.query_results = outer.query_results
221
            self.query_lookup = None
222
            self.result = []
223
        
224
        def execute(self, query):
225
            self._is_insert = query.startswith('INSERT')
226
            self.query_lookup = query
227
            try:
228
                try:
229
                    cur = self.inner.execute(query)
230
                    self.outer.do_autocommit()
231
                finally: self.query = get_cur_query(self.inner, query)
232
            except Exception, e:
233
                _add_cursor_info(e, self, query)
234
                self.result = e # cache the exception as the result
235
                self._cache_result()
236
                raise
237
            
238
            # Always cache certain queries
239
            if query.startswith('CREATE') or query.startswith('ALTER'):
240
                # structural changes
241
                if query.find('ADD COLUMN') < 0: # rest of query must be unique
242
                    self._cache_result()
243
            elif self.rowcount == 0 and query.startswith('SELECT'): # empty
244
                consume_rows(self) # fetch all rows so result will be cached
245
            
246
            return cur
247
        
248
        def fetchone(self):
249
            row = self.inner.fetchone()
250
            if row != None: self.result.append(row)
251
            # otherwise, fetched all rows
252
            else: self._cache_result()
253
            return row
254
        
255
        def _cache_result(self):
256
            # For inserts that return a result set, don't cache result set since
257
            # inserts are not idempotent. Other non-SELECT queries don't have
258
            # their result set read, so only exceptions will be cached (an
259
            # invalid query will always be invalid).
260
            if self.query_results != None and (not self._is_insert
261
                or isinstance(self.result, Exception)):
262
                
263
                assert self.query_lookup != None
264
                self.query_results[self.query_lookup] = self.CacheCursor(
265
                    util.dict_subset(dicts.AttrsDictView(self),
266
                    ['query', 'result', 'rowcount', 'description']))
267
        
268
        class CacheCursor:
269
            def __init__(self, cached_result): self.__dict__ = cached_result
270
            
271
            def execute(self, *args, **kw_args):
272
                if isinstance(self.result, Exception): raise self.result
273
                # otherwise, result is a rows list
274
                self.iter = iter(self.result)
275
            
276
            def fetchone(self):
277
                try: return self.iter.next()
278
                except StopIteration: return None
279
    
280
    def esc_value(self, value):
281
        try: str_ = self.mogrify('%s', [value])
282
        except NotImplementedError, e:
283
            module = util.root_module(self.db)
284
            if module == 'MySQLdb':
285
                import _mysql
286
                str_ = _mysql.escape_string(value)
287
            else: raise e
288
        return strings.to_unicode(str_)
289
    
290
    def esc_name(self, name): return esc_name(self, name) # calls global func
291
    
292
    def std_code(self, str_):
293
        '''Standardizes SQL code.
294
        * Ensures that string literals are prefixed by `E`
295
        '''
296
        if str_.startswith("'"): str_ = 'E'+str_
297
        return str_
298
    
299
    def can_mogrify(self):
300
        module = util.root_module(self.db)
301
        return module == 'psycopg2'
302
    
303
    def mogrify(self, query, params=None):
304
        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
305
        else: raise NotImplementedError("Can't mogrify query")
306
    
307
    def print_notices(self):
308
        if hasattr(self.db, 'notices'):
309
            for msg in self.db.notices:
310
                if msg not in self._notices_seen:
311
                    self._notices_seen.add(msg)
312
                    self.log_debug(msg, level=2)
313
    
314
    def run_query(self, query, cacheable=False, log_level=2,
315
        debug_msg_ref=None):
316
        '''
317
        @param log_ignore_excs The log_level will be increased by 2 if the query
318
            throws one of these exceptions.
319
        @param debug_msg_ref If specified, the log message will be returned in
320
            this instead of being output. This allows you to filter log messages
321
            depending on the result of the query.
322
        '''
323
        assert query != None
324
        
325
        if not self.caching: cacheable = False
326
        used_cache = False
327
        
328
        def log_msg(query):
329
            if used_cache: cache_status = 'cache hit'
330
            elif cacheable: cache_status = 'cache miss'
331
            else: cache_status = 'non-cacheable'
332
            return 'DB query: '+cache_status+':\n'+strings.as_code(query, 'SQL')
333
        
334
        try:
335
            # Get cursor
336
            if cacheable:
337
                try:
338
                    cur = self.query_results[query]
339
                    used_cache = True
340
                except KeyError: cur = self.DbCursor(self)
341
            else: cur = self.db.cursor()
342
            
343
            # Log query
344
            if self.debug and debug_msg_ref == None: # log before running
345
                self.log_debug(log_msg(query), log_level)
346
            
347
            # Run query
348
            cur.execute(query)
349
        finally:
350
            self.print_notices()
351
            if self.debug and debug_msg_ref != None: # return after running
352
                debug_msg_ref[0] = log_msg(str(get_cur_query(cur, query)))
353
        
354
        return cur
355
    
356
    def is_cached(self, query): return query in self.query_results
357
    
358
    def with_autocommit(self, func):
359
        import psycopg2.extensions
360
        
361
        prev_isolation_level = self.db.isolation_level
362
        self.db.set_isolation_level(
363
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
364
        try: return func()
365
        finally: self.db.set_isolation_level(prev_isolation_level)
366
    
367
    def with_savepoint(self, func):
368
        savepoint = 'level_'+str(self._savepoint)
369
        self.run_query('SAVEPOINT '+savepoint, log_level=4)
370
        self._savepoint += 1
371
        try: return func()
372
        except:
373
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint, log_level=4)
374
            raise
375
        finally:
376
            # Always release savepoint, because after ROLLBACK TO SAVEPOINT,
377
            # "The savepoint remains valid and can be rolled back to again"
378
            # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).
379
            self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
380
            
381
            self._savepoint -= 1
382
            assert self._savepoint >= 0
383
            
384
            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT
385
    
386
    def do_autocommit(self):
387
        '''Autocommits if outside savepoint'''
388
        assert self._savepoint >= 0
389
        if self.autocommit and self._savepoint == 0:
390
            self.log_debug('Autocommitting', level=4)
391
            self.db.commit()
392
    
393
    def col_info(self, col):
394
        table = sql_gen.Table('columns', 'information_schema')
395
        cols = ['data_type', 'column_default',
396
            cast(self, 'boolean', 'is_nullable')]
397
        
398
        conds = [('table_name', col.table.name), ('column_name', col.name)]
399
        schema = col.table.schema
400
        if schema != None: conds.append(('table_schema', schema))
401
        
402
        type_, default, nullable = row(select(self, table, cols, conds,
403
            order_by='table_schema', limit=1, log_level=4))
404
            # TODO: order_by search_path schema order
405
        default = sql_gen.as_Code(default, self)
406
        
407
        return sql_gen.TypedCol(col.name, type_, default, nullable)
408
    
409
    def TempFunction(self, name):
410
        if self.debug_temp: schema = None
411
        else: schema = 'pg_temp'
412
        return sql_gen.Function(name, schema)
413

    
414
connect = DbConn
415

    
416
##### Recoverable querying
417

    
418
def with_savepoint(db, func): return db.with_savepoint(func)
419

    
420
def run_query(db, query, recover=None, cacheable=False, log_level=2,
421
    log_ignore_excs=None, **kw_args):
422
    '''For params, see DbConn.run_query()'''
423
    if recover == None: recover = False
424
    if log_ignore_excs == None: log_ignore_excs = ()
425
    log_ignore_excs = tuple(log_ignore_excs)
426
    
427
    debug_msg_ref = None # usually, db.run_query() logs query before running it
428
    # But if filtering with log_ignore_excs, wait until after exception parsing
429
    if log_ignore_excs != () or not db.can_mogrify(): debug_msg_ref = [None]
430
    
431
    try:
432
        try:
433
            def run(): return db.run_query(query, cacheable, log_level,
434
                debug_msg_ref, **kw_args)
435
            if recover and not db.is_cached(query):
436
                return with_savepoint(db, run)
437
            else: return run() # don't need savepoint if cached
438
        except Exception, e:
439
            if not recover: raise # need savepoint to run index_cols()
440
            msg = exc.str_(e)
441
            
442
            match = re.search(r'duplicate key value violates unique constraint '
443
                r'"((_?[^\W_]+)_.+?)"', msg)
444
            if match:
445
                constraint, table = match.groups()
446
                try: cols = index_cols(db, table, constraint)
447
                except NotImplementedError: raise e
448
                else: raise DuplicateKeyException(constraint, cols, e)
449
            
450
            match = re.search(r'null value in column "(.+?)" violates not-null'
451
                r' constraint', msg)
452
            if match: raise NullValueException('NOT NULL', [match.group(1)], e)
453
            
454
            match = re.search(r'\b(?:invalid input (?:syntax|value)\b.*?'
455
                r'|date/time field value out of range): "(.+?)"\n'
456
                r'(?:(?s).*?)\bfunction "(.+?)"', msg)
457
            if match:
458
                value, name = match.groups()
459
                raise FunctionValueException(name, strings.to_unicode(value), e)
460
            
461
            match = re.search(r'column "(.+?)" is of type (.+?) but expression '
462
                r'is of type', msg)
463
            if match:
464
                col, type_ = match.groups()
465
                raise MissingCastException(type_, col, e)
466
            
467
            match = re.search(r'\b(\S+) "(.+?)".*? already exists', msg)
468
            if match:
469
                type_, name = match.groups()
470
                raise DuplicateException(type_, name, e)
471
            
472
            raise # no specific exception raised
473
    except log_ignore_excs:
474
        log_level += 2
475
        raise
476
    finally:
477
        if debug_msg_ref != None and debug_msg_ref[0] != None:
478
            db.log_debug(debug_msg_ref[0], log_level)
479

    
480
##### Basic queries
481

    
482
def next_version(name):
483
    version = 1 # first existing name was version 0
484
    match = re.match(r'^(.*)#(\d+)$', name)
485
    if match:
486
        name, version = match.groups()
487
        version = int(version)+1
488
    return sql_gen.concat(name, '#'+str(version))
489

    
490
def lock_table(db, table, mode):
491
    table = sql_gen.as_Table(table)
492
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
493

    
494
def run_query_into(db, query, into=None, add_indexes_=False, **kw_args):
495
    '''Outputs a query to a temp table.
496
    For params, see run_query().
497
    '''
498
    if into == None: return run_query(db, query, **kw_args)
499
    
500
    assert isinstance(into, sql_gen.Table)
501
    
502
    into.is_temp = True
503
    # "temporary tables cannot specify a schema name", so remove schema
504
    into.schema = None
505
    
506
    kw_args['recover'] = True
507
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
508
    
509
    temp = not db.debug_temp # tables are permanent in debug_temp mode
510
    
511
    # Create table
512
    while True:
513
        create_query = 'CREATE'
514
        if temp: create_query += ' TEMP'
515
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
516
        
517
        try:
518
            cur = run_query(db, create_query, **kw_args)
519
                # CREATE TABLE AS sets rowcount to # rows in query
520
            break
521
        except DuplicateException, e:
522
            into.name = next_version(into.name)
523
            # try again with next version of name
524
    
525
    if add_indexes_: add_indexes(db, into)
526
    
527
    return cur
528

    
529
order_by_pkey = object() # tells mk_select() to order by the pkey
530

    
531
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
532

    
533
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
534
    start=None, order_by=order_by_pkey, default_table=None):
535
    '''
536
    @param tables The single table to select from, or a list of tables to join
537
        together, with tables after the first being sql_gen.Join objects
538
    @param fields Use None to select all fields in the table
539
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
540
        * container can be any iterable type
541
        * compare_left_side: sql_gen.Code|str (for col name)
542
        * compare_right_side: sql_gen.ValueCond|literal value
543
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
544
        use all columns
545
    @return query
546
    '''
547
    # Parse tables param
548
    tables = lists.mk_seq(tables)
549
    tables = list(tables) # don't modify input! (list() copies input)
550
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
551
    
552
    # Parse other params
553
    if conds == None: conds = []
554
    elif dicts.is_dict(conds): conds = conds.items()
555
    conds = list(conds) # don't modify input! (list() copies input)
556
    assert limit == None or type(limit) == int
557
    assert start == None or type(start) == int
558
    if order_by is order_by_pkey:
559
        if distinct_on != []: order_by = None
560
        else: order_by = pkey(db, table0, recover=True)
561
    
562
    query = 'SELECT'
563
    
564
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
565
    
566
    # DISTINCT ON columns
567
    if distinct_on != []:
568
        query += '\nDISTINCT'
569
        if distinct_on is not distinct_on_all:
570
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
571
    
572
    # Columns
573
    query += '\n'
574
    if fields == None: query += '*'
575
    else:
576
        assert fields != []
577
        query += '\n, '.join(map(parse_col, fields))
578
    
579
    # Main table
580
    query += '\nFROM '+table0.to_str(db)
581
    
582
    # Add joins
583
    left_table = table0
584
    for join_ in tables:
585
        table = join_.table
586
        
587
        # Parse special values
588
        if join_.type_ is sql_gen.filter_out: # filter no match
589
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
590
                sql_gen.CompareCond(None, '~=')))
591
        
592
        query += '\n'+join_.to_str(db, left_table)
593
        
594
        left_table = table
595
    
596
    missing = True
597
    if conds != []:
598
        if len(conds) == 1: whitespace = ' '
599
        else: whitespace = '\n'
600
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
601
            .to_str(db) for l, r in conds], 'WHERE')
602
        missing = False
603
    if order_by != None:
604
        query += '\nORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
605
    if limit != None: query += '\nLIMIT '+str(limit); missing = False
606
    if start != None:
607
        if start != 0: query += '\nOFFSET '+str(start)
608
        missing = False
609
    if missing: warnings.warn(DbWarning(
610
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
611
    
612
    return query
613

    
614
def select(db, *args, **kw_args):
615
    '''For params, see mk_select() and run_query()'''
616
    recover = kw_args.pop('recover', None)
617
    cacheable = kw_args.pop('cacheable', True)
618
    log_level = kw_args.pop('log_level', 2)
619
    
620
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
621
        log_level=log_level)
622

    
623
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
624
    embeddable=False, ignore=False):
625
    '''
626
    @param returning str|None An inserted column (such as pkey) to return
627
    @param embeddable Whether the query should be embeddable as a nested SELECT.
628
        Warning: If you set this and cacheable=True when the query is run, the
629
        query will be fully cached, not just if it raises an exception.
630
    @param ignore Whether to ignore duplicate keys.
631
    '''
632
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
633
    if cols == []: cols = None # no cols (all defaults) = unknown col names
634
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
635
    if select_query == None: select_query = 'DEFAULT VALUES'
636
    if returning != None: returning = sql_gen.as_Col(returning, table)
637
    
638
    first_line = 'INSERT INTO '+table.to_str(db)
639
    
640
    def mk_insert(select_query):
641
        query = first_line
642
        if cols != None:
643
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
644
        query += '\n'+select_query
645
        
646
        if returning != None:
647
            returning_name_col = sql_gen.to_name_only_col(returning)
648
            query += '\nRETURNING '+returning_name_col.to_str(db)
649
        
650
        return query
651
    
652
    return_type = 'unknown'
653
    if returning != None: return_type = returning.to_str(db)+'%TYPE'
654
    
655
    lang = 'sql'
656
    if ignore:
657
        assert cols != None
658
        # Always return something to set the correct rowcount
659
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
660
        
661
        embeddable = True # must use function
662
        lang = 'plpgsql'
663
        row = [sql_gen.Col(c.name, 'row') for c in cols]
664
        
665
        query = '''\
666
DECLARE
667
    row '''+table.to_str(db)+'''%ROWTYPE;
668
    cur CURSOR FOR
669
'''+select_query+'''
670
;
671
BEGIN
672
    OPEN cur;
673
    
674
    /* Need an EXCEPTION block for each individual row because "When an error is
675
    caught by an EXCEPTION clause, [...] all changes to persistent database
676
    state within the block are rolled back."
677
    This is unfortunate because "A block containing an EXCEPTION clause is
678
    significantly more expensive to enter and exit than a block without one."
679
    (http://www.postgresql.org/docs/8.3/static/plpgsql-control-structures.html\
680
#PLPGSQL-ERROR-TRAPPING)
681
    */
682
    WHILE true LOOP
683
        BEGIN
684
            FETCH FROM cur INTO \
685
'''+(', '.join((c.to_str(db) for c in row)))+''';
686
            EXIT WHEN NOT FOUND;
687
            
688
            RETURN QUERY
689
'''+mk_insert(sql_gen.Values(row).to_str(db))+'''
690
;
691
        EXCEPTION
692
            /* Note that cursor stays at current position in spite of automatic
693
            ROLLBACK to before EXCEPTION block: "If a previously opened cursor
694
            is affected by a FETCH command inside a savepoint that is later
695
            rolled back, the cursor position remains at the position that FETCH
696
            left it pointing to (that is, FETCH is not rolled back)."
697
            (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html)
698
            */
699
            WHEN unique_violation THEN NULL; -- continue to next row
700
        END;
701
    END LOOP;
702
    
703
    CLOSE cur;
704
END;\
705
'''
706
    else: query = mk_insert(select_query)
707
    
708
    if embeddable:
709
        # Create function
710
        function_name = sql_gen.clean_name(first_line)
711
        while True:
712
            try:
713
                function = db.TempFunction(function_name)
714
                
715
                function_query = '''\
716
CREATE FUNCTION '''+function.to_str(db)+'''()
717
RETURNS SETOF '''+return_type+'''
718
LANGUAGE '''+lang+'''
719
AS $$
720
'''+query+'''
721
$$;
722
'''
723
                run_query(db, function_query, recover=True, cacheable=True,
724
                    log_ignore_excs=(DuplicateException,))
725
                break # this version was successful
726
            except DuplicateException, e:
727
                function_name = next_version(function_name)
728
                # try again with next version of name
729
        
730
        # Return query that uses function
731
        cols = None
732
        if returning != None: cols = [returning]
733
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(function),
734
            cols) # AS clause requires function alias
735
        return mk_select(db, func_table, start=0, order_by=None)
736
    
737
    return query
738

    
739
def insert_select(db, *args, **kw_args):
740
    '''For params, see mk_insert_select() and run_query_into()
741
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
742
        values in
743
    '''
744
    into = kw_args.pop('into', None)
745
    if into != None: kw_args['embeddable'] = True
746
    recover = kw_args.pop('recover', None)
747
    if kw_args.get('ignore', False): recover = True
748
    cacheable = kw_args.pop('cacheable', True)
749
    log_level = kw_args.pop('log_level', 2)
750
    
751
    return run_query_into(db, mk_insert_select(db, *args, **kw_args), into,
752
        recover=recover, cacheable=cacheable, log_level=log_level)
753

    
754
default = sql_gen.default # tells insert() to use the default value for a column
755

    
756
def insert(db, table, row, *args, **kw_args):
757
    '''For params, see insert_select()'''
758
    if lists.is_seq(row): cols = None
759
    else:
760
        cols = row.keys()
761
        row = row.values()
762
    row = list(row) # ensure that "== []" works
763
    
764
    if row == []: query = None
765
    else: query = sql_gen.Values(row).to_str(db)
766
    
767
    return insert_select(db, table, cols, query, *args, **kw_args)
768

    
769
def mk_update(db, table, changes=None, cond=None):
770
    '''
771
    @param changes [(col, new_value),...]
772
        * container can be any iterable type
773
        * col: sql_gen.Code|str (for col name)
774
        * new_value: sql_gen.Code|literal value
775
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
776
    @return str query
777
    '''
778
    query = 'UPDATE '+sql_gen.as_Table(table).to_str(db)+'\nSET\n'
779
    query += ',\n'.join((sql_gen.to_name_only_col(col, table).to_str(db)+' = '
780
        +sql_gen.as_Value(new_value).to_str(db) for col, new_value in changes))
781
    if cond != None: query += '\nWHERE\n'+cond.to_str(db)
782
    
783
    return query
784

    
785
def update(db, *args, **kw_args):
786
    '''For params, see mk_update() and run_query()'''
787
    recover = kw_args.pop('recover', None)
788
    
789
    return run_query(db, mk_update(db, *args, **kw_args), recover)
790

    
791
def last_insert_id(db):
792
    module = util.root_module(db.db)
793
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
794
    elif module == 'MySQLdb': return db.insert_id()
795
    else: return None
796

    
797
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
798
    '''Creates a mapping from original column names (which may have collisions)
799
    to names that will be distinct among the columns' tables.
800
    This is meant to be used for several tables that are being joined together.
801
    @param cols The columns to combine. Duplicates will be removed.
802
    @param into The table for the new columns.
803
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
804
        columns will be included in the mapping even if they are not in cols.
805
        The tables of the provided Col objects will be changed to into, so make
806
        copies of them if you want to keep the original tables.
807
    @param as_items Whether to return a list of dict items instead of a dict
808
    @return dict(orig_col=new_col, ...)
809
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
810
        * new_col: sql_gen.Col(orig_col_name, into)
811
        * All mappings use the into table so its name can easily be
812
          changed for all columns at once
813
    '''
814
    cols = lists.uniqify(cols)
815
    
816
    items = []
817
    for col in preserve:
818
        orig_col = copy.copy(col)
819
        col.table = into
820
        items.append((orig_col, col))
821
    preserve = set(preserve)
822
    for col in cols:
823
        if col not in preserve:
824
            items.append((col, sql_gen.Col(str(col), into, col.srcs)))
825
    
826
    if not as_items: items = dict(items)
827
    return items
828

    
829
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
830
    '''For params, see mk_flatten_mapping()
831
    @return See return value of mk_flatten_mapping()
832
    '''
833
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
834
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
835
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
836
        into=into, add_indexes_=True)
837
    return dict(items)
838

    
839
def mk_track_data_error(db, errors_table, cols, value, error_code, error):
840
    assert cols != ()
841
    
842
    cols = map(sql_gen.to_name_only_col, cols)
843
    
844
    columns_cols = ['column']
845
    columns = sql_gen.NamedValues('columns', columns_cols,
846
        [[c.name] for c in cols])
847
    values_cols = ['value', 'error_code', 'error']
848
    values = sql_gen.NamedValues('values', values_cols,
849
        [value, error_code, error])
850
    
851
    select_cols = columns_cols+values_cols
852
    name_only_cols = map(sql_gen.to_name_only_col, select_cols)
853
    errors_table = sql_gen.NamedTable('errors', errors_table)
854
    joins = [columns, sql_gen.Join(values, type_='CROSS'),
855
        sql_gen.Join(errors_table, dict(zip(name_only_cols, select_cols)),
856
        sql_gen.filter_out)]
857
    
858
    return mk_insert_select(db, errors_table, name_only_cols,
859
        mk_select(db, joins, select_cols, order_by=None))
860

    
861
def track_data_error(db, errors_table, cols, *args, **kw_args):
862
    '''
863
    @param errors_table If None, does nothing.
864
    '''
865
    if errors_table == None or cols == (): return
866
    run_query(db, mk_track_data_error(db, errors_table, cols, *args, **kw_args),
867
        cacheable=True, log_level=4)
868

    
869
def cast(db, type_, col, errors_table=None):
870
    '''Casts an (unrenamed) column or value.
871
    If errors_table set and col has srcs, saves errors in errors_table (using
872
    col's srcs attr as the source columns) and converts errors to warnings.
873
    @param col str|sql_gen.Col|sql_gen.Literal
874
    @param errors_table None|sql_gen.Table|str
875
    '''
876
    col = sql_gen.as_Col(col)
877
    save_errors = (errors_table != None and isinstance(col, sql_gen.Col)
878
        and col.srcs != ())
879
    if not save_errors: return sql_gen.Cast(type_, col) # can't save errors
880
    
881
    assert not isinstance(col, sql_gen.NamedCol)
882
    
883
    errors_table = sql_gen.as_Table(errors_table)
884
    srcs = map(sql_gen.to_name_only_col, col.srcs)
885
    function_name = str(sql_gen.FunctionCall(type_, *srcs))
886
    function = db.TempFunction(function_name)
887
    
888
    while True:
889
        # Create function definition
890
        query = '''\
891
CREATE FUNCTION '''+function.to_str(db)+'''(value text)
892
RETURNS '''+type_+'''
893
LANGUAGE plpgsql
894
STRICT
895
AS $$
896
BEGIN
897
    /* The explicit cast to the return type is needed to make the cast happen
898
    inside the try block. (Implicit casts to the return type happen at the end
899
    of the function, outside any block.) */
900
    RETURN value::'''+type_+''';
901
EXCEPTION
902
    WHEN data_exception THEN
903
        -- Save error in errors table.
904
        -- Insert the value and error for *each* source column.
905
'''+mk_track_data_error(db, errors_table, srcs,
906
    *map(sql_gen.CustomCode, ['value', 'SQLSTATE', 'SQLERRM']))+''';
907
        
908
        RAISE WARNING '%', SQLERRM;
909
        RETURN NULL;
910
END;
911
$$;
912
'''
913
        
914
        # Create function
915
        try:
916
            run_query(db, query, recover=True, cacheable=True,
917
                log_ignore_excs=(DuplicateException,))
918
            break # successful
919
        except DuplicateException:
920
            function.name = next_version(function.name)
921
            # try again with next version of name
922
    
923
    return sql_gen.FunctionCall(function, col)
924

    
925
##### Database structure queries
926

    
927
def table_row_count(db, table, recover=None):
928
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
929
        order_by=None, start=0), recover=recover, log_level=3))
930

    
931
def table_cols(db, table, recover=None):
932
    return list(col_names(select(db, table, limit=0, order_by=None,
933
        recover=recover, log_level=4)))
934

    
935
def pkey(db, table, recover=None):
936
    '''Assumed to be first column in table'''
937
    return table_cols(db, table, recover)[0]
938

    
939
not_null_col = 'not_null_col'
940

    
941
def table_not_null_col(db, table, recover=None):
942
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
943
    if not_null_col in table_cols(db, table, recover): return not_null_col
944
    else: return pkey(db, table, recover)
945

    
946
def index_cols(db, table, index):
947
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
948
    automatically created. When you don't know whether something is a UNIQUE
949
    constraint or a UNIQUE index, use this function.'''
950
    module = util.root_module(db.db)
951
    if module == 'psycopg2':
952
        return list(values(run_query(db, '''\
953
SELECT attname
954
FROM
955
(
956
        SELECT attnum, attname
957
        FROM pg_index
958
        JOIN pg_class index ON index.oid = indexrelid
959
        JOIN pg_class table_ ON table_.oid = indrelid
960
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
961
        WHERE
962
            table_.relname = '''+db.esc_value(table)+'''
963
            AND index.relname = '''+db.esc_value(index)+'''
964
    UNION
965
        SELECT attnum, attname
966
        FROM
967
        (
968
            SELECT
969
                indrelid
970
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
971
                    AS indkey
972
            FROM pg_index
973
            JOIN pg_class index ON index.oid = indexrelid
974
            JOIN pg_class table_ ON table_.oid = indrelid
975
            WHERE
976
                table_.relname = '''+db.esc_value(table)+'''
977
                AND index.relname = '''+db.esc_value(index)+'''
978
        ) s
979
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
980
) s
981
ORDER BY attnum
982
'''
983
            , cacheable=True, log_level=4)))
984
    else: raise NotImplementedError("Can't list index columns for "+module+
985
        ' database')
986

    
987
def constraint_cols(db, table, constraint):
988
    module = util.root_module(db.db)
989
    if module == 'psycopg2':
990
        return list(values(run_query(db, '''\
991
SELECT attname
992
FROM pg_constraint
993
JOIN pg_class ON pg_class.oid = conrelid
994
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
995
WHERE
996
    relname = '''+db.esc_value(table)+'''
997
    AND conname = '''+db.esc_value(constraint)+'''
998
ORDER BY attnum
999
'''
1000
            )))
1001
    else: raise NotImplementedError("Can't list constraint columns for "+module+
1002
        ' database')
1003

    
1004
row_num_col = '_row_num'
1005

    
1006
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1007
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1008
    Currently, only function calls are supported as expressions.
1009
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1010
        This allows indexes to be used for comparisons where NULLs are equal.
1011
    '''
1012
    exprs = lists.mk_seq(exprs)
1013
    
1014
    # Parse exprs
1015
    old_exprs = exprs[:]
1016
    exprs = []
1017
    cols = []
1018
    for i, expr in enumerate(old_exprs):
1019
        expr = sql_gen.as_Col(expr, table)
1020
        
1021
        # Handle nullable columns
1022
        if ensure_not_null_:
1023
            try: expr = ensure_not_null(db, expr)
1024
            except KeyError: pass # unknown type, so just create plain index
1025
        
1026
        # Extract col
1027
        expr = copy.deepcopy(expr) # don't modify input!
1028
        if isinstance(expr, sql_gen.FunctionCall):
1029
            col = expr.args[0]
1030
            expr = sql_gen.Expr(expr)
1031
        else: col = expr
1032
        assert isinstance(col, sql_gen.Col)
1033
        
1034
        # Extract table
1035
        if table == None:
1036
            assert sql_gen.is_table_col(col)
1037
            table = col.table
1038
        
1039
        col.table = None
1040
        
1041
        exprs.append(expr)
1042
        cols.append(col)
1043
    
1044
    table = sql_gen.as_Table(table)
1045
    index = sql_gen.Table(str(sql_gen.Col(','.join(map(str, cols)), table)))
1046
    
1047
    # Add index
1048
    while True:
1049
        str_ = 'CREATE'
1050
        if unique: str_ += ' UNIQUE'
1051
        str_ += ' INDEX '+index.to_str(db)+' ON '+table.to_str(db)+' ('+(
1052
            ', '.join((v.to_str(db) for v in exprs)))+')'
1053
        
1054
        try:
1055
            run_query(db, str_, recover=True, cacheable=True, log_level=3,
1056
                log_ignore_excs=(DuplicateException,))
1057
            break
1058
        except DuplicateException:
1059
            index.name = next_version(index.name)
1060
            # try again with next version of name
1061

    
1062
def add_pkey(db, table, cols=None, recover=None):
1063
    '''Adds a primary key.
1064
    @param cols [sql_gen.Col,...] The columns in the primary key.
1065
        Defaults to the first column in the table.
1066
    @pre The table must not already have a primary key.
1067
    '''
1068
    table = sql_gen.as_Table(table)
1069
    if cols == None: cols = [pkey(db, table, recover)]
1070
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1071
    
1072
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1073
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1074
        log_ignore_excs=(DuplicateException,))
1075

    
1076
def add_not_null(db, col):
1077
    table = col.table
1078
    col = sql_gen.to_name_only_col(col)
1079
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1080
        +col.to_str(db)+' SET NOT NULL', cacheable=True)
1081

    
1082
def add_index_col(db, col, suffix, expr, nullable=True):
1083
    if sql_gen.index_col(col) != None: return # already has index col
1084
    
1085
    new_col = sql_gen.suffixed_col(col, suffix)
1086
    
1087
    # Add column
1088
    while True:
1089
        new_typed_col = sql_gen.TypedCol(new_col.name, db.col_info(col).type)
1090
        try:
1091
            add_col(db, col.table, new_typed_col,
1092
                log_ignore_excs=(DuplicateException,))
1093
            break
1094
        except DuplicateException:
1095
            new_col.name = next_version(new_col.name)
1096
            # try again with next version of name
1097
    
1098
    update(db, col.table, [(new_col, expr)])
1099
    if not nullable: add_not_null(db, new_col)
1100
    add_index(db, new_col)
1101
    
1102
    col.table.index_cols[col.name] = new_col
1103

    
1104
def ensure_not_null(db, col):
1105
    '''For params, see sql_gen.ensure_not_null()'''
1106
    expr = sql_gen.ensure_not_null(db, col)
1107
    
1108
    # If nullable column in a temp table, add separate column instead
1109
    if sql_gen.is_temp_col(col) and isinstance(expr, sql_gen.EnsureNotNull):
1110
        add_index_col(db, col, '::NOT NULL', expr, nullable=False)
1111
        expr = sql_gen.index_col(col)
1112
    
1113
    return expr
1114

    
1115
already_indexed = object() # tells add_indexes() the pkey has already been added
1116

    
1117
def add_indexes(db, table, has_pkey=True):
1118
    '''Adds an index on all columns in a table.
1119
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1120
        index should be added on the first column.
1121
        * If already_indexed, the pkey is assumed to have already been added
1122
    '''
1123
    cols = table_cols(db, table)
1124
    if has_pkey:
1125
        if has_pkey is not already_indexed: add_pkey(db, table)
1126
        cols = cols[1:]
1127
    for col in cols: add_index(db, col, table)
1128

    
1129
def add_col(db, table, col, **kw_args):
1130
    assert isinstance(col, sql_gen.TypedCol)
1131
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '
1132
        +col.to_str(db), recover=True, cacheable=True, **kw_args)
1133

    
1134
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1135
    constraints='PRIMARY KEY')
1136

    
1137
def add_row_num(db, table):
1138
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1139
    be the primary key.'''
1140
    add_col(db, table, row_num_typed_col, log_level=3)
1141

    
1142
def cast_temp_col(db, type_, col, errors_table=None):
1143
    '''Like cast(), but creates a new column with the cast values if the input
1144
    is a column.
1145
    @return The new column or cast value
1146
    '''
1147
    def cast_(col): return cast(db, type_, col, errors_table)
1148
    
1149
    try: col = sql_gen.underlying_col(col)
1150
    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)
1151
    
1152
    table = col.table
1153
    new_col = sql_gen.Col(sql_gen.concat(col.name, '::'+type_), table, col.srcs)
1154
    expr = cast_(col)
1155
    add_col(db, table, sql_gen.TypedCol(new_col.name, type_))
1156
    update(db, table, [(new_col, expr)])
1157
    add_index(db, new_col)
1158
    
1159
    return new_col
1160

    
1161
def drop_table(db, table):
1162
    table = sql_gen.as_Table(table)
1163
    return run_query(db, 'DROP TABLE IF EXISTS '+table.to_str(db)+' CASCADE')
1164

    
1165
def create_table(db, table, cols, has_pkey=True, col_indexes=True):
1166
    '''Creates a table.
1167
    @param cols [sql_gen.TypedCol,...] The column names and types
1168
    @param has_pkey If set, the first column becomes the primary key.
1169
    @param col_indexes bool|[ref]
1170
        * If True, indexes will be added on all non-pkey columns.
1171
        * If a list reference, [0] will be set to a function to do this.
1172
          This can be used to delay index creation until the table is populated.
1173
    '''
1174
    table = sql_gen.as_Table(table)
1175
    
1176
    if has_pkey:
1177
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1178
        pkey.constraints = 'PRIMARY KEY'
1179
    
1180
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
1181
    str_ += '\n, '.join(v.to_str(db) for v in cols)
1182
    str_ += '\n);\n'
1183
    run_query(db, str_, cacheable=True, log_level=2)
1184
    
1185
    # Add indexes
1186
    if has_pkey: has_pkey = already_indexed
1187
    def add_indexes_(): add_indexes(db, table, has_pkey)
1188
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1189
    elif col_indexes: add_indexes_() # add now
1190

    
1191
def vacuum(db, table):
1192
    table = sql_gen.as_Table(table)
1193
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1194
        log_level=3))
1195

    
1196
def truncate(db, table, schema='public', **kw_args):
1197
    '''For params, see run_query()'''
1198
    table = sql_gen.as_Table(table, schema)
1199
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1200

    
1201
def empty_temp(db, tables):
1202
    if db.debug_temp: return # leave temp tables there for debugging
1203
    tables = lists.mk_seq(tables)
1204
    for table in tables: truncate(db, table, log_level=3)
1205

    
1206
def tables(db, schema_like='public', table_like='%', exact=False):
1207
    if exact: compare = '='
1208
    else: compare = 'LIKE'
1209
    
1210
    module = util.root_module(db.db)
1211
    if module == 'psycopg2':
1212
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1213
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1214
        return values(select(db, 'pg_tables', ['tablename'], conds,
1215
            order_by='tablename', log_level=4))
1216
    elif module == 'MySQLdb':
1217
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1218
            , cacheable=True, log_level=4))
1219
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1220

    
1221
def table_exists(db, table):
1222
    table = sql_gen.as_Table(table)
1223
    return list(tables(db, table.schema, table.name, exact=True)) != []
1224

    
1225
def function_exists(db, function):
1226
    function = sql_gen.as_Function(function)
1227
    
1228
    info_table = sql_gen.Table('routines', 'information_schema')
1229
    conds = [('routine_name', function.name)]
1230
    schema = function.schema
1231
    if schema != None: conds.append(('routine_schema', schema))
1232
    # Exclude trigger functions, since they cannot be called directly
1233
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1234
    
1235
    return list(values(select(db, info_table, ['routine_name'], conds,
1236
        order_by='routine_schema', limit=1, log_level=4))) != []
1237
        # TODO: order_by search_path schema order
1238

    
1239
def errors_table(db, table, if_exists=True):
1240
    '''
1241
    @param if_exists If set, returns None if the errors table doesn't exist
1242
    @return None|sql_gen.Table
1243
    '''
1244
    table = sql_gen.as_Table(table)
1245
    if table.srcs != (): table = table.srcs[0]
1246
    
1247
    errors_table = sql_gen.suffixed_table(table, '.errors')
1248
    if if_exists and not table_exists(db, errors_table): return None
1249
    return errors_table
1250

    
1251
##### Database management
1252

    
1253
def empty_db(db, schema='public', **kw_args):
1254
    '''For kw_args, see tables()'''
1255
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1256

    
1257
##### Heuristic queries
1258

    
1259
def put(db, table, row, pkey_=None, row_ct_ref=None):
1260
    '''Recovers from errors.
1261
    Only works under PostgreSQL (uses INSERT RETURNING).
1262
    '''
1263
    row = sql_gen.ColDict(db, table, row)
1264
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1265
    
1266
    try:
1267
        cur = insert(db, table, row, pkey_, recover=True)
1268
        if row_ct_ref != None and cur.rowcount >= 0:
1269
            row_ct_ref[0] += cur.rowcount
1270
        return value(cur)
1271
    except DuplicateKeyException, e:
1272
        row = sql_gen.ColDict(db, table,
1273
            util.dict_subset_right_join(row, e.cols))
1274
        return value(select(db, table, [pkey_], row, recover=True))
1275

    
1276
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1277
    '''Recovers from errors'''
1278
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1279
    except StopIteration:
1280
        if not create: raise
1281
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1282

    
1283
def is_func_result(col):
1284
    return col.table.name.find('(') >= 0 and col.name == 'result'
1285

    
1286
def into_table_name(out_table, in_tables0, mapping, is_func):
1287
    def in_col_str(in_col):
1288
        in_col = sql_gen.remove_col_rename(in_col)
1289
        if isinstance(in_col, sql_gen.Col):
1290
            table = in_col.table
1291
            if table == in_tables0:
1292
                in_col = sql_gen.to_name_only_col(in_col)
1293
            elif is_func_result(in_col): in_col = table # omit col name
1294
        return str(in_col)
1295
    
1296
    str_ = str(out_table)
1297
    if is_func:
1298
        str_ += '('
1299
        
1300
        try: value_in_col = mapping['value']
1301
        except KeyError:
1302
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1303
                for k, v in mapping.iteritems()))
1304
        else: str_ += in_col_str(value_in_col)
1305
        
1306
        str_ += ')'
1307
    else:
1308
        out_col = 'rank'
1309
        try: in_col = mapping[out_col]
1310
        except KeyError: str_ += '_pkeys'
1311
        else: # has a rank column, so hierarchical
1312
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
1313
    return str_
1314

    
1315
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
1316
    default=None, is_func=False, on_error=exc.raise_):
1317
    '''Recovers from errors.
1318
    Only works under PostgreSQL (uses INSERT RETURNING).
1319
    IMPORTANT: Must be run at the *beginning* of a transaction.
1320
    @param in_tables The main input table to select from, followed by a list of
1321
        tables to join with it using the main input table's pkey
1322
    @param mapping dict(out_table_col=in_table_col, ...)
1323
        * out_table_col: str (*not* sql_gen.Col)
1324
        * in_table_col: sql_gen.Col|literal-value
1325
    @param into The table to contain the output and input pkeys.
1326
        Defaults to `out_table.name+'_pkeys'`.
1327
    @param default The *output* column to use as the pkey for missing rows.
1328
        If this output column does not exist in the mapping, uses None.
1329
    @param is_func Whether out_table is the name of a SQL function, not a table
1330
    @return sql_gen.Col Where the output pkeys are made available
1331
    '''
1332
    out_table = sql_gen.as_Table(out_table)
1333
    
1334
    def log_debug(msg): db.log_debug(msg, level=1.5)
1335
    def col_ustr(str_):
1336
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
1337
    
1338
    log_debug('********** New iteration **********')
1339
    log_debug('Inserting these input columns into '+strings.as_tt(
1340
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
1341
    
1342
    is_function = function_exists(db, out_table)
1343
    
1344
    if is_function: out_pkey = 'result'
1345
    else: out_pkey = pkey(db, out_table, recover=True)
1346
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
1347
    
1348
    if mapping == {}: # need at least one column for INSERT SELECT
1349
        mapping = {out_pkey: None} # ColDict will replace with default value
1350
    
1351
    # Create input joins from list of input tables
1352
    in_tables_ = in_tables[:] # don't modify input!
1353
    in_tables0 = in_tables_.pop(0) # first table is separate
1354
    errors_table_ = errors_table(db, in_tables0)
1355
    in_pkey = pkey(db, in_tables0, recover=True)
1356
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
1357
    input_joins = [in_tables0]+[sql_gen.Join(v,
1358
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
1359
    
1360
    if into == None:
1361
        into = into_table_name(out_table, in_tables0, mapping, is_func)
1362
    into = sql_gen.as_Table(into)
1363
    
1364
    # Set column sources
1365
    in_cols = filter(sql_gen.is_table_col, mapping.values())
1366
    for col in in_cols:
1367
        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
1368
    
1369
    log_debug('Joining together input tables into temp table')
1370
    # Place in new table for speed and so don't modify input if values edited
1371
    in_table = sql_gen.Table('in')
1372
    mapping = dicts.join(mapping, flatten(db, in_table, input_joins, in_cols,
1373
        preserve=[in_pkey_col], start=0))
1374
    input_joins = [in_table]
1375
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
1376
    
1377
    mapping = sql_gen.ColDict(db, out_table, mapping)
1378
        # after applying dicts.join() because that returns a plain dict
1379
    
1380
    # Resolve default value column
1381
    if default != None:
1382
        try: default = mapping[default]
1383
        except KeyError:
1384
            db.log_debug('Default value column '
1385
                +strings.as_tt(strings.repr_no_u(default))
1386
                +' does not exist in mapping, falling back to None', level=2.1)
1387
            default = None
1388
    
1389
    pkeys_names = [in_pkey, out_pkey]
1390
    pkeys_cols = [in_pkey_col, out_pkey_col]
1391
    
1392
    pkeys_table_exists_ref = [False]
1393
    def insert_into_pkeys(joins, cols, distinct=False):
1394
        kw_args = {}
1395
        if distinct: kw_args.update(dict(distinct_on=[in_pkey_col]))
1396
        query = mk_select(db, joins, cols, order_by=None, start=0, **kw_args)
1397
        
1398
        if pkeys_table_exists_ref[0]:
1399
            insert_select(db, into, pkeys_names, query)
1400
        else:
1401
            run_query_into(db, query, into=into)
1402
            pkeys_table_exists_ref[0] = True
1403
    
1404
    limit_ref = [None]
1405
    conds = set()
1406
    distinct_on = sql_gen.ColDict(db, out_table)
1407
    def mk_main_select(joins, cols):
1408
        distinct_on_cols = [c.to_Col() for c in distinct_on.values()]
1409
        return mk_select(db, joins, cols, conds, distinct_on_cols,
1410
            limit=limit_ref[0], start=0)
1411
    
1412
    exc_strs = set()
1413
    def log_exc(e):
1414
        e_str = exc.str_(e, first_line_only=True)
1415
        log_debug('Caught exception: '+e_str)
1416
        assert e_str not in exc_strs # avoid infinite loops
1417
        exc_strs.add(e_str)
1418
    
1419
    def remove_all_rows():
1420
        log_debug('Ignoring all rows')
1421
        limit_ref[0] = 0 # just create an empty pkeys table
1422
    
1423
    def ignore(in_col, value, e):
1424
        track_data_error(db, errors_table_, in_col.srcs, value, e.cause.pgcode,
1425
            e.cause.pgerror)
1426
        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '
1427
            +strings.as_tt(repr(value)))
1428
    
1429
    def remove_rows(in_col, value, e):
1430
        ignore(in_col, value, e)
1431
        cond = (in_col, sql_gen.CompareCond(value, '!='))
1432
        assert cond not in conds # avoid infinite loops
1433
        conds.add(cond)
1434
    
1435
    def invalid2null(in_col, value, e):
1436
        ignore(in_col, value, e)
1437
        update(db, in_table, [(in_col, None)],
1438
            sql_gen.ColValueCond(in_col, value))
1439
    
1440
    def insert_pkeys_table(which):
1441
        return sql_gen.Table(sql_gen.concat(in_table.name,
1442
            '_insert_'+which+'_pkeys'))
1443
    insert_out_pkeys = insert_pkeys_table('out')
1444
    insert_in_pkeys = insert_pkeys_table('in')
1445
    
1446
    # Do inserts and selects
1447
    join_cols = sql_gen.ColDict(db, out_table)
1448
    while True:
1449
        if limit_ref[0] == 0: # special case
1450
            log_debug('Creating an empty pkeys table')
1451
            cur = run_query_into(db, mk_select(db, out_table, [out_pkey],
1452
                limit=limit_ref[0]), into=insert_out_pkeys)
1453
            break # don't do main case
1454
        
1455
        has_joins = join_cols != {}
1456
        
1457
        log_debug('Trying to insert new rows')
1458
        
1459
        # Prepare to insert new rows
1460
        insert_joins = input_joins[:] # don't modify original!
1461
        insert_args = dict(recover=True, cacheable=False)
1462
        if has_joins:
1463
            insert_joins.append(sql_gen.Join(out_table, join_cols,
1464
                sql_gen.filter_out))
1465
        else:
1466
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
1467
        main_select = mk_main_select(insert_joins, mapping.values())
1468
        
1469
        def main_insert():
1470
            if is_function:
1471
                log_debug('Calling function on input rows')
1472
                args = dict(((k.name, v) for k, v in mapping.iteritems()))
1473
                func_call = sql_gen.NamedCol(out_pkey,
1474
                    sql_gen.FunctionCall(out_table, **args))
1475
                insert_into_pkeys(input_joins, [in_pkey_col, func_call])
1476
                return None
1477
            else:
1478
                lock_table(db, out_table, 'EXCLUSIVE')
1479
                return insert_select(db, out_table, mapping.keys(), main_select,
1480
                    **insert_args)
1481
        
1482
        try:
1483
            cur = with_savepoint(db, main_insert)
1484
            break # insert successful
1485
        except MissingCastException, e:
1486
            log_exc(e)
1487
            
1488
            out_col = e.col
1489
            type_ = e.type
1490
            
1491
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
1492
                +strings.as_tt(type_))
1493
            mapping[out_col] = cast_temp_col(db, type_, mapping[out_col],
1494
                errors_table_)
1495
        except DuplicateKeyException, e:
1496
            log_exc(e)
1497
            
1498
            old_join_cols = join_cols.copy()
1499
            distinct_on.update(util.dict_subset(mapping, e.cols))
1500
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
1501
            log_debug('Ignoring existing rows, comparing on these columns:\n'
1502
                +strings.as_inline_table(join_cols, ustr=col_ustr))
1503
            assert join_cols != old_join_cols # avoid infinite loops
1504
        except NullValueException, e:
1505
            log_exc(e)
1506
            
1507
            out_col, = e.cols
1508
            try: in_col = mapping[out_col]
1509
            except KeyError:
1510
                log_debug('Missing mapping for NOT NULL column '+out_col)
1511
                remove_all_rows()
1512
            else: remove_rows(in_col, None, e)
1513
        except FunctionValueException, e:
1514
            log_exc(e)
1515
            
1516
            func_name = e.name
1517
            value = e.value
1518
            for out_col, in_col in mapping.iteritems():
1519
                in_col = sql_gen.unwrap_func_call(in_col, func_name)
1520
                invalid2null(in_col, value, e)
1521
        except DatabaseErrors, e:
1522
            log_exc(e)
1523
            
1524
            log_debug('No handler for exception')
1525
            on_error(e)
1526
            remove_all_rows()
1527
        # after exception handled, rerun loop with additional constraints
1528
    
1529
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
1530
        row_ct_ref[0] += cur.rowcount
1531
    
1532
    if is_function: pass # pkeys table already created
1533
    elif has_joins:
1534
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
1535
        log_debug('Getting output table pkeys of existing/inserted rows')
1536
        insert_into_pkeys(select_joins, pkeys_cols, distinct=True)
1537
    else:
1538
        add_row_num(db, insert_out_pkeys) # for joining with input pkeys
1539
        
1540
        log_debug('Getting input table pkeys of inserted rows')
1541
        run_query_into(db, mk_main_select(input_joins, [in_pkey]),
1542
            into=insert_in_pkeys)
1543
        add_row_num(db, insert_in_pkeys) # for joining with output pkeys
1544
        
1545
        assert table_row_count(db, insert_out_pkeys) == table_row_count(db,
1546
            insert_in_pkeys)
1547
        
1548
        log_debug('Combining output and input pkeys in inserted order')
1549
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
1550
            {row_num_col: sql_gen.join_same_not_null})]
1551
        insert_into_pkeys(pkey_joins, pkeys_names)
1552
        
1553
        empty_temp(db, [insert_out_pkeys, insert_in_pkeys])
1554
    
1555
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
1556
    add_pkey(db, into)
1557
    
1558
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
1559
    missing_rows_joins = input_joins+[sql_gen.Join(into,
1560
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
1561
        # must use join_same_not_null or query will take forever
1562
    insert_into_pkeys(missing_rows_joins,
1563
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
1564
    
1565
    assert table_row_count(db, into) == table_row_count(db, in_table)
1566
    
1567
    empty_temp(db, in_table)
1568
    
1569
    srcs = []
1570
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
1571
    return sql_gen.Col(out_pkey, into, srcs)
1572

    
1573
##### Data cleanup
1574

    
1575
def cleanup_table(db, table, cols):
1576
    table = sql_gen.as_Table(table)
1577
    cols = map(sql_gen.as_Col, cols)
1578
    
1579
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1580
        +db.esc_value(r'\N')+')')
1581
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1582
        for v in cols]
1583
    
1584
    update(db, table, changes)
(24-24/36)