/lib/sql_gen.py - Annotate - BIEN 3 - NCEAS Projects

root/lib/sql_gen.py @ 5289

-            aaronmk
+# SQL code generation
-            aaronmk
+import copy
-            aaronmk
+import itertools
-            aaronmk
+import operator
-            aaronmk
+from ordereddict import OrderedDict
-            aaronmk
+import re
-            aaronmk
+import UserDict
-            aaronmk
+import warnings
             aaronmk
-            aaronmk
+import dicts
-            aaronmk
+import exc
-            aaronmk
+import iters
 import lists
-            aaronmk
+import objects
-            aaronmk
+import strings
-            aaronmk
+import util
             aaronmk
-            aaronmk
+##### Names
             aaronmk
-            aaronmk
+identifier_max_len = 63 # works for both PostgreSQL and MySQL
             aaronmk
-            aaronmk
+def concat(str_, suffix):
-            aaronmk
+    '''Preserves version so that it won't be truncated off the string, leading
     to collisions.'''
-            aaronmk
+    # Preserve version
-            aaronmk
+    match = re.match(r'^(.*?)((?:(?:#\d+)?\)?)*(?:\.\w+)?(?:::[\w ]+)*)$', str_)
-            aaronmk
+    if match:
         str_, old_suffix = match.groups()
         suffix = old_suffix+suffix
             aaronmk
-            aaronmk
+    return strings.concat(str_, suffix, identifier_max_len)
             aaronmk
-            aaronmk
+def truncate(str_): return concat(str_, '')
             aaronmk
-            aaronmk
+def is_safe_name(name):
-            aaronmk
+    '''A name is safe *and unambiguous* if it:
     * contains only *lowercase* word (\w) characters
     * doesn't start with a digit
     * contains "_", so that it's not a keyword
-            aaronmk
+    '''
     return re.match(r'^(?=.*_)(?!\d)[^\WA-Z]+$', name)
             aaronmk
-            aaronmk
+def esc_name(name, quote='"'):
     return quote + name.replace(quote, quote+quote) + quote
         # doubling an embedded quote escapes it in both PostgreSQL and MySQL
-            aaronmk
+def unesc_name(name, quote='"'):
     removed_ref = [False]
     name = strings.remove_prefix(quote, name, removed_ref)
     if removed_ref[0]:
         name = strings.remove_suffix(quote, name, removed_ref)
         assert removed_ref[0]
         name = name.replace(quote+quote, quote)
     return name
-            aaronmk
+def clean_name(name): return name.replace('"', '').replace('`', '')
-            aaronmk
+def esc_comment(comment): return '/*'+comment.replace('*/', '* /')+'*/'
-            aaronmk
+def lstrip(str_):
     '''Also removes comments.'''
     if str_.startswith('/*'): comment, sep, str_ = str_.partition('*/')
     return str_.lstrip()
-            aaronmk
+##### General SQL code objects
             aaronmk
-            aaronmk
+class MockDb:
-            aaronmk
+    def esc_value(self, value): return strings.repr_no_u(value)
             aaronmk
-            aaronmk
+    def esc_name(self, name): return esc_name(name)
             aaronmk
     def col_info(self, col):
         return TypedCol(col.name, '<type>', CustomCode('<default>'), True)
-            aaronmk
+mockDb = MockDb()
-            aaronmk
+class BasicObject(objects.BasicObject):
     def __str__(self): return clean_name(strings.repr_no_u(self))
-            aaronmk
+##### Unparameterized code objects
-            aaronmk
+class Code(BasicObject):
-            aaronmk
+    def __init__(self, lang='sql'):
         self.lang = lang
             aaronmk
-            aaronmk
+    def to_str(self, db): raise NotImplementedError()
             aaronmk
-            aaronmk
+    def __repr__(self): return self.to_str(mockDb)
             aaronmk
-            aaronmk
+class CustomCode(Code):
-            aaronmk
+    def __init__(self, str_):
         Code.__init__(self)
         self.str_ = str_
             aaronmk
     def to_str(self, db): return self.str_
-            aaronmk
+def as_Code(value, db=None):
     '''
     @param db If set, runs db.std_code() on the value.
     '''
-            aaronmk
+    if isinstance(value, Code): return value
-            aaronmk
+    if util.is_str(value):
         if db != None: value = db.std_code(value)
         return CustomCode(value)
-            aaronmk
+    else: return Literal(value)
-            aaronmk
+class Expr(Code):
-            aaronmk
+    def __init__(self, expr):
         Code.__init__(self)
         self.expr = expr
             aaronmk
     def to_str(self, db): return '('+self.expr.to_str(db)+')'
-            aaronmk
+##### Names
 class Name(Code):
-            aaronmk
+    def __init__(self, name):
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        name = truncate(name)
         self.name = name
             aaronmk
     def to_str(self, db): return db.esc_name(self.name)
 def as_Name(value):
     if isinstance(value, Code): return value
     else: return Name(value)
-            aaronmk
+##### Literal values
-            aaronmk
+#### Primitives
-            aaronmk
+class Literal(Code):
-            aaronmk
+    def __init__(self, value):
         Code.__init__(self)
         self.value = value
             aaronmk
     def to_str(self, db): return db.esc_value(self.value)
             aaronmk
-            aaronmk
+def as_Value(value):
     if isinstance(value, Code): return value
     else: return Literal(value)
-            aaronmk
+def is_literal(value): return isinstance(value, Literal)
             aaronmk
-            aaronmk
+def is_null(value): return is_literal(value) and value.value == None
-            aaronmk
+#### Composites
-            aaronmk
+class List(Code):
     def __init__(self, values):
-            aaronmk
+        Code.__init__(self)
         self.values = values
-            aaronmk
+    def to_str(self, db): return ', '.join((v.to_str(db) for v in self.values))
             aaronmk
-            aaronmk
+class Tuple(List):
     def __init__(self, *values):
         List.__init__(self, values)
     def to_str(self, db): return '('+List.to_str(self, db)+')'
-            aaronmk
+class Row(Tuple):
     def to_str(self, db): return 'ROW'+Tuple.to_str(self, db)
             aaronmk
-            aaronmk
+### Arrays
 class Array(List):
     def __init__(self, values):
         values = map(remove_col_rename, values)
         List.__init__(self, values)
     def to_str(self, db): return 'ARRAY['+List.to_str(self, db)+']'
 def to_Array(value):
     if isinstance(value, Array): return value
     return Array(lists.mk_seq(value))
-            aaronmk
+##### Derived elements
 src_self = object() # tells Col that it is its own source column
 class Derived(Code):
     def __init__(self, srcs):
-            aaronmk
+        '''An element which was derived from some other element(s).
-            aaronmk
+        @param srcs See self.set_srcs()
         '''
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        self.set_srcs(srcs)
-            aaronmk
+    def set_srcs(self, srcs, overwrite=True):
-            aaronmk
+        '''
         @param srcs (self_type...)|src_self The element(s) this is derived from
         '''
-            aaronmk
+        if not overwrite and self.srcs != (): return # already set
-            aaronmk
+        if srcs == src_self: srcs = (self,)
         srcs = tuple(srcs) # make Col hashable
         self.srcs = srcs
     def _compare_on(self):
         compare_on = self.__dict__.copy()
         del compare_on['srcs'] # ignore
         return compare_on
 def cols_srcs(cols): return lists.uniqify(iters.flatten((v.srcs for v in cols)))
-            aaronmk
+##### Tables
-            aaronmk
+class Table(Derived):
-            aaronmk
+    def __init__(self, name, schema=None, srcs=(), is_temp=False):
-            aaronmk
+        '''
         @param schema str|None (for no schema)
-            aaronmk
+        @param srcs (Table...)|src_self See Derived.set_srcs()
-            aaronmk
+        '''
-            aaronmk
+        Derived.__init__(self, srcs)
-            aaronmk
+        if util.is_str(name): name = truncate(name)
             aaronmk
-            aaronmk
+        self.name = name
         self.schema = schema
-            aaronmk
+        self.is_temp = is_temp
-            aaronmk
+        self.index_cols = {}
             aaronmk
-            aaronmk
+    def to_str(self, db):
         str_ = ''
-            aaronmk
+        if self.schema != None: str_ += as_Name(self.schema).to_str(db)+'.'
         str_ += as_Name(self.name).to_str(db)
-            aaronmk
+        return str_
             aaronmk
     def to_Table(self): return self
             aaronmk
     def _compare_on(self):
         compare_on = Derived._compare_on(self)
         del compare_on['index_cols'] # ignore
         return compare_on
             aaronmk
-            aaronmk
+def is_underlying_table(table):
     return isinstance(table, Table) and table.to_Table() is table
             aaronmk
-            aaronmk
+class NoUnderlyingTableException(Exception):
     def __init__(self, ref):
         Exception.__init__(self, 'for: '+strings.as_tt(strings.urepr(ref)))
         self.ref = ref
             aaronmk
 def underlying_table(table):
     table = remove_table_rename(table)
-            aaronmk
+    if table != None and table.srcs:
         table, = table.srcs # for derived tables or row vars
-            aaronmk
+    if not is_underlying_table(table): raise NoUnderlyingTableException(table)
-            aaronmk
+    return table
-            aaronmk
+def as_Table(table, schema=None):
-            aaronmk
+    if table == None or isinstance(table, Code): return table
-            aaronmk
+    else: return Table(table, schema)
             aaronmk
-            aaronmk
+def suffixed_table(table, suffix):
-            aaronmk
+    table = copy.copy(table) # don't modify input!
     table.name = concat(table.name, suffix)
     return table
             aaronmk
-            aaronmk
+class NamedTable(Table):
     def __init__(self, name, code, cols=None):
         Table.__init__(self, name)
-            aaronmk
+        code = as_Table(code)
-            aaronmk
+        if not isinstance(code, (Table, FunctionCall, Expr)): code = Expr(code)
-            aaronmk
+        if cols != None: cols = [to_name_only_col(c).to_Col() for c in cols]
             aaronmk
         self.code = code
         self.cols = cols
     def to_str(self, db):
-            aaronmk
+        str_ = self.code.to_str(db)
         if str_.find('\n') >= 0: whitespace = '\n'
         else: whitespace = ' '
         str_ += whitespace+'AS '+Table.to_str(self, db)
-            aaronmk
+        if self.cols != None:
             str_ += ' ('+(', '.join((c.to_str(db) for c in self.cols)))+')'
-            aaronmk
+        return str_
     def to_Table(self): return Table(self.name)
-            aaronmk
+def remove_table_rename(table):
     if isinstance(table, NamedTable): table = table.code
     return table
-            aaronmk
+##### Columns
-            aaronmk
+class Col(Derived):
-            aaronmk
+    def __init__(self, name, table=None, srcs=()):
-            aaronmk
+        '''
         @param table Table|None (for no table)
-            aaronmk
+        @param srcs (Col...)|src_self See Derived.set_srcs()
-            aaronmk
+        '''
-            aaronmk
+        Derived.__init__(self, srcs)
-            aaronmk
+        if util.is_str(name): name = truncate(name)
-            aaronmk
+        if util.is_str(table): table = Table(table)
-            aaronmk
+        assert table == None or isinstance(table, Table)
         self.name = name
         self.table = table
-            aaronmk
+    def to_str(self, db, for_str=False):
-            aaronmk
+        str_ = as_Name(self.name).to_str(db)
-            aaronmk
+        if for_str: str_ = clean_name(str_)
-            aaronmk
+        if self.table != None:
-            aaronmk
+            table = self.table.to_Table()
-            aaronmk
+            if for_str: str_ = concat(strings.ustr(table), '.'+str_)
-            aaronmk
+            else: str_ = table.to_str(db)+'.'+str_
-            aaronmk
+        return str_
             aaronmk
-            aaronmk
+    def __str__(self): return self.to_str(mockDb, for_str=True)
             aaronmk
-            aaronmk
+    def to_Col(self): return self
             aaronmk
-            aaronmk
+def is_col(col): return isinstance(col, Col)
             aaronmk
-            aaronmk
+def is_table_col(col): return is_col(col) and col.table != None
-            aaronmk
+def index_col(col):
     if not is_table_col(col): return None
             aaronmk
     table = col.table
     try: name = table.index_cols[col.name]
     except KeyError: return None
     else: return Col(name, table, col.srcs)
             aaronmk
-            aaronmk
+def is_temp_col(col): return is_table_col(col) and col.table.is_temp
             aaronmk
-            aaronmk
+def as_Col(col, table=None, name=None):
     '''
     @param name If not None, any non-Col input will be renamed using NamedCol.
     '''
     if name != None:
         col = as_Value(col)
         if not isinstance(col, Col): col = NamedCol(name, col)
             aaronmk
     if isinstance(col, Code): return col
-            aaronmk
+    elif util.is_str(col): return Col(col, table)
     else: return Literal(col)
             aaronmk
-            aaronmk
+def with_table(col, table):
-            aaronmk
+    if isinstance(col, NamedCol): pass # doesn't take a table
     elif isinstance(col, FunctionCall):
         col = copy.deepcopy(col) # don't modify input!
         col.args[0].table = table
-            aaronmk
+    elif isinstance(col, Col):
-            aaronmk
+        col = copy.copy(col) # don't modify input!
         col.table = table
-            aaronmk
+    return col
-            aaronmk
+def with_default_table(col, table):
-            aaronmk
+    col = as_Col(col)
-            aaronmk
+    if col.table == None: col = with_table(col, table)
-            aaronmk
+    return col
-            aaronmk
+def set_cols_table(table, cols):
     table = as_Table(table)
     for i, col in enumerate(cols):
         col = cols[i] = as_Col(col)
         col.table = table
-            aaronmk
+def to_name_only_col(col, check_table=None):
     col = as_Col(col)
-            aaronmk
+    if not is_table_col(col): return col
             aaronmk
     if check_table != None:
         table = col.table
         assert table == None or table == check_table
     return Col(col.name)
-            aaronmk
+def suffixed_col(col, suffix):
     return Col(concat(col.name, suffix), col.table, col.srcs)
-            aaronmk
+def has_srcs(col): return is_col(col) and col.srcs
-            aaronmk
+def cross_join_srcs(cols):
     cols = filter(has_srcs, cols) # empty srcs will mess up the cross join
     srcs = [[s.name for s in c.srcs] for c in cols]
     return [Col(','.join(s)) for s in itertools.product(*srcs)]
             aaronmk
-            aaronmk
+class NamedCol(Col):
-            aaronmk
+    def __init__(self, name, code):
-            aaronmk
+        Col.__init__(self, name)
-            aaronmk
+        code = as_Value(code)
             aaronmk
         self.code = code
     def to_str(self, db):
-            aaronmk
+        return self.code.to_str(db)+' AS '+Col.to_str(self, db)
             aaronmk
     def to_Col(self): return Col(self.name)
             aaronmk
-            aaronmk
+def remove_col_rename(col):
     if isinstance(col, NamedCol): col = col.code
     return col
-            aaronmk
+def underlying_col(col):
     col = remove_col_rename(col)
-            aaronmk
+    if not isinstance(col, Col): raise NoUnderlyingTableException(col)
             aaronmk
-            aaronmk
+    return Col(col.name, underlying_table(col.table), col.srcs)
             aaronmk
-            aaronmk
+def wrap(wrap_func, value):
     '''Wraps a value, propagating any column renaming to the returned value.'''
     if isinstance(value, NamedCol):
         return NamedCol(value.name, wrap_func(value.code))
     else: return wrap_func(value)
-            aaronmk
+class ColDict(dicts.DictProxy):
-            aaronmk
+    '''A dict that automatically makes inserted entries Col objects.
     Anything that isn't a column is wrapped in a NamedCol with the key's column
     name by `as_Col(value, name=key.name)`.
     '''
             aaronmk
-            aaronmk
+    def __init__(self, db, keys_table, dict_={}):
-            aaronmk
+        dicts.DictProxy.__init__(self, OrderedDict())
             aaronmk
-            aaronmk
+        keys_table = as_Table(keys_table)
-            aaronmk
+        self.db = db
-            aaronmk
+        self.table = keys_table
-            aaronmk
+        self.update(dict_) # after setting vars because __setitem__() needs them
             aaronmk
-            aaronmk
+    def copy(self): return ColDict(self.db, self.table, self.inner.copy())
             aaronmk
-            aaronmk
+    def __getitem__(self, key):
         return dicts.DictProxy.__getitem__(self, self._key(key))
             aaronmk
-            aaronmk
+    def __setitem__(self, key, value):
-            aaronmk
+        key = self._key(key)
-            aaronmk
+        if value == None:
             try: value = self.db.col_info(key).default
             except NoUnderlyingTableException: pass # not a table column
-            aaronmk
+        dicts.DictProxy.__setitem__(self, key, as_Col(value, name=key.name))
             aaronmk
-            aaronmk
+    def _key(self, key): return as_Col(key, self.table)
             aaronmk
-            aaronmk
+##### Definitions
             aaronmk
-            aaronmk
+class TypedCol(Col):
     def __init__(self, name, type_, default=None, nullable=True,
         constraints=None):
         assert default == None or isinstance(default, Code)
         Col.__init__(self, name)
         self.type = type_
         self.default = default
         self.nullable = nullable
         self.constraints = constraints
     def to_str(self, db):
-            aaronmk
+        str_ = Col.to_str(self, db)+' '+as_Code(self.type).to_str(db)
-            aaronmk
+        if not self.nullable: str_ += ' NOT NULL'
         if self.default != None: str_ += ' DEFAULT '+self.default.to_str(db)
         if self.constraints != None: str_ += ' '+self.constraints
         return str_
     def to_Col(self): return Col(self.name)
-            aaronmk
+class SetOf(Code):
     def __init__(self, type_):
         Code.__init__(self)
         self.type = type_
     def to_str(self, db):
         return 'SETOF '+self.type.to_str(db)
-            aaronmk
+class RowType(Code):
     def __init__(self, table):
         Code.__init__(self)
         self.table = table
     def to_str(self, db):
         return self.table.to_str(db)+'%ROWTYPE'
-            aaronmk
+class ColType(Code):
     def __init__(self, col):
         Code.__init__(self)
         self.col = col
     def to_str(self, db):
         return self.col.to_str(db)+'%TYPE'
-            aaronmk
+class ArrayType(Code):
     def __init__(self, elem_type):
         Code.__init__(self)
         elem_type = as_Code(elem_type)
         self.elem_type = elem_type
     def to_str(self, db):
         return self.elem_type.to_str(db)+'[]'
-            aaronmk
+##### Functions
-            aaronmk
+Function = Table
-            aaronmk
+as_Function = as_Table
-            aaronmk
+class InternalFunction(CustomCode): pass
-            aaronmk
+#### Calls
-            aaronmk
+class NamedArg(NamedCol):
     def __init__(self, name, value):
         NamedCol.__init__(self, name, value)
     def to_str(self, db):
         return Col.to_str(self, db)+' := '+self.code.to_str(db)
-            aaronmk
+class FunctionCall(Code):
-            aaronmk
+    def __init__(self, function, *args, **kw_args):
-            aaronmk
+        '''
-            aaronmk
+        @param args [Code|literal-value...] The function's arguments
-            aaronmk
+        '''
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        function = as_Function(function)
-            aaronmk
+        def filter_(arg): return remove_col_rename(as_Value(arg))
         args = map(filter_, args)
         args += [NamedArg(k, filter_(v)) for k, v in kw_args.iteritems()]
             aaronmk
         self.function = function
         self.args = args
     def to_str(self, db):
         args_str = ', '.join((v.to_str(db) for v in self.args))
         return self.function.to_str(db)+'('+args_str+')'
-            aaronmk
+def wrap_in_func(function, value):
     '''Wraps a value inside a function call.
     Propagates any column renaming to the returned value.
     '''
-            aaronmk
+    return wrap(lambda v: FunctionCall(function, v), value)
             aaronmk
-            aaronmk
+def unwrap_func_call(func_call, check_name=None):
     '''Unwraps any function call to its first argument.
     Also removes any column renaming.
     '''
     func_call = remove_col_rename(func_call)
     if not isinstance(func_call, FunctionCall): return func_call
     if check_name != None:
         name = func_call.function.name
         assert name == None or name == check_name
     return func_call.args[0]
-            aaronmk
+#### Definitions
 class FunctionDef(Code):
-            aaronmk
+    def __init__(self, function, return_type, body, params=[], modifiers=None):
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        return_type = as_Code(return_type)
-            aaronmk
+        body = as_Code(body)
-            aaronmk
+        self.function = function
         self.return_type = return_type
         self.body = body
-            aaronmk
+        self.params = params
-            aaronmk
+        self.modifiers = modifiers
             aaronmk
     def to_str(self, db):
-            aaronmk
+        params_str = (', '.join((p.to_str(db) for p in self.params)))
-            aaronmk
+        str_ = '''\
-            aaronmk
+CREATE FUNCTION '''+self.function.to_str(db)+'''('''+params_str+''')
-            aaronmk
+RETURNS '''+self.return_type.to_str(db)+'''
-            aaronmk
+LANGUAGE '''+self.body.lang+'''
-            aaronmk
+'''
         if self.modifiers != None: str_ += self.modifiers+'\n'
         str_ += '''\
-            aaronmk
+AS $$
-            aaronmk
+'''+self.body.to_str(db)+'''
-            aaronmk
+$$;
 '''
         return str_
-            aaronmk
+class FunctionParam(TypedCol):
     def __init__(self, name, type_, default=None, out=False):
         TypedCol.__init__(self, name, type_, default)
         self.out = out
     def to_str(self, db):
         str_ = TypedCol.to_str(self, db)
         if self.out: str_ = 'OUT '+str_
         return str_
     def to_Col(self): return Col(self.name)
-            aaronmk
+### PL/pgSQL
-            aaronmk
+class ReturnQuery(Code):
     def __init__(self, query):
         Code.__init__(self)
         query = as_Code(query)
         self.query = query
     def to_str(self, db):
         return 'RETURN QUERY\n'+strings.indent(self.query.to_str(db))+';\n'
-            aaronmk
+## Exceptions
-            aaronmk
+class BaseExcHandler(BasicObject):
     def to_str(self, db, body): raise NotImplementedError()
             aaronmk
     def __repr__(self): return self.to_str(mockDb, '<body>')
             aaronmk
-            aaronmk
+suppress_exc = 'NULL;\n';
-            aaronmk
+reraise_exc = 'RAISE USING ERRCODE = SQLSTATE, MESSAGE = SQLERRM;\n';
-            aaronmk
+class ExcHandler(BaseExcHandler):
-            aaronmk
+    def __init__(self, exc, handler=None):
         if handler != None: handler = as_Code(handler)
         self.exc = exc
         self.handler = handler
     def to_str(self, db, body):
         body = as_Code(body)
-            aaronmk
+        if self.handler != None:
             handler_str = '\n'+strings.indent(self.handler.to_str(db), 2)
-            aaronmk
+        else: handler_str = ' '+suppress_exc
             aaronmk
         str_ = '''\
 BEGIN
-            aaronmk
+'''+strings.indent(body.to_str(db))+'''\
-            aaronmk
+EXCEPTION
-            aaronmk
+    WHEN '''+self.exc+''' THEN'''+handler_str+'''\
-            aaronmk
+END;\
 '''
         return str_
-            aaronmk
+class NestedExcHandler(BaseExcHandler):
     def __init__(self, *handlers):
         '''
         @param handlers Sorted from outermost to innermost
         '''
         self.handlers = handlers
     def to_str(self, db, body):
         for handler in reversed(self.handlers): body = handler.to_str(db, body)
         return body
-            aaronmk
+class ExcToWarning(Code):
     def __init__(self, return_):
         '''
         @param return_ Statement to return a default value in case of error
         '''
         Code.__init__(self)
         return_ = as_Code(return_)
         self.return_ = return_
     def to_str(self, db):
         return "RAISE WARNING '%', SQLERRM;\n"+self.return_.to_str(db)
-            aaronmk
+unique_violation_handler = ExcHandler('unique_violation')
-            aaronmk
+# Note doubled "\"s because inside Python string
-            aaronmk
+plpythonu_error_handler = ExcHandler('internal_error', '''\
-            aaronmk
+-- Handle PL/Python exceptions
-            aaronmk
+DECLARE
-            aaronmk
+    matches text[] := regexp_matches(SQLERRM,
         E'^(?:PL/Python: )?(\\\\w+): (.*)$'); -- .* also matches \\n
-            aaronmk
+    exc_name text := matches[1];
     msg text := matches[2];
 BEGIN
-            aaronmk
+    /* Re-raise PL/Python exceptions with the PL/Python prefix removed.
     This allows the exception to be parsed like a native exception.
     Always raise as data_exception so it goes in the errors table. */
-            aaronmk
+    IF exc_name IS NOT NULL THEN
-            aaronmk
+        RAISE data_exception USING MESSAGE = msg;
-            aaronmk
+    -- Re-raise non-PL/Python exceptions
     ELSE
         '''+reraise_exc+'''\
-            aaronmk
+    END IF;
 END;
-            aaronmk
+''')
-            aaronmk
+def data_exception_handler(handler):
     return ExcHandler('data_exception', handler)
-            aaronmk
+row_var = Table('row')
-            aaronmk
+class RowExcIgnore(Code):
     def __init__(self, row_type, select_query, with_row, cols=None,
-            aaronmk
+        exc_handler=unique_violation_handler, row_var=row_var):
-            aaronmk
+        '''
         @param row_type Ignored if a custom row_var is used.
         @pre If a custom row_var is used, it must already be defined.
         '''
-            aaronmk
+        Code.__init__(self, lang='plpgsql')
-            aaronmk
+        row_type = as_Code(row_type)
-            aaronmk
+        select_query = as_Code(select_query)
         with_row = as_Code(with_row)
-            aaronmk
+        row_var = as_Table(row_var)
             aaronmk
         self.row_type = row_type
         self.select_query = select_query
         self.with_row = with_row
         self.cols = cols
-            aaronmk
+        self.exc_handler = exc_handler
-            aaronmk
+        self.row_var = row_var
             aaronmk
     def to_str(self, db):
-            aaronmk
+        if self.cols == None: row_vars = [self.row_var]
         else: row_vars = [Col(c.name, self.row_var) for c in self.cols]
             aaronmk
-            aaronmk
+        # Need an EXCEPTION block for each individual row because "When an error
         # is caught by an EXCEPTION clause, [...] all changes to persistent
         # database state within the block are rolled back."
         # This is unfortunate because "A block containing an EXCEPTION clause is
         # significantly more expensive to enter and exit than a block without
         # one."
         # (http://www.postgresql.org/docs/8.3/static/\
         # plpgsql-control-structures.html#PLPGSQL-ERROR-TRAPPING)
-            aaronmk
+        str_ = '''\
-            aaronmk
+FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN
 '''+strings.indent(self.select_query.to_str(db))+'''\
 LOOP
 '''+strings.indent(self.exc_handler.to_str(db, self.with_row))+'''\
 END LOOP;
 '''
-            aaronmk
+        if self.row_var == row_var:
-            aaronmk
+            str_ = '''\
-            aaronmk
+DECLARE
-            aaronmk
+    '''+self.row_var.to_str(db)+''' '''+self.row_type.to_str(db)+''';
-            aaronmk
+BEGIN
-            aaronmk
+'''+strings.indent(str_)+'''\
 END;
-            aaronmk
+'''
         return str_
-            aaronmk
+##### Casts
 class Cast(FunctionCall):
     def __init__(self, type_, value):
-            aaronmk
+        type_ = as_Code(type_)
-            aaronmk
+        value = as_Value(value)
         self.type_ = type_
         self.value = value
     def to_str(self, db):
-            aaronmk
+        return 'CAST('+self.value.to_str(db)+' AS '+self.type_.to_str(db)+')'
             aaronmk
-            aaronmk
+def cast_literal(value):
-            aaronmk
+    if not is_literal(value): return value
             aaronmk
     if util.is_str(value.value): value = Cast('text', value)
     return value
-            aaronmk
+##### Conditions
             aaronmk
-            aaronmk
+class NotCond(Code):
     def __init__(self, cond):
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        self.cond = cond
     def to_str(self, db): return 'NOT '+self.cond.to_str(db)
-            aaronmk
+class ColValueCond(Code):
     def __init__(self, col, value):
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        value = as_ValueCond(value)
         self.col = col
         self.value = value
     def to_str(self, db): return self.value.to_str(db, self.col)
-            aaronmk
+def combine_conds(conds, keyword=None):
     '''
     @param keyword The keyword to add before the conditions, if any
     '''
     str_ = ''
     if keyword != None:
         if conds == []: whitespace = ''
         elif len(conds) == 1: whitespace = ' '
         else: whitespace = '\n'
         str_ += keyword+whitespace
     str_ += '\nAND '.join(conds)
     return str_
-            aaronmk
+##### Condition column comparisons
-            aaronmk
+class ValueCond(BasicObject):
-            aaronmk
+    def __init__(self, value):
-            aaronmk
+        value = remove_col_rename(as_Value(value))
             aaronmk
         self.value = value
             aaronmk
-            aaronmk
+    def to_str(self, db, left_value):
-            aaronmk
+        '''
-            aaronmk
+        @param left_value The Code object that the condition is being applied on
-            aaronmk
+        '''
         raise NotImplemented()
             aaronmk
-            aaronmk
+    def __repr__(self): return self.to_str(mockDb, '<left_value>')
             aaronmk
 class CompareCond(ValueCond):
     def __init__(self, value, operator='='):
-            aaronmk
+        '''
         @param operator By default, compares NULL values literally. Use '~=' or
             '~!=' to pass NULLs through.
         '''
-            aaronmk
+        ValueCond.__init__(self, value)
         self.operator = operator
-            aaronmk
+    def to_str(self, db, left_value):
-            aaronmk
+        left_value = remove_col_rename(as_Col(left_value))
             aaronmk
-            aaronmk
+        right_value = self.value
         # Parse operator
-            aaronmk
+        operator = self.operator
-            aaronmk
+        passthru_null_ref = [False]
         operator = strings.remove_prefix('~', operator, passthru_null_ref)
         neg_ref = [False]
         operator = strings.remove_prefix('!', operator, neg_ref)
-            aaronmk
+        equals = operator.endswith('=') # also includes <=, >=
             aaronmk
-            aaronmk
+        # Handle nullable columns
         check_null = False
-            aaronmk
+        if not passthru_null_ref[0]: # NULLs compare equal
-            aaronmk
+            try: left_value = ensure_not_null(db, left_value)
-            aaronmk
+            except ensure_not_null_excs: # fall back to alternate method
                 check_null = equals and isinstance(right_value, Col)
-            aaronmk
+            else:
-            aaronmk
+                if isinstance(left_value, EnsureNotNull):
                     right_value = ensure_not_null(db, right_value,
                         left_value.type) # apply same function to both sides
             aaronmk
-            aaronmk
+        if equals and is_null(right_value): operator = 'IS'
-            aaronmk
+        left = left_value.to_str(db)
         right = right_value.to_str(db)
-            aaronmk
+        # Create str
         str_ = left+' '+operator+' '+right
-            aaronmk
+        if check_null:
-            aaronmk
+            str_ = '('+str_+' OR ('+left+' IS NULL AND '+right+' IS NULL))'
         if neg_ref[0]: str_ = 'NOT '+str_
-            aaronmk
+        return str_
             aaronmk
-            aaronmk
+# Tells as_ValueCond() to assume a non-ValueCond is a literal value
 assume_literal = object()
 def as_ValueCond(value, default_table=assume_literal):
     if not isinstance(value, ValueCond):
         if default_table is not assume_literal:
-            aaronmk
+            value = with_default_table(value, default_table)
-            aaronmk
+        return CompareCond(value)
-            aaronmk
+    else: return value
             aaronmk
-            aaronmk
+##### Joins
-            aaronmk
+join_same = object() # tells Join the left and right columns have the same name
             aaronmk
-            aaronmk
+# Tells Join the left and right columns have the same name and are never NULL
 join_same_not_null = object()
-            aaronmk
+filter_out = object() # tells Join to filter out rows that match the join
-            aaronmk
+class Join(BasicObject):
-            aaronmk
+    def __init__(self, table, mapping={}, type_=None):
-            aaronmk
+        '''
         @param mapping dict(right_table_col=left_table_col, ...)
-            aaronmk
+            * if left_table_col is join_same: left_table_col = right_table_col
-            aaronmk
+              * Note that right_table_col must be a string
             * if left_table_col is join_same_not_null:
               left_table_col = right_table_col and both have NOT NULL constraint
               * Note that right_table_col must be a string
-            aaronmk
+        @param type_ None (for plain join)|str (e.g. 'LEFT')|filter_out
             * filter_out: equivalent to 'LEFT' with the query filtered by
               `table_pkey IS NULL` (indicating no match)
         '''
         if util.is_str(table): table = Table(table)
         assert type_ == None or util.is_str(type_) or type_ is filter_out
         self.table = table
         self.mapping = mapping
         self.type_ = type_
-            aaronmk
+    def to_str(self, db, left_table_):
-            aaronmk
+        def join(entry):
             '''Parses non-USING joins'''
             right_table_col, left_table_col = entry
-            aaronmk
+            # Switch order (right_table_col is on the left in the comparison)
             left = right_table_col
             right = left_table_col
-            aaronmk
+            left_table = self.table
             right_table = left_table_
             aaronmk
-            aaronmk
+            # Parse left side
-            aaronmk
+            left = with_default_table(left, left_table)
             aaronmk
-            aaronmk
+            # Parse special values
-            aaronmk
+            left_on_right = Col(left.name, right_table)
             if right is join_same: right = left_on_right
-            aaronmk
+            elif right is join_same_not_null:
-            aaronmk
+                right = CompareCond(left_on_right, '~=')
             aaronmk
-            aaronmk
+            # Parse right side
-            aaronmk
+            right = as_ValueCond(right, right_table)
             aaronmk
             return right.to_str(db, left)
             aaronmk
-            aaronmk
+        # Create join condition
         type_ = self.type_
-            aaronmk
+        joins = self.mapping
-            aaronmk
+        if joins == {}: join_cond = None
         elif type_ is not filter_out and reduce(operator.and_,
-            aaronmk
+            (v is join_same_not_null for v in joins.itervalues())):
-            aaronmk
+            # all cols w/ USING, so can use simpler USING syntax
-            aaronmk
+            cols = map(to_name_only_col, joins.iterkeys())
             join_cond = 'USING ('+(', '.join((c.to_str(db) for c in cols)))+')'
-            aaronmk
+        else: join_cond = combine_conds(map(join, joins.iteritems()), 'ON')
             aaronmk
-            aaronmk
+        if isinstance(self.table, NamedTable): whitespace = '\n'
         else: whitespace = ' '
-            aaronmk
+        # Create join
         if type_ is filter_out: type_ = 'LEFT'
-            aaronmk
+        str_ = ''
         if type_ != None: str_ += type_+' '
-            aaronmk
+        str_ += 'JOIN'+whitespace+self.table.to_str(db)
         if join_cond != None: str_ += whitespace+join_cond
-            aaronmk
+        return str_
             aaronmk
-            aaronmk
+    def __repr__(self): return self.to_str(mockDb, '<left_table>')
             aaronmk
 ##### Value exprs
-            aaronmk
+all_cols = CustomCode('*')
-            aaronmk
+default = CustomCode('DEFAULT')
-            aaronmk
+row_count = FunctionCall(InternalFunction('COUNT'), all_cols)
             aaronmk
-            aaronmk
+class Coalesce(FunctionCall):
     def __init__(self, *args):
         FunctionCall.__init__(self, InternalFunction('COALESCE'), *args)
             aaronmk
-            aaronmk
+class Nullif(FunctionCall):
     def __init__(self, *args):
         FunctionCall.__init__(self, InternalFunction('NULLIF'), *args)
-            aaronmk
+null_as_str = Cast('text', 'NULL')
 def to_text(value): return Coalesce(Cast('text', value), null_as_str)
-            aaronmk
+# See <http://www.postgresql.org/docs/8.3/static/datatype-numeric.html>
-            aaronmk
+null_sentinels = {
     'character varying': r'\N',
     'double precision': 'NaN',
     'integer': 2147483647,
     'text': r'\N',
-            aaronmk
+    'timestamp with time zone': 'infinity',
     'taxonrank': 'unknown',
             aaronmk
             aaronmk
-            aaronmk
+class EnsureNotNull(Coalesce):
-            aaronmk
+    def __init__(self, value, type_):
-            aaronmk
+        if isinstance(type_, ArrayType): null = []
         else: null = null_sentinels[type_]
         Coalesce.__init__(self, as_Col(value), Cast(type_, null))
             aaronmk
         self.type = type_
             aaronmk
     def to_str(self, db):
         col = self.args[0]
         index_col_ = index_col(col)
         if index_col_ != None: return index_col_.to_str(db)
-            aaronmk
+        return Coalesce.to_str(self, db)
             aaronmk
-            aaronmk
+#### Arrays
-            aaronmk
+class ArrayMerge(FunctionCall):
-            aaronmk
+    def __init__(self, sep, array):
         array = to_Array(array)
         FunctionCall.__init__(self, InternalFunction('array_to_string'), array,
             sep)
-            aaronmk
+def merge_not_null(db, sep, values):
-            aaronmk
+    return ArrayMerge(sep, map(to_text, values))
             aaronmk
-            aaronmk
+##### Table exprs
 class Values(Code):
     def __init__(self, values):
-            aaronmk
+        '''
         @param values [...]|[[...], ...] Can be one or multiple rows.
         '''
-            aaronmk
+        Code.__init__(self)
-            aaronmk
+        rows = values
         if len(values) >= 1 and not lists.is_seq(values[0]): # only one row
             rows = [values]
         for i, row in enumerate(rows):
             rows[i] = map(remove_col_rename, map(as_Value, row))
             aaronmk
-            aaronmk
+        self.rows = rows
             aaronmk
     def to_str(self, db):
-            aaronmk
+        return 'VALUES '+(', '.join((Tuple(*r).to_str(db) for r in self.rows)))
             aaronmk
-            aaronmk
+def NamedValues(name, cols, values):
-            aaronmk
+    '''
-            aaronmk
+    @param cols None|[...]
-            aaronmk
+    @post `cols` will be changed to Col objects with the table set to `name`.
     '''
-            aaronmk
+    table = NamedTable(name, Values(values), cols)
-            aaronmk
+    if cols != None: set_cols_table(table, cols)
-            aaronmk
+    return table
             aaronmk
-            aaronmk
+##### Database structure
-            aaronmk
+def is_nullable(db, value):
     try: return is_null(value) or db.col_info(value).nullable
     except NoUnderlyingTableException: return True # not a table column
-            aaronmk
+text_types = set(['character varying', 'text'])
             aaronmk
-            aaronmk
+def is_text_col(db, col): return db.col_info(col).type in text_types
-            aaronmk
+ensure_not_null_excs = (NoUnderlyingTableException, KeyError)
-            aaronmk
+def ensure_not_null(db, col, type_=None):
-            aaronmk
+    '''
-            aaronmk
+    @param col If type_ is not set, must have an underlying column.
-            aaronmk
+    @param type_ If set, overrides the underlying column's type and casts the
         column to it if needed.
-            aaronmk
+    @return EnsureNotNull|Col
     @throws ensure_not_null_excs
     '''
-            aaronmk
+    nullable = True
     try: typed_col = db.col_info(underlying_col(col))
     except NoUnderlyingTableException:
-            aaronmk
+        col = remove_col_rename(col)
-            aaronmk
+        if is_literal(col) and not is_null(col): nullable = False
-            aaronmk
+        elif type_ == None: raise
-            aaronmk
+    else:
         if type_ == None: type_ = typed_col.type
-            aaronmk
+        elif type_ != typed_col.type: col = Cast(type_, col)
-            aaronmk
+        nullable = typed_col.nullable
-            aaronmk
+    if nullable:
         try: col = EnsureNotNull(col, type_)
         except KeyError, e:
             # Warn of no null sentinel for type, even if caller catches error
             warnings.warn(UserWarning(exc.str_(e)))
             raise
-            aaronmk
+    return col
             aaronmk
 def try_mk_not_null(db, value):
     '''
     Warning: This function does not guarantee that its result is NOT NULL.
     '''
     try: return ensure_not_null(db, value)
     except ensure_not_null_excs: return value

Project

General

Profile

root/lib/sql_gen.py @ 5289